//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2015 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 "StackProcessMerge.hpp"
#include "Information.hpp"
#include "Progress.hpp"

namespace mgx {
  namespace {
    
    void trilinear_filtering(
      Stack* target, Store* target_store, const HVecUS& source, const Point3u& ssize,
      const BoundingBox3i& bbox,       // Bounding box of the source in the target canvas
      const Matrix4f& toSourceImage)   // Transform world coordinates of target to image coordinate of source
    {
      HVecUS& data = target_store->data();
      const Point3i& bmin = bbox.pmin();
      const Point3i& bmax = bbox.pmax();
      #define SOFFSET(X, Y, Z) ((size_t(Z) * ssize.y() + size_t(Y)) * ssize.x() + size_t(X))
      #pragma omp parallel for
      for(int z = bmin.z(); z < bmax.z(); ++z)
        for(int y = bmin.y(); y < bmax.y(); ++y)
          for(int x = bmin.x(); x < bmax.x(); ++x) {
            Point4f lp = homogeneous(target->imageToWorld(Point3i(x, y, z)));
            Point3f ip = cartesian(toSourceImage * lp);
            // Perform tri-linear interpolation
            float dx = ip.x();
            float dy = ip.y();
            float dz = ip.z();
            int x_base = (int)floorf(dx);
            int y_base = (int)floorf(dy);
            int z_base = (int)floorf(dz);
            dx -= x_base;
            dy -= y_base;
            dz -= z_base;
    
            // Check if we are in the stack
            if(x_base < (int)ssize.x() and y_base < (int)ssize.y() and z_base < (int)ssize.z() and (x_base + 1) >= 0
               and (y_base + 1) >= 0 and (z_base + 1) >= 0) {
              // Find values. The cube is represented as:
              //
              //         6---7
              //        /|  /|
              //       3---5 |
              //       | 2-|-4
              // z y   |/  |/
              // |/    0---1
              // *--x
              float vs[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
    
              // Set values checking bounds
              if(x_base >= 0) {
                if(y_base >= 0) {
                  if(z_base >= 0)
                    vs[0] = float(source[SOFFSET(x_base, y_base, z_base)]);
                  if(z_base + 1 < (int)ssize.z())
                    vs[3] = float(source[SOFFSET(x_base, y_base, z_base + 1)]);
                }
                if(y_base + 1 < (int)ssize.y()) {
                  if(z_base >= 0)
                    vs[2] = float(source[SOFFSET(x_base, y_base + 1, z_base)]);
                  if(z_base + 1 < (int)ssize.z())
                    vs[6] = float(source[SOFFSET(x_base, y_base + 1, z_base + 1)]);
                }
              }
              if(x_base + 1 < (int)ssize.x()) {
                if(y_base >= 0) {
                  if(z_base >= 0)
                    vs[1] = float(source[SOFFSET(x_base + 1, y_base, z_base)]);
                  if(z_base + 1 < (int)ssize.z())
                    vs[5] = float(source[SOFFSET(x_base + 1, y_base, z_base + 1)]);
                }
                if(y_base + 1 < (int)ssize.y()) {
                  if(z_base >= 0)
                    vs[4] = float(source[SOFFSET(x_base + 1, y_base + 1, z_base)]);
                  if(z_base + 1 < (int)ssize.z())
                    vs[7] = float(source[SOFFSET(x_base + 1, y_base + 1, z_base + 1)]);
                }
              }
    
              float v01 = (1 - dx) * vs[0] + dx * vs[1];
              float v35 = (1 - dx) * vs[3] + dx * vs[5];
              float v67 = (1 - dx) * vs[6] + dx * vs[7];
              float v24 = (1 - dx) * vs[2] + dx * vs[4];
    
              float y_bottom = (1 - dy) * v01 + dy * v24;
              float y_top = (1 - dy) * v35 + dy * v67;
    
              ushort v = ushort((1 - dz) * y_bottom + dz * y_top);
              size_t off = target->offset(x, y, z);
              data[off] = v;
            }
          }
    #undef SOFFSET
    }
    
    void
    nearest_filtering(Stack* target, Store* target_store, const HVecUS& source, const Point3u& ssize,
                      const BoundingBox3i& bbox,     // Bounding box of the source in the target canvas
                      const Matrix4f& toSourceImage) // Transform world coordinates of target to image coordinate of source
    {
      HVecUS& data = target_store->data();
      const Point3i& bmin = bbox.pmin();
      const Point3i& bmax = bbox.pmax();
      #define SOFFSET(X, Y, Z) ((size_t(Z) * ssize.y() + size_t(Y)) * ssize.x() + size_t(X))
      #pragma omp parallel for
      for(int z = bmin.z(); z < bmax.z(); ++z)
        for(int y = bmin.y(); y < bmax.y(); ++y)
          for(int x = bmin.x(); x < bmax.x(); ++x) {
            Point4f lp = homogeneous(target->imageToWorld(Point3i(x, y, z)));
            Point3f ip = cartesian(toSourceImage * lp);
            // Perform tri-linear interpolation
            float dx = ip.x();
            float dy = ip.y();
            float dz = ip.z();
            int x_base = (int)floorf(dx);
            int y_base = (int)floorf(dy);
            int z_base = (int)floorf(dz);
            dx -= x_base;
            dy -= y_base;
            dz -= z_base;
    
            if(dx > 0.5)
              x_base++;
            if(dy > 0.5)
              y_base++;
            if(dz > -.5)
              z_base++;
    
            // Check if we are in the stack
            if(x_base < (int)ssize.x() and y_base < (int)ssize.y() and z_base < (int)ssize.z() and x_base >= 0
               and y_base >= 0 and z_base >= 0) {
              ushort value = source[SOFFSET(x_base, y_base, z_base)];
              data[target->offset(x, y, z)] = value;
            }
          }
    #undef SOFFSET
    }
    
    void filtering(Stack* target, Store* target_store, const HVecUS& source, const Point3u& ssize,
                   const BoundingBox3i& bbox,     // Bounding box of the source in the target canvas
                   const Matrix4f& toSourceImage) // Transform world coordinates of target to image coordinate of source
    {
      if(target_store->labels())
        nearest_filtering(target, target_store, source, ssize, bbox, toSourceImage);
      else
        trilinear_filtering(target, target_store, source, ssize, bbox, toSourceImage);
    }
  } 
  
