//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2016 Richard S. Smith and collaborators.
//
// If you use MorphoGraphX in your work, please cite:
//   http://dx.doi.org/10.7554/eLife.05864
//
// MorphoGraphX is free software, and is licensed under under the terms of the 
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
// 
#include <StackProcessMorphology.hpp>

#include <StackProcess.hpp>
#include <cuda/CudaExport.hpp>
#include <Image.hpp>

#include <CImg.h>

using namespace cimg_library;

namespace mgx 
{
  typedef CImg<ushort> CImgUS;
  
  bool EdgeDetectProcess::run(Store *input, Store *output, float threshold, float multiplier, float factor,
                                     uint fillValue)
  {
    ushort lowthresh = trim(ushort(float(threshold)  *factor), ushort(0), ushort(0xFFFF));
    ushort fill = trim(fillValue, uint(0), uint(0xFFFF));
    Point3i size(input->stack()->size());
    edgeDetectGPU(size, lowthresh, threshold, multiplier, fill, input->data(), output->data());
  
    output->copyMetaData(input);
    output->changed();
    return true;
  }
  
  REGISTER_PROCESS(EdgeDetectProcess);
  
  bool DilateStack::run(Stack *stack, const Store *input, Store *output, 
         uint xradius, uint yradius, uint zradius, bool auto_resize, bool roundNhbd)
  {
    Point3i radius(xradius, yradius, zradius);
    if(auto_resize) {
      ResizeCanvas resize(*this);
      resize.run(stack, true, true, 2 * radius);
    }
    Point3i size(stack->size());
    if(dilateGPU(size, radius, roundNhbd, input->data(), output->data())) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    output->copyMetaData(input);
    output->changed();
    return true;
  }
  REGISTER_PROCESS(DilateStack);
  
