//
// 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 <StackProcessLabels.hpp>

#include <cuda/CudaExport.hpp>
#include <Information.hpp>
#include <Progress.hpp>
#include <Thrust.hpp>
#include <Image.hpp>

#include <algorithm>
#include <climits>
#include <QThread>

#include <CImg.h>

using namespace cimg_library;

namespace mgx 
{
  typedef CImg<ushort> CImgUS;
  
  bool WatershedStack::run(Stack* stack, Store* main, Store* labels)
  {
    progressStart(QString("Stack %1 - Watershed transform").arg(stack->userId()), 0);
    if(progressCanceled())
      userCancel();
    Point3u size = stack->size();
    CImgUS image(size.x(), size.y(), size.z(), 1, 0xFFFF);
    image -= CImgUS(main->data().data(), size.x(), size.y(), size.z(), 1, true);
    if(progressCanceled())
      userCancel();
    CImgUS seeds(labels->data().data(), size.x(), size.y(), size.z(), 1, true);
    seeds.watershed(image);
    labels->changed();
    return true;
  }
  
  REGISTER_PROCESS(WatershedStack);
  
  bool ConsolidateRegions::run(Store* data, Store* labels, uint threshold, uint minvoxels)
  {
    // Make a map of light and voxel count at shared walls
    typedef Vector<2, uint> Point2u;
    typedef std::pair<ushort, Point2u> UShortPoint2uPair;
    typedef std::map<ushort, Point2u> UShortPoint2uMap;
  
    const size_t LAB_SZ = 1 << 16;
  
    const Stack* stack = data->stack();
    std::vector<UShortPoint2uMap> light(LAB_SZ);
  
    HVecUS& Labels = labels->data();
    HVecUS& Data = data->data();
  
    progressStart(QString("Consolidating regions in stack %1").arg(stack->userId()), 0);
    uint regions = 0, totregions = 0;
    Point3u size = stack->size();
  
    do {
      // Clear data structures
      for(size_t i = 0; i < LAB_SZ; i++)
        light[i].clear();
  
      for(uint z = 0; z < size.z(); z++) {
        for(uint y = 0; y < size.y(); y++) {
          for(uint x = 0; x < size.x(); x++) {
            // Find pixel offset and label
            size_t off = stack->offset(x, y, z);
            uint label = Labels[off];
            if(label == 0)
              continue;
  
            if(x + 1 < size.x()) {
              size_t noff = stack->offset(x + 1, y, z);
              uint nlabel = Labels[noff];
              if(nlabel != 0 and label != nlabel)
                // Record neighbor, light and voxels
                light[label][nlabel] += Point2u(Data[off] + Data[noff], 2);
            }
            if(x > 0) {
              size_t noff = stack->offset(x - 1, y, z);
              uint nlabel = Labels[noff];
              if(nlabel != 0 and label != nlabel)
                // Record neighbor, light and voxels
                light[label][nlabel] += Point2u(Data[off] + Data[noff], 2);
            }
            if(y + 1 < size.y()) {
              size_t noff = stack->offset(x, y + 1, z);
              uint nlabel = Labels[noff];
              if(nlabel != 0 and label != nlabel)
                // Record neighbor, light and voxels
                light[label][nlabel] += Point2u(Data[off] + Data[noff], 2);
            }
            if(y > 0) {
              size_t noff = stack->offset(x, y - 1, z);
              uint nlabel = Labels[noff];
              if(nlabel != 0 and label != nlabel)
                // Record neighbor, light and voxels
                light[label][nlabel] += Point2u(Data[off] + Data[noff], 2);
            }
            if(z + 1 < size.z()) {
              size_t noff = stack->offset(x, y, z + 1);
              uint nlabel = Labels[noff];
              if(nlabel != 0 and label != nlabel)
                // Record neighbor, light and voxels
                light[label][nlabel] += Point2u(Data[off] + Data[noff], 2);
            }
            if(z > 0) {
              size_t noff = stack->offset(x, y, z - 1);
              uint nlabel = Labels[noff];
              if(nlabel != 0 and label != nlabel)
                // Record neighbor, light and voxels
                light[label][nlabel] += Point2u(Data[off] + Data[noff], 2);
            }
          }
        }
        if(!progressAdvance(0))
          userCancel();
      }
  
      // Find regions (this should be possible to collapse into 1 pass).
      std::vector<ushort> merge(LAB_SZ, 0);
      for(size_t i = 0; i < LAB_SZ; i++) {
        if(light[i].empty())
          continue;
        forall(const UShortPoint2uPair& p, light[i]) {
          if(i < p.first and merge[i] == 0 and p.second.y() > minvoxels
             and p.second.x() / p.second.y() < threshold)
            merge[i] = p.first;
        }
      }
      regions = 0;
      for(size_t i = 0; i < LAB_SZ; i++)
        if(merge[i] > 0)
          regions++;
  
      // Collapse regions
      for(uint i = 0; i < Labels.size(); i++) {
        ushort label = Labels[i];
        ushort mlabel = merge[label];
        if(label > 0 and mlabel > 0)
          Labels[i] = mlabel;
        if(i % size.x() * size.y() == 0 and !progressAdvance(0))
          userCancel();
      }
      totregions += regions;
      SETSTATUS("Stack " << stack->userId() << " " << regions << " regions consolidated");
  
    } while(regions > 0);
  
    labels->changed();
    return true;
  }
  