  bool AlignCanvas::run(Stack* target, Stack* other, bool change_both, QString interp)
  {
    if(change_both)
      return projectGlobal(target, other, interp);
    else
      return projectOnStack(target, other, interp);
  }
  
  bool AlignCanvas::projectOnStack(Stack* target, const Stack* other, QString interp)
  {
    // First, find the new bounding box
    BoundingBox3i bbox = target->boundingBox();
  
    // Bounding box of the other image
    BoundingBox3i obbox = other->boundingBox();
  
    Matrix4f toTarget, toOther;
  
    {
      qglviewer::Frame local_frame = target->getFrame();
      qglviewer::Frame other_frame = other->getFrame();
  
      double gl_m[16];
      local_frame.inverse().getMatrix(gl_m);
      toTarget = transpose(Matrix4f(gl_m));
      other_frame.getMatrix(gl_m);
      toTarget *= transpose(Matrix4f(gl_m));
  
      other_frame.inverse().getMatrix(gl_m);
      toOther = transpose(Matrix4f(gl_m));
      local_frame.getMatrix(gl_m);
      toOther *= transpose(Matrix4f(gl_m));
    }
  
    // Update obbox to be bbox of the other image in the target frame
    {
      BoundingBox3f world_obbox = other->imageToWorld(obbox);
      Point3f pmin = world_obbox.pmin();
      Point3f pmax = world_obbox.pmax() + other->step();
  
      // Create the 8 corners
      // The cube is represented as:
      //
      //         6---7
      //        /|  /|
      //       3---5 |
      //       | 2-|-4
      // z y   |/  |/
      // |/    0---1
      // *--x
      Point3f ps[8] = { pmin, pmin, pmin, pmin, pmax, pmax, pmax, pmax };
      ps[1].x() = pmax.x();
      ps[2].y() = pmax.y();
      ps[3].z() = pmax.z();
  
      ps[6].x() = pmin.x();
      ps[5].y() = pmin.y();
      ps[4].z() = pmin.z();
  
      // Find the new BoundingBox
      obbox.reset();
      for(size_t i = 0; i < 8; ++i) {
        Point4f lp = homogeneous(ps[i]);
        obbox |= target->worldToImagei(cartesian(toTarget * lp));
      }
    }
  
    Information::out << "obbox = " << obbox.pmin() << " -- " << obbox.pmax() << endl;
  
    bbox |= obbox;
  
    Information::out << "bbox = " << bbox.pmin() << " -- " << bbox.pmax() << endl;
  
    // Resize the image
    Point3s s = Point3s(bbox.size());
    Point3i shift = -bbox.pmin();
  
    Information::out << "new size = " << s << endl << "shift = " << shift << endl;
  
    if(shift.x() < 0 or shift.y() < 0 or shift.z() < 0)
      return setErrorMessage(
        QString("Error, computed invalid shift of (%1,%2,%3)").arg(shift.x()).arg(shift.y()).arg(shift.z()));
  
    Point3u tsize = target->size();
    Point3f new_origin = target->imageToWorld(Point3f(bbox.pmin()) - 0.5f);
  
    Point3i bmin = obbox.pmin() + shift;
    Point3i bmax = obbox.pmax() + shift;
  
    Information::out << "box from " << bmin << " to " << bmax << endl;
    {
      Point4f lmin = homogeneous(target->imageToWorld(bmin));
      Point3f ipmin = other->worldToImagef(cartesian(toOther * lmin));
      Point4f lmax = homogeneous(target->imageToWorld(bmax));
      Point3f ipmax = other->worldToImagef(cartesian(toOther * lmax));
  
      Information::out << "Translate into " << ipmin << " to " << ipmax << endl;
    }
  
    Matrix4f toOtherImage = other->worldToImage() * toOther;
  
    HVecUS data = target->main()->data();
  
    // Resize stack
    target->setSize(Point3u(s));
    target->setOrigin(new_origin);
  
    HVecUS& main_data = target->main()->data();
    HVecUS& work_data = target->work()->data();
  
    // Reset to 0 both stacks
    for(size_t k = 0; k < main_data.size(); ++k) {
      main_data[k] = 0;
      work_data[k] = 0;
    }
  
    #pragma omp parallel for
    for(uint z = 0; z < tsize.z(); ++z) {
      size_t k = size_t(z) * tsize.y() * tsize.x();
      for(uint y = 0; y < tsize.y(); ++y)
        for(uint x = 0; x < tsize.x(); ++x, ++k) {
          size_t k1 = target->offset(x + shift.x(), y + shift.y(), z + shift.z());
          main_data[k1] = data[k];
        }
    }
  
    const HVecUS& other_data = other->currentStore()->data();
    Point3u osize = other->size();
    target->work()->copyMetaData(other->currentStore());

    if(interp == "Nearest")
      nearest_filtering(target, target->work(), other_data, osize, BoundingBox3i(bmin, bmax), toOtherImage);
    else
      filtering(target, target->work(), other_data, osize, BoundingBox3i(bmin, bmax), toOtherImage);
  
    target->main()->changed();
    target->work()->changed();
  
    return true;
  }
  