  bool ErodeStack::run(Stack *stack, const Store *input, Store *output, uint xradius, uint yradius, uint zradius,
                              bool auto_resize, bool roundNhbd)
  {
    Point3i radius(xradius, yradius, zradius);
    Point3i size(stack->size());
    if(erodeGPU(size, radius, input->labels(), roundNhbd, input->data(), output->data())) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    if(auto_resize) {
      ResizeCanvas resize(*this);
      resize.run(stack, true, true, -2 * radius);
    }
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(ErodeStack);
  
  bool OpeningStack::run(Stack *stack, const Store *input, Store *output, 
                                     uint xradius, uint yradius, uint zradius, bool roundNhbd)
  {
    Point3i radius(xradius, yradius, zradius);
    Point3i size(stack->size());
    if(erodeGPU(size, radius, input->labels(), roundNhbd, input->data(), output->data())) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    if(dilateGPU(size, radius, roundNhbd, output->data(), output->data())) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(OpeningStack);
  
  bool ClosingStack::run(Stack *stack, const Store *input, Store *output, 
                                uint xradius, uint yradius, uint zradius, bool roundNhbd)
  {
    Point3i radius(xradius, yradius, zradius);
    Point3i size(stack->size());
    Point3i newSize = size + 2  *radius;
  
    const HVecUS &data = input->data();
    HVecUS &outputData = output->data();
  
    {
      HVecUS resized = resize(data, size, newSize, true);
      outputData.swap(resized);
    }
  
    if(dilateGPU(newSize, radius, roundNhbd, outputData, outputData)) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    if(erodeGPU(newSize, radius, false, roundNhbd, outputData, outputData)) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
  
    {
      HVecUS resized = resize(outputData, newSize, size, true);
      outputData.swap(resized);
    }
  
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(ClosingStack);

  #define OFFSET(_x, _y) ((_y)*size.x() + (_x))
  bool FillHolesProcess::run(const HVecUS &input, HVecUS &output, const Point3i &size, 
                                    int xradius, int yradius, int threshold, int depth, int fillval)
  {
    //HVecUS input = _input;
    if(&input != &output)
      output = input;
    fillval = trim(fillval, 0, 0xFFFF);
    depth = trim(depth, 0, size.z());
  
    // First create height map, allocate update map
    std::vector<int> hmap(size.x() * size.y());
    std::vector<int> umap(size.x() * size.y());
    for(int x = 0; x < size.x(); x++) {
      for(int y = 0; y < size.y(); y++) {
        uint idx = OFFSET(x, y);
        hmap[idx] = umap[idx] = 0;
        for(int z = 0; z < size.z(); z++)
          if(output[getOffset(x, y, z, size)] >= threshold)
            hmap[idx] = z;
      }
    }
  
    // Scan in X direction
    if(xradius > 0) {
      for(int y = 0; y < size.y(); y++) {
        // first look for range of interest
        int firstx = 0;
        while(firstx < int(size.x()) and hmap[OFFSET(firstx, y)] == 0)
          firstx++;
        if(firstx >= int(size.x() - 1))
          continue;
        int lastx = size.x() - 1;
        while(lastx > firstx and hmap[OFFSET(lastx, y)] == 0)
          lastx--;
        if(--lastx <= ++firstx)
          continue;
  
        for(int x = firstx; x < lastx; x++) {
          uint idx = OFFSET(x, y);
          // max prev and next z's in range
          int pz = hmap[idx], nz = hmap[idx];
          int px = x, nx = x;
          int startx = x - xradius;
          int endx = x + xradius;
          if(startx < firstx - 1)
            startx = firstx - 1;
          if(endx > lastx + 1)
            endx = lastx + 1;
          for(int ix = x - 1; ix >= startx; ix--)
            if(hmap[OFFSET(ix, y)] > pz + depth) {
              pz = hmap[OFFSET(ix, y)];
              px = ix;
              if(depth == 0)
                break;
            }
          for(int ix = x + 1; ix <= endx; ix++)
            if(hmap[OFFSET(ix, y)] > nz + depth) {
              nz = hmap[OFFSET(ix, y)];
              nx = ix;
              if(depth == 0)
                break;
            }
          // Calculate line intersection with pixel
          if(px != x and nx != x) {
            int newz = pz + int(float(x - px) * float(nz - pz) / float(nx - px) + .5);
            if(newz > hmap[idx])
              umap[idx] = newz;
          }
        }
      }
    }
  
    // Scan in Y direction
    if(yradius > 0) {
      for(int x = 0; x < size.x(); x++) {
        // first look for range of interest
        int firsty = 0;
        while(firsty < int(size.y()) and hmap[OFFSET(x, firsty)] == 0)
          firsty++;
        if(firsty >= int(size.y() - 1))
          continue;
        int lasty = size.y() - 1;
        while(lasty > firsty and hmap[OFFSET(x, lasty)] == 0)
          lasty--;
        if(--lasty <= ++firsty)
          continue;
  
        for(int y = firsty; y < lasty; y++) {
          uint idx = OFFSET(x, y);
          // max prev and next z's in range
          int pz = hmap[idx], nz = hmap[idx];
          int py = y, ny = y;
          int starty = y - yradius;
          int endy = y + yradius;
          if(starty < firsty - 1)
            starty = firsty - 1;
          if(endy > lasty + 1)
            endy = lasty + 1;
          for(int iy = y - 1; iy >= starty; iy--)
            if(hmap[OFFSET(x, iy)] > pz + depth) {
              pz = hmap[OFFSET(x, iy)];
              py = iy;
              if(depth == 0)
                break;
            }
          for(int iy = y + 1; iy <= endy; iy++)
            if(hmap[OFFSET(x, iy)] > nz + depth) {
              nz = hmap[OFFSET(x, iy)];
              ny = iy;
              if(depth == 0)
                break;
            }
          // Calculate line intersection with pixel
          if(py != y and ny != y) {
            int newz = pz + int(float(y - py) * float(nz - pz) / float(ny - py));
            if(newz > hmap[idx] and newz > umap[idx])
              umap[idx] = newz;
          }
        }
      }
    }
  
    // Update stack
    for(int x = 0; x < int(size.x()); x++)
      for(int y = 0; y < int(size.y()); y++)
        for(int z = umap[OFFSET(x, y)]; z >= 0; z--)
          output[getOffset(x, y, z, size)] = fillval;
  
    return true;
  }
  #undef OFFSET
  REGISTER_PROCESS(FillHolesProcess);
  
  static bool openCloseStack(Stack *stack, const Store *input, Store *output, 
    uint xradius, uint yradius, uint zradius, bool roundNhbd, bool is_opening, Process *proc)
  {
    Point3i radius(xradius, yradius, zradius);
    Point3u size = stack->size();
  
    const HVecUS &data = input->data();
    HVecUS &outputData = output->data();
    std::vector<BoundingBox3u> bBoxImg(65536, BoundingBox3u());
  
    if(input != output)
      outputData = data;
  
    // Progress progress(QString("Finding object bounds-%1").arg(mesh->userId()), data.size());
    // int progress_step = data.size()/100;
  
    const ushort *pdata = data.data();
    for(uint z = 0; z < size.z(); z++)
      for(uint y = 0; y < size.y(); y++)
        for(uint x = 0; x < size.x(); x++, pdata++) {
          int label = *pdata;
          if(label == 0)
            continue;
          bBoxImg[label] |= Point3u(x, y, z);
        }
  
    #define OFFSET(x, y, z, xsz, ysz) ((size_t(z)  *(ysz) + (y))  *(xsz) + (x))
    for(int i = 0; i < 65536; ++i) {
      // Check for valid label
      if(bBoxImg[i].empty())
        continue;
      // Copy data to temp area
      Point3i bSize(bBoxImg[i].pmax() - bBoxImg[i].pmin() + Point3u(1,1,1));
      HVecUS tempData(size_t(bSize.x())  *bSize.y()  *bSize.z());
      uint x, y, z, xb, yb, zb;
      for(zb = 0, z = bBoxImg[i].pmin().z(); z <= bBoxImg[i].pmax().z(); z++, zb++)
        for(yb = 0, y = bBoxImg[i].pmin().y(); y <= bBoxImg[i].pmax().y(); y++, yb++)
          for(xb = 0, x = bBoxImg[i].pmin().x(); x <= bBoxImg[i].pmax().x(); x++, xb++)
            if(data[stack->offset(x, y, z)] == i)
              tempData[OFFSET(xb, yb, zb, bSize.x(), bSize.y())] = 1;
            else
              tempData[OFFSET(xb, yb, zb, bSize.x(), bSize.y())] = 0;
  
      if(is_opening) {
        if(erodeGPU(bSize, radius, input->labels(), roundNhbd, tempData, tempData)) {
          proc->setErrorMessage("Error while running CUDA process");
          return false;
        }
        if(dilateGPU(bSize, radius, roundNhbd, tempData, tempData)) {
          proc->setErrorMessage("Error while running CUDA process");
          return false;
        }
        for(zb = 0, z = bBoxImg[i].pmin().z(); z <= bBoxImg[i].pmax().z(); z++, zb++)
          for(yb = 0, y = bBoxImg[i].pmin().y(); y <= bBoxImg[i].pmax().y(); y++, yb++)
            for(xb = 0, x = bBoxImg[i].pmin().x(); x <= bBoxImg[i].pmax().x(); x++, xb++) {
              size_t off = stack->offset(x, y, z);
              if(outputData[off] == i)
                outputData[off] *= tempData[OFFSET(xb, yb, zb, bSize.x(), bSize.y())];
            }
      } else {
        if(dilateGPU(bSize, radius, roundNhbd, tempData, tempData)) {
          proc->setErrorMessage("Error while running CUDA process");
          return false;
        }
        if(erodeGPU(bSize, radius, false, roundNhbd, tempData, tempData)) {
          proc->setErrorMessage("Error while running CUDA process");
          return false;
        }
        for(zb = 0, z = bBoxImg[i].pmin().z(); z <= bBoxImg[i].pmax().z(); z++, zb++)
          for(yb = 0, y = bBoxImg[i].pmin().y(); y <= bBoxImg[i].pmax().y(); y++, yb++)
            for(xb = 0, x = bBoxImg[i].pmin().x(); x <= bBoxImg[i].pmax().x(); x++, xb++)
              if(tempData[OFFSET(xb, yb, zb, bSize.x(), bSize.y())] == 1)
                outputData[stack->offset(x, y, z)] = i;
      }
    }
  #undef OFFSET
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  
  bool OpenStackLabel::run(Stack *stack, const Store *input, Store *output, 
                                  uint xradius, uint yradius, uint zradius, bool roundNhbd)
  {
    return openCloseStack(stack, input, output, xradius, yradius, zradius, roundNhbd, true, this);
  }
  REGISTER_PROCESS(OpenStackLabel);
  
  bool CloseStackLabel::run(Stack *stack, const Store *input, Store *output, 
                                  uint xradius, uint yradius, uint zradius, bool roundNhbd)
  {
    return openCloseStack(stack, input, output, xradius, yradius, zradius, roundNhbd, false, this);
  }
  REGISTER_PROCESS(CloseStackLabel);
  
  bool ApplyMaskToStack::run(Stack*, const Store *input, Store *output, const QString &mode, uint threshold)
  {
    const ushort *pin = &input->data()[0];
    ushort *pout = &output->data()[0];
    if(mode == "Normal") {
      for(size_t i = 0; i < output->size(); ++i, ++pin, ++pout)
        if(*pout > threshold)
          *pout = *pin;
        else
          *pout = 0;
    } else if(mode == "Invert") {
      for(size_t i = 0; i < output->size(); ++i, ++pin, ++pout)
        if(*pout <= threshold)
          *pout = *pin;
        else
          *pout = 0;
    } else if(mode == "Combine") {
      for(size_t i = 0; i < output->size(); ++i, ++pin, ++pout)
        if(*pout <= threshold)
          *pout = *pin;
    }
  
    output->changed();
    return true;
  }
  
  REGISTER_PROCESS(ApplyMaskToStack);
  
  bool ApplyMaskLabels::run(Stack *, const Store *input, Store *output, bool invert, uint threshold)
  {
    const ushort *pin = &input->data()[0];
    ushort *pout = &output->data()[0];
    std::set<int> labelSet;
    for(size_t i = 0; i < output->size(); ++i, ++pin, ++pout)
      if(*pout >= threshold)
        labelSet.insert(*pin);
  
    pin = &input->data()[0];
    pout = &output->data()[0];
    if(!invert) {
      for(size_t i = 0; i < output->size(); ++i, ++pin, ++pout)
        if(labelSet.count(*pin) > 0)
          *pout = *pin;
        else
          *pout = 0;
    } else {
      for(size_t i = 0; i < output->size(); ++i, ++pin, ++pout)
        if(labelSet.count(*pin) > 0)
          *pout = 0;
        else
          *pout = *pin;
    }
  
    output->changed();
    return true;
  }
  REGISTER_PROCESS(ApplyMaskLabels);

  bool TrimStackToBezier::run(Stack *stack, const Store *input, Store *output, double distance, int points, bool mask, int fillValue, QString axes)
  {
    if(distance <= 0)
      throw QString("%1::run Distance must be > 0").arg(name());
    if(points <= 0)
      throw QString("%1::run Points must be > 0").arg(name());
    fillValue = trim(fillValue, 0, 65535);
    axes = axes.toLower();
    if(axes != "u" and axes != "v" and axes != "uv")
      throw QString("%1::run Axes must be u, v or uv").arg(name());

    typedef std::map<Point2i, Point3d> P2iP3dMap;
    
    // get the bezier
    CuttingSurface* cutSurf = cuttingSurface();

    P2iP3dMap bezGridMap, diffBezGridMapX, diffBezGridMapY;

    double stepSize = 1.0/(double)(points-1);

    // correct directions, take into account the rotations/translations
    Matrix4d rotMatrixS1, rotMatrixCS;
    stack->getFrame().getMatrix(rotMatrixS1.data());
    cutSurf->frame().getMatrix(rotMatrixCS.data());
    Matrix4d mGLTot = transpose(inverse(rotMatrixS1)) * transpose(rotMatrixCS);

    std::unordered_set<Point3d> bezPoints;

    int uPoints = points;
    int vPoints = points;
    if(axes == "u")
      vPoints = 1;
    if(axes == "v")
      uPoints = 1;

    // create points on the bezier to test
    for(int i = 0; i < uPoints; i++) {
      double u = i * stepSize;
      for(int j = 0; j < vPoints; j++){
        double v = j * stepSize;
        Point3d p(cutSurf->evalCoord(u,v)); // get bez world coordinate
        p = multMatrix4Point3(mGLTot, p); // correct rotations
        bezPoints.insert(p);
      }
    }
    Information::out << QString("%1 points created to test").arg(bezPoints.size()) << endl;

    auto size = stack->size();
    auto &inData = input->data();
    auto &outData = output->data();

    // Loop through voxels and decide what to keep
    #pragma omp parallel for
    for(uint z = 0; z < size.z(); z++) {
      if(!progressAdvance())
        continue;
      for(uint y = 0; y < size.y(); y++) {
        for(uint x = 0; x < size.x(); x++) {
          bool keep = false;
          for(auto &p : bezPoints)
            if(norm(Point3d(stack->imageToWorld(Point3i(x, y, z))) - p) < distance) {
              keep = true;
              break;
            }
          if(keep) {
            if(mask)
              outData[stack->offset(x, y, z)] = fillValue;
            else
              outData[stack->offset(x, y, z)] = inData[stack->offset(x, y, z)];
          } else if(!mask)
            outData[stack->offset(x, y, z)] = 0;
        }
      }
    }

    output->changed();
    return true;
  }
  REGISTER_PROCESS(TrimStackToBezier);
}