  REGISTER_PROCESS(ConsolidateRegions);
  
  bool ConsolidateRegionsNormalized::run(Store* data, Store* labels, float tolerance)
  {
    // Make a map of light at shared walls
    typedef std::pair<ushort, ushort> UShortUShortPair;
    typedef std::pair<UShortUShortPair, float> UShortUShortFloatPair;
    typedef std::pair<ushort, float> UShortFloatPair;
  
    const Stack* stack = data->stack();
    HVecUS& Labels = labels->data();
    HVecUS& Data = data->data();
  
    progressStart(QString("Consolidating regions in stack %1").arg(stack->userId()), 0);
    uint regions, totregions = 0;
    Point3u size = stack->size();
    do {
      std::map<UShortUShortPair, float> WallLight;
      std::map<UShortUShortPair, uint> WallPix;
      std::map<ushort, float> CellIntLight;
      std::map<ushort, float> CellIntPix;
  
      for(uint z = 0; z < size.z(); z++) {
        for(uint y = 0; y < size.y(); y++)
          for(uint x = 0; x < size.x(); x++) {
            // Find pixel offset
            size_t off = stack->offset(x, y, z);
            uint label = Labels[off];
  
            bool interior = true;
            if(x > 0) {
              size_t noff = stack->offset(x - 1, y, z);
              uint nlabel = Labels[noff];
              if(label != nlabel) {
                interior = false;
  
                UShortUShortPair pr(label, nlabel);
                UShortUShortPair npr(nlabel, label);
  
                float light = Data[off] + Data[noff];
                WallLight[pr] += light;
                WallLight[npr] += light;
                WallPix[pr]++;
                WallPix[npr]++;
              }
            }
            if(x < size.x() - 1) {
              size_t noff = stack->offset(x + 1, y, z);
              uint nlabel = Labels[noff];
              if(label != nlabel) {
                interior = false;
  
                UShortUShortPair pr(label, nlabel);
                UShortUShortPair npr(nlabel, label);
  
                float light = Data[off] + Data[noff];
                WallLight[pr] += light;
                WallLight[npr] += light;
                WallPix[pr]++;
                WallPix[npr]++;
              }
            }
            if(y > 0) {
              size_t noff = stack->offset(x, y - 1, z);
              uint nlabel = Labels[noff];
              if(label != nlabel) {
                interior = false;
  
                UShortUShortPair pr(label, nlabel);
                UShortUShortPair npr(nlabel, label);
  
                float light = Data[off] + Data[noff];
                WallLight[pr] += light;
                WallLight[npr] += light;
                WallPix[pr]++;
                WallPix[npr]++;
              }
            }
            if(y < size.y() - 1) {
              size_t noff = stack->offset(x, y + 1, z);
              uint nlabel = Labels[noff];
              if(label != nlabel) {
                interior = false;
  
                UShortUShortPair pr(label, nlabel);
                UShortUShortPair npr(nlabel, label);
  
                float light = Data[off] + Data[noff];
                WallLight[pr] += light;
                WallLight[npr] += light;
                WallPix[pr]++;
                WallPix[npr]++;
              }
            }
            if(z > 0) {
              size_t noff = stack->offset(x, y, z - 1);
              uint nlabel = Labels[noff];
              if(label != nlabel) {
                interior = false;
  
                UShortUShortPair pr(label, nlabel);
                UShortUShortPair npr(nlabel, label);
  
                float light = Data[off] + Data[noff];
                WallLight[pr] += light;
                WallLight[npr] += light;
                WallPix[pr]++;
                WallPix[npr]++;
              }
            }
            if(z < size.z() - 1) {
              size_t noff = stack->offset(x, y, z + 1);
              uint nlabel = Labels[noff];
              if(label != nlabel) {
                interior = false;
  
                UShortUShortPair pr(label, nlabel);
                UShortUShortPair npr(nlabel, label);
  
                float light = Data[off] + Data[noff];
                WallLight[pr] += light;
                WallLight[npr] += light;
                WallPix[pr]++;
                WallPix[npr]++;
              }
            }
            if(interior) {
              CellIntLight[label] += Data[off];
              CellIntPix[label]++;
            }
          }
        if(!progressAdvance(0))
          userCancel();
      }
  
      if(WallLight.size() == 0)
        throw QString("Error::Consolidate regions: No cells");
  
      // Get average light at wall and for whole cell
      std::map<ushort, float> CellWallLight;     // Average light at cell walls
      std::map<ushort, uint> CellWallPix;
      uint totalpix = 0;
      float avglight = 0;
      forall(const UShortUShortFloatPair& pr, WallLight) {
        CellWallLight[pr.first.first] += pr.second;
        CellWallPix[pr.first.first] += WallPix[pr.first];
        avglight += pr.second;
        totalpix += WallPix[pr.first];
      }
      // Check for light in image
      if(totalpix != 0)
        avglight /= totalpix;
      if(avglight == 0)
        throw QString("Error::Consolidate regions: No light in image");
  
      // Calculate total light for cell minus the wall
      // Average light at cell wall not including current wall
      std::map<UShortUShortPair, float> CellBordLight;
      std::map<UShortUShortPair, uint> CellBordPix;
      forall(const UShortUShortFloatPair& pr, WallLight) {
        CellBordLight[pr.first] = CellWallLight[pr.first.first] - pr.second;
        CellBordPix[pr.first] = CellWallPix[pr.first.first] - WallPix[pr.first];
      }
      // Divide by Pix counts
      forall(const UShortUShortFloatPair& pr, WallLight) {
        WallLight[pr.first] /= WallPix[pr.first];
        CellBordLight[pr.first] /= CellBordPix[pr.first];
      }
      forall(const UShortFloatPair& pr, CellWallLight)
        CellWallLight[pr.first] /= float(CellWallPix[pr.first]);
      forall(const UShortFloatPair& pr, CellIntLight)
        CellIntLight[pr.first] /= float(CellIntPix[pr.first]);
  
      // Count neighbors for each cell
      std::map<ushort, uint> Nhbs;
      forall(const UShortUShortFloatPair& pr, WallLight) {
        Nhbs[pr.first.first]++;
        Nhbs[pr.first.second]++;
      }
  
      // Choose which labels to merge
      std::map<ushort, ushort> MergeMap;
      std::set<ushort> MergeSet;
  
      // Add cells with low light interfaces, handle when cells encased by a single cell
      forall(const UShortUShortFloatPair& pr, WallLight) {
        if(pr.first.first == 0 or pr.first.second == 0)
          continue;
        float wallLight = pr.second - CellIntLight[pr.first.first];
        float nbwallLight = pr.second - CellIntLight[pr.first.second];
        if(wallLight < 0)
          wallLight = 0;
        if(nbwallLight < 0)
          nbwallLight = 0;
  
        // Check if wall is unlikely from both directions
        bool chkout, chkin;
        UShortUShortPair rev(pr.first.second, pr.first.first);
        if(Nhbs[pr.first.first] == 1)
          chkout = wallLight / CellWallLight[pr.first.second] < tolerance;
        else
          chkout = wallLight / CellBordLight[pr.first] < tolerance;
  
        // If there is only one neighbor, CellBordLight will be zero
        if(Nhbs[pr.first.second] == 1)
          chkin = wallLight / CellWallLight[pr.first.first] < tolerance;
        else
          chkin = wallLight / CellBordLight[rev] < tolerance;
  
        if(chkin and chkout and MergeSet.find(pr.first.first) == MergeSet.end()) {
          MergeMap[pr.first.first] = pr.first.second;
          MergeSet.insert(pr.first.second);
        }
      }
  
      // Collapse regions
      for(uint i = 0; i < Labels.size(); i++) {
        if(Labels[i] > 0 and MergeMap[Labels[i]] > 0)
          Labels[i] = MergeMap[Labels[i]];
        if(i % size.x() * size.y() == 0 and !progressAdvance(0))
          userCancel();
      }
      regions = MergeSet.size();
      totregions += regions;
      SETSTATUS("Stack " << stack->userId() << " " << totregions << " regions consolidated");
    } while(regions > 0);
    labels->changed();
    return true;
  }
  