  namespace {
    bool projectStack(Stack* stack, const Matrix4f& toStack, const Point3f& new_origin, const Point3u& new_size,
                      const Point3f& new_step, const BoundingBox3i& bbox, QString interp = "Linear")
    {
      HVecUS main, work;
      using std::swap;
    
      // Swap to save time
      swap(main, stack->main()->data());
      swap(work, stack->work()->data());
    
      Matrix4f toImage = stack->worldToImage() * toStack;
    
      Point3u old_size = stack->size();
    
      QString fmain = stack->main()->file();
      QString fwork = stack->work()->file();
    
      stack->setSize(new_size);
      stack->setStep(new_step);
      stack->setOrigin(new_origin);
    
      if(interp == "Nearest") {
        nearest_filtering(stack, stack->main(), main, old_size, bbox, toImage);
        nearest_filtering(stack, stack->work(), work, old_size, bbox, toImage);
      } else {
        filtering(stack, stack->main(), main, old_size, bbox, toImage);
        filtering(stack, stack->work(), work, old_size, bbox, toImage);
      }
    
      stack->main()->setFile(fmain);
      stack->work()->setFile(fwork);
    
      stack->main()->changed();
      stack->work()->changed();
      return true;
    }
  } 
  
  bool AlignCanvas::projectGlobal(Stack* s1, Stack* s2, QString interp)
  {
    // First, find the new bounding box, aligned with the axis
    Matrix4f toS1, toS2;
    Matrix4f fromS1, fromS2;
  
    {
      double gl_m[16];
      s1->getFrame().getMatrix(gl_m);
      fromS1 = transpose(Matrix4f(gl_m));
      s1->getFrame().inverse().getMatrix(gl_m);
      toS1 = transpose(Matrix4f(gl_m));
  
      s2->getFrame().getMatrix(gl_m);
      fromS2 = transpose(Matrix4f(gl_m));
      s2->getFrame().inverse().getMatrix(gl_m);
      toS2 = transpose(Matrix4f(gl_m));
    }
  
    BoundingBox3i bbox1 = s1->boundingBox();
    BoundingBox3i bbox2 = s2->boundingBox();
    BoundingBox3f world_bbox1;
    BoundingBox3f world_bbox2;
  
    // Find the best bounding box and step
    {
      BoundingBox3f wbbox1 = s1->imageToWorld(bbox1);
      BoundingBox3f wbbox2 = s2->imageToWorld(bbox2);
  
      // Create the 8 corners
      // The cube is represented as:
      //
      //         6---7
      //        /|  /|
      //       3---5 |
      //       | 2-|-4
      // z y   |/  |/
      // |/    0---1
      // *--x
  
      Point3f pmin1 = wbbox1.pmin();
      Point3f pmax1 = wbbox1.pmax() + s1->step();
  
      Point3f pmin2 = wbbox2.pmin();
      Point3f pmax2 = wbbox2.pmax() + s2->step();
  
      Point3f ps1[8] = { pmin1, pmin1, pmin1, pmin1, pmax1, pmax1, pmax1, pmax1 };
      ps1[1].x() = pmax1.x();
      ps1[2].y() = pmax1.y();
      ps1[3].z() = pmax1.z();
  
      ps1[6].x() = pmin1.x();
      ps1[5].y() = pmin1.y();
      ps1[4].z() = pmin1.z();
  
      Point3f ps2[8] = { pmin2, pmin2, pmin2, pmin2, pmax2, pmax2, pmax2, pmax2 };
      ps2[1].x() = pmax2.x();
      ps2[2].y() = pmax2.y();
      ps2[3].z() = pmax2.z();
  
      ps2[6].x() = pmin2.x();
      ps2[5].y() = pmin2.y();
      ps2[4].z() = pmin2.z();
  
      for(size_t i = 0; i < 8; ++i) {
        Point3f p1 = cartesian(fromS1 * homogeneous(ps1[i]));
        Point3f p2 = cartesian(fromS2 * homogeneous(ps2[i]));
        world_bbox1 |= p1;
        world_bbox2 |= p2;
      }
    }
  
    Information::out << "BBox1 = " << world_bbox1 << endl;
    Information::out << "BBox2 = " << world_bbox2 << endl;
  
    Point3f step = min(s1->step(), s2->step());
    BoundingBox3f world_bbox = world_bbox1 | world_bbox2;
  
    Information::out << "BBox = " << world_bbox << endl;
  
    Point3f world_size = world_bbox.pmax() - world_bbox.pmin();
    Point3i new_size = Point3i(map(ceilf, world_size / step));
    Point3f new_origin = world_bbox.pmin();
  
    bbox1.pmin() = Point3i(map(floorf, (world_bbox1.pmin() - new_origin) / step));
    bbox1.pmax() = Point3i(map(ceilf, (world_bbox1.pmax() - new_origin) / step));
    bbox2.pmin() = Point3i(map(floorf, (world_bbox2.pmin() - new_origin) / step));
    bbox2.pmax() = Point3i(map(ceilf, (world_bbox2.pmax() - new_origin) / step));
  
    // First, copy the work and main stores of the current stack
    if(!projectStack(s1, toS1, new_origin, Point3u(new_size), step, bbox1, interp))
      return setErrorMessage("Could not project stack 1");
    if(!projectStack(s2, toS2, new_origin, Point3u(new_size), step, bbox2, interp))
      return setErrorMessage("Could not project stack 2");
  
    Matrix4d unit = Matrix4d::identity();
    s1->getFrame().setFromMatrix(unit.c_data());
    s2->getFrame().setFromMatrix(unit.c_data());
    return true;
  }
  REGISTER_PROCESS(AlignCanvas);
  