  REGISTER_PROCESS(ConsolidateRegionsNormalized);
  
  bool LocalMaximaStack::run(const Store* input, Store* labels, Point3f radius, uint startLabel, uint minColor, uint value)
  {
    uint label = startLabel;
    const Stack* stack = input->stack();
  
    HVecUS& output = labels->data();
    const HVecUS& data = input->data();
  
    HVecUS tdata(input->data().size());
    Point3i imgSize(stack->size());
    Point3i step(toVoxelsCeil(radius, Point3f(stack->step())));
    dilateGPU(imgSize, step, false, input->data(), tdata);
    ulong count = 0;
    for(ulong i = 0; i < data.size(); i++) {
      if(tdata[i] == 0)
        continue;
      if(tdata[i] == data[i] and data[i] >= minColor) {
        if(startLabel > 0)
          output[i] = label++;
        else
          output[i] = value;
        count++;
      } else
        output[i] = 0;
      if(startLabel > 0 and label >= uint(0xFFFF))
        throw QString("Too many local maxima to label");
    }
    SETSTATUS("Stack " << stack->userId() << " " << count << " pixels in local maxima");
    labels->changed();
    if(startLabel > 0)
      labels->setLabels(true);
    else
      labels->setLabels(false);
    labels->setFile("");
    return true;
  }
  REGISTER_PROCESS(LocalMaximaStack);
  