  namespace {
  namespace combine {
  
  void max(ushort& s, ushort o)
  {
    if(o > s)
      s = o;
  }
  
  void min(ushort& s, ushort o)
  {
    if(o < s)
      s = 0;
  }
  
  void average(ushort& s, ushort o)
  {
    uint32_t ss = s;
    uint32_t oo = o;
    s = ushort((ss + oo) / 2);
  }
  
  void product(ushort& s, ushort o)
  {
    float ss = float(s) / 65535.f;
    float oo = float(o) / 65535.f;
    s = ushort((ss * oo) * 65535);
  }
  
  void add(ushort& s, ushort o) {
    s = ushort(trim(int(o) + int(s), 0, 65535));
  }
  
  void subtract(ushort& s, ushort o) {
    s = ushort(trim(int(o) - int(s), 0, 65535));
  }
  } // namespace combine
  
  template <typename Fct> void combineStores(const HVecUS& main, HVecUS& work, const Fct& fct)
  {
  #pragma omp parallel for
    for(size_t k = 0; k < main.size(); ++k)
      fct(work[k], main[k]);
  }
  } // namespace
  
  bool CombineStacks::run(Stack* target, QString method)
  {
    const HVecUS& main = target->main()->data();
    HVecUS& work = target->work()->data();
    method = method.toLower();
    if(method == "max")
      combineStores(main, work, combine::max);
    else if(method == "min")
      combineStores(main, work, combine::min);
    else if(method == "average")
      combineStores(main, work, combine::average);
    else if(method == "product")
      combineStores(main, work, combine::product);
    else if(method == "add")
      combineStores(main, work, combine::add);
    else if(method == "subtract")
      combineStores(main, work, combine::subtract);
    else
      return setErrorMessage(QString("Unknown method '%1'").arg(method));
    target->work()->changed();
    return true;
  }
  
  REGISTER_PROCESS(CombineStacks);
  
  bool MergeStacks::run(Stack* target, const Stack* other, QString method, QString interp)
  {
    AlignCanvas align(*this);
    if(not align.projectOnStack(target, other, interp))
      return false;
    CombineStacks combine(*this);
    if(not combine.run(target, method))
      return false;
    return true;
  }
  
  REGISTER_PROCESS(MergeStacks);
}