  bool ThresholdLabelDelete::run(const Store* input, Store* labels, uint minVoxels, uint maxVoxels)
  {
    const Stack* stack = input->stack();
  
    HVecUS& output = labels->data();
    const HVecUS& data = input->data();
  
    const uint LAB_SZ = 1 << 16;
    std::vector<uint> voxels(LAB_SZ, 0);
    for(uint i = 0; i < data.size(); i++)
      voxels[data[i]]++;
  
    for(uint i = 0; i < data.size(); i++) {
      ushort label = data[i];
      if(voxels[label] < minVoxels or (maxVoxels > 0 and voxels[label] > maxVoxels))
        output[i] = 0;
      else
        output[i] = label;
    }
  
    uint count = 0;
    for(uint i = 0; i < LAB_SZ; i++)
      if(voxels[i] > 0 and (voxels[i] < minVoxels or (maxVoxels > 0 and voxels[i] > maxVoxels)))
        count++;
  
    SETSTATUS("Stack " << stack->userId() << " " << count << " labels cleared");
    labels->changed();
    labels->setLabels(true);
    labels->setFile("");
    return true;
  }
  
  REGISTER_PROCESS(ThresholdLabelDelete);
  
  struct fillLabelCopy {
    fillLabelCopy(ushort tf, ushort nv, Store* out, uint ss = 0)
      : toFill(tf)
      , newValue(nv)
      , output(out->data())
      , slice_size(ss)
    {
      if(slice_size == 0)
        slice_size = output.size();
    }
  
    void operator()(const uint& idx)
    {
      BoundingBox3i bbox;
      uint start = idx * slice_size;
      uint end = std::min((idx + 1) * slice_size, (uint)output.size());
      for(uint i = start; i < end; ++i)
        if(output[i] == toFill)
          output[i] = newValue;
    }
  
    ushort toFill;
    ushort newValue;
    HVecUS& output;
    uint slice_size;
  };
  
  struct fillLabel {
    fillLabel(ushort tf, ushort nv, Store* st, uint ss = 0)
      : data(st->data())
      , stk(st->stack())
      , toFill(tf)
      , newValue(nv)
      , slice_size(ss)
    {
      if(slice_size == 0)
        slice_size = st->size();
    }
  
    BoundingBox3i operator()(uint idx) const
    {
      BoundingBox3i bbox;
      uint start = idx * slice_size;
      uint end = std::min((idx + 1) * slice_size, (uint)data.size());
      for(uint i = start; i < end; ++i) {
        if(data[i] == toFill) {
          data[i] = newValue;
          bbox |= Point3i(stk->position(i));
        }
      }
      return bbox;
    }
  
    HVecUS& data;
    const Stack* stk;
    ushort toFill;
    ushort newValue;
    uint slice_size;
    // BoundingBox3i bbox;
  };
  
  bool FillLabelStack::run(const Store* input, Store* output, ushort filledLabel, ushort newLabel)
  {
    int ns = QThread::idealThreadCount();
    uint nb_slices = (ns < 0 ? 128 : ns);
    uint size = output->size();
    uint slice = size / nb_slices;
    if(slice * nb_slices < size)
      slice++;
    if(input != output) {
      memcpy(&output->data()[0], &input->data()[0], sizeof(ushort) * size);
      thrust::for_each(THRUST_RETAG(thrust::counting_iterator<uint>(0)),
                       THRUST_RETAG(thrust::counting_iterator<uint>(nb_slices)),
                       fillLabelCopy(filledLabel, newLabel, output, slice));
      output->changed();
    } else {
      thrust::host_vector<BoundingBox3i> bboxes(nb_slices);
      thrust::transform(THRUST_RETAG(thrust::counting_iterator<uint>(0)),
                        THRUST_RETAG(thrust::counting_iterator<uint>(nb_slices)), THRUST_RETAG(bboxes.begin()),
                        fillLabel(filledLabel, newLabel, output, slice));
      BoundingBox3i bbox;
      for(size_t i = 0; i < bboxes.size(); ++i)
        bbox |= bboxes[i];
  
      output->changed(bbox);
    }
    output->copyMetaData(input);
    return true;
  }
  
  REGISTER_PROCESS(FillLabelStack);
  
  bool EraseAtBorderStack::run(const Store* input, Store* output)
  {
    const Stack* stk = output->stack();
    Point3u size = stk->size();
    const HVecUS& idata = input->data();
    HVecUS& data = output->data();
    std::vector<ushort> new_label(1 << 16);
  
  #pragma omp parallel for
    for(uint i = 0; i < 1 << 16; ++i)
      new_label[i] = ushort(i);
  
    progressStart("Erasing border labels", data.size());
  // First, process z planes
  #pragma omp parallel for
    for(uint y = 0; y < size.y(); ++y)
      for(uint x = 0; x < size.x(); ++x) {
        uint k1 = stk->offset(x, y, 0);
        uint k2 = stk->offset(x, y, size.z() - 1);
        new_label[idata[k1]] = 0;
        new_label[idata[k2]] = 0;
      }
  // First, process y planes
  #pragma omp parallel for
    for(uint z = 0; z < size.z(); ++z)
      for(uint x = 0; x < size.x(); ++x) {
        uint k1 = stk->offset(x, size.y() - 1, z);
        uint k2 = stk->offset(x, size.y() - 1, z);
        new_label[idata[k1]] = 0;
        new_label[idata[k2]] = 0;
      }
  // First, process x planes
  #pragma omp parallel for
    for(uint z = 0; z < size.z(); ++z)
      for(uint y = 0; y < size.y(); ++y) {
        uint k1 = stk->offset(size.x() - 1, y, z);
        uint k2 = stk->offset(size.x() - 1, y, z);
        new_label[idata[k1]] = 0;
        new_label[idata[k2]] = 0;
      }
  
  #pragma omp parallel for
    for(uint k = 0; k < data.size(); ++k)
      data[k] = new_label[idata[k]];
  
    output->changed();
    output->copyMetaData(input);
  
    return true;
  }
  REGISTER_PROCESS(EraseAtBorderStack);

  bool StackExportByLabel::run()
  {
    auto *stk = currentStack();
    if(!stk->currentStore()->labels())
      throw QString("%1: Stack must be labeled").arg(name());

    QString prefix = parm("File Prefix");

    std::map<int, BoundingBox<3, uint> > bBox;
    const HVecUS &data = stk->currentStore()->data();
    Point3u size = stk->size();
    for(uint z = 0; z < size.z(); z++)
      for(uint y = 0; y < size.y(); y++)
        for(uint x = 0; x < size.x(); x++)
          if(data[stk->offset(x,y,z)] > 0)
            bBox[data[stk->offset(x,y,z)]] |= Point3u(x,y,z);

    for(const auto &pr : bBox) {
      Point3u sz = pr.second.size() + Point3u(1,1,1);
      HVecUS labelData(sz.x()*sz.y()*sz.z());
      const Point3u &base = pr.second[0];
      for(uint z = 0; z < sz.z(); z++)
        for(uint y = 0; y < sz.y(); y++)
          for(uint x = 0; x < sz.x(); x++) {
            uint label = data[stk->offset(base.x()+x,base.y()+y,base.z()+z)];
            if(label == pr.first)
              labelData[offset(x,y,z,sz.x(),sz.y())] = label;
            else
              labelData[offset(x,y,z,sz.x(),sz.y())] = 0;
          }

      Image3D image(labelData, sz, stk->step(), stk->currentStore()->labels());
      saveTIFFImage(QString("%1-%2.tif").arg(prefix).arg(pr.first), image);
    }

    Information::out << "Labels found:" << bBox.size() << endl;

    return true;
  }
  REGISTER_PROCESS(StackExportByLabel);


}
