//
// 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 <MeshProcessCurv.hpp>
#include <Information.hpp>
#include <Progress.hpp>
#include <Curvature.hpp>
#include <vector>
#include <QFile>
#include <QTextStream>
#include <Misc.hpp>
#include <algorithm>
#include <SetVector.hpp>
#include <GraphUtils.hpp>

using std::swap;

namespace mgx
{
  struct ComputeTissueCurvatureKernel {
    ComputeTissueCurvatureKernel(const vvGraph& _S, const std::vector<vertex>& _cells, float nh2,
                          std::vector<Curvature>& cs, std::vector<set_vector<vertex> >& _vert)
      : S(_S), cells(_cells), neighborhood_sq(nh2), curvs(cs), vert(_vert) {}
  
    void operator()(int ci)
    {
      const vertex& center = cells[ci];
      std::vector<Point3f> points, normals;
  
      // the cell center has to be the first element of points and normals
      // it will be used to calculate the local coordinate system in curvature computation
      points.push_back(Point3f(center->pos));
      normals.push_back(Point3f(center->nrml));
  
      set_vector<vertex> neighCenters;
      set_vector<vertex> neighVertices;
      set_vector<vertex> newJunctions;
      neighCenters.insert(center);
      do {
        newJunctions.clear();
        // in neighboring cells, look for all the junctions within neighborhood radius
        forall(const vertex& c, neighCenters)
          forall(const vertex& v, S.neighbors(c)) {
            if(normsq(v->pos - center->pos) > neighborhood_sq or neighVertices.count(v) != 0)
              continue;
            neighVertices.insert(v);
            newJunctions.insert(v);
          }
  
        // look for next neighbors cells connected to junctions
        set_vector<vertex> newCenters;
        forall(const vertex& j, newJunctions)
          forall(const vertex& v, S.neighbors(j))
            if(v->label > 0 and neighCenters.count(v) == 0)
              newCenters.insert(v);
  
        swap(neighCenters, newCenters);
      } while(newJunctions.size() != 0);
  
      forall(vertex v, neighVertices) {
        if(v->label > 0)
          continue;
        points.push_back(Point3f(v->pos));
        normals.push_back(Point3f(v->nrml));
      }
  
      if(points.size() > 3)
        curvs[ci].update(points, normals);
  
      // store all the vertices used to compute the curvature (center + neighboring junctions)
      // this is used when we want to check how the computation was made (checkLabel option)
      neighVertices.insert(center);
      vert[ci] = neighVertices;
    }
  
    const vvGraph& S;
    const std::vector<vertex>& cells;
    float neighborhood_sq;
    std::vector<Curvature>& curvs;
    std::vector<set_vector<vertex> >& vert;
  };
  
  bool TissueCurvature::run(Mesh* mesh, float neighborhood, bool checkLabel)
  {
  
    AttrMap<int, SymmetricTensor>& curvAttr = mesh->attributes().attrMap<int, SymmetricTensor>("Measure Label Tensor Curvature");
    curvAttr.clear();

    if(neighborhood <= 0)
      throw QString("Error, neighborhood argument must be strictly positive.");
    if(checkLabel and !mesh->showMesh() and selectedLabel() > 0)
      throw QString("The mesh must be visible. Vertices used to compute curvature "
                    "for chosen label will be selected.");
   
    // Leave the original mesh is unchanged unless we use the option "checkLabel" to 
		// check which vertices are used for computation
    vvGraph T = mesh->graph();
    if(mesh->meshType() == "MGXM")
      if(!mgxmToMgxc(T, T, neighborhood / 10.0))
				throw(QString("Error converting to cell mesh"));

    // First extract the centers of the cells from the mesh surface.
    // We will use these centers to compute curvature, because the centers given by
    // mgxmToMgxc tend to be away from the surface for large cells in curved tissues.
    mesh->clearCellAxis();
    mesh->updateCentersNormals();
    IntPoint3fAttr& LabelCenter = mesh->labelCenterVis();
    //IntPoint3fAttr& LabelCenter = (mesh->useParents() ? mesh->parentCenterVis() : mesh->labelCenterVis());// to make it work on parents


    // The curvature computation is done on each cell center, but takes into account only
    // cell borders
    VtxVec cells;
  
    forall(const vertex& c, T)
      if(c->label > 0) {
        c->pos = Point3d(LabelCenter[c->label]);
        cells.push_back(c);
      }
    std::vector<Curvature> curvs(cells.size());
    std::vector<set_vector<vertex> > vert(cells.size());
  
    float neighborhood_sq = neighborhood * neighborhood;
  
    ComputeTissueCurvatureKernel op(T, cells, neighborhood_sq, curvs, vert);
    #pragma omp parallel for
    for(size_t i = 0; i < cells.size(); i++)
      op(i);
  
    // Select the vertices used for curvature computation in a particular cell (checkLabel)
    set_vector<vertex> vertLabel;
    if(checkLabel == true and selectedLabel() > 0) {
      for(uint i = 0; i < cells.size(); ++i) {
        const vertex& c = cells[i];
        if(c->label == selectedLabel())
          vertLabel = vert[i];
      }
      // select vertices used in ComputeTissueCurvatureKernel for particular label
      forall(const vertex& v, mesh->graph()) {
        v->selected = false;
        if(std::find(vertLabel.begin(), vertLabel.end(), v) != vertLabel.end())
          v->selected = true;
      }
      mesh->updateSelection();
    }
  
    // fill in the cell axis for display of curvature directions
    IntSymTensorAttr& cellTensors = mesh->cellAxis();
  
    for(uint i = 0; i < cells.size(); ++i) {
      const vertex& c = cells[i];
      Curvature& curv = curvs[i];
      Point3f e1, e2;
      float c1, c2;
      curv.get_eigenVectors(e1, e2);
      curv.get_eigenValues(c1, c2);
      // make sure the largest curvature (norm) is MaxCurv
      if(fabs(c2) > fabs(c1)) {
        std::swap(c1, c2);
        std::swap(e1, e2);
      }
      // Store as a 3D tensor with the third component to 0
      cellTensors[c->label].ev1() = e1;
      cellTensors[c->label].ev2() = e2;
      cellTensors[c->label].evals() = Point3f(c1, c2, 0);

      curvAttr[c->label] = cellTensors[c->label];
    }
  
    mesh->setCellAxisType("curvature");
    mesh->setCellAxisUnit(QString("1/") + UM);
  
    // Run curvature display, with the parameters from the GUI
    DisplayTissueCurvature *dcurv;
	if(!getProcess("Mesh/Cell Axis/Curvature/Display Tissue Curvature", dcurv))  
      throw(QString("MeshProcessHeatMap:: Unable to make DisplayTissueCurvature process")); 
    return dcurv->run();
    
  }
  REGISTER_PROCESS(TissueCurvature);
  
  // Display curvature axis (separatly from each other) and heatmaps of gaussian curv etc.
  bool DisplayTissueCurvature::run(Mesh* mesh, QString displayHeatMap, bool symmetricScale,
	  float heatmapPercentile, QString displayCurv, const QColor& qColorPositive, const QColor& qColorNegative,
		float axisLineWidth, float axisLineScale, float axisOffset, float curvThreshold)
  
  {
    // if there is no cell axis stored, display error message
    const IntSymTensorAttr& cellAxis = mesh->cellAxis();
    if(cellAxis.size() == 0 or mesh->cellAxisType() != "curvature") {
      setErrorMessage(QString("No curvature cell axis stored in active mesh!"));
      return false;
    }
    if(heatmapPercentile < 1)
      heatmapPercentile = 1;
    else if(heatmapPercentile > 100)
      heatmapPercentile = 100;
  
  
    // Get attribute maps
    IntPoint3fAttr& LabelCenter = mesh->labelCenterVis();
    IntMatrix3fAttr& cellAxisVis = mesh->cellAxisVis();
    IntVec3ColorbAttr& cellAxisColor = mesh->cellAxisColor();

    // Clear cell Axis
    cellAxisVis.clear();
    cellAxisColor.clear();

    // Scale curvature for visualization
    mesh->setAxisWidth(axisLineWidth);
    mesh->setAxisOffset(axisOffset);
    Colorb colorPositive(qColorPositive);
    Colorb colorNegative(qColorNegative);

    // Check if there is a cell center for each cell axis
    int nProblems = 0;
    forall(const IntSymTensorPair& p, cellAxis) {
      int label = p.first;
      if(LabelCenter.count(label) == 0)
        nProblems++;
    }
    if(nProblems != 0)
      SETSTATUS("Warning: non-existing cell center found for " << nProblems << " cell axis.");

    // Populate cellAxisVis for display

    forall(const IntSymTensorPair& p, cellAxis) {
      int cell = p.first;
      const SymmetricTensor& tensor = p.second;
      Point3b showAxis;
      showAxis[0] = (displayCurv == "CurvMax" or displayCurv == "Both");
      showAxis[1] = (displayCurv == "CurvMin" or displayCurv == "Both");
      showAxis[2] = false;

      // if curvature lower than threshold, don't display axis
      if(curvThreshold != 0 and curvThreshold > (fabs(tensor.evals()[0]) + fabs(tensor.evals()[1])) / 2)
        showAxis[0] = showAxis[1] = false;

      cellAxisVis[cell][0] = (showAxis[0] ? tensor.evals()[0] : 0.f) * tensor.ev1() * axisLineScale;
      cellAxisVis[cell][1] = (showAxis[1] ? tensor.evals()[1] : 0.f) * tensor.ev2() * axisLineScale;
      cellAxisVis[cell][2] = (showAxis[2] ? tensor.evals()[2] : 0.f) * tensor.ev3() * axisLineScale;

      for(size_t i = 0; i < 3; ++i) {
        if(tensor.evals()[i] < 0)
          cellAxisColor[cell][i] = colorNegative;
        else
          cellAxisColor[cell][i] = colorPositive;
      }
    }
    mesh->setShowAxis("Cell Axis");
  
    displayHeatMap = displayHeatMap.toLower();

    if(displayHeatMap != "none") {
      IntFloatAttr& labelHeatMap = mesh->labelHeat();
      labelHeatMap.clear();
      std::vector<float> values;
      values.reserve(cellAxis.size());
  
      Information::out << "displayHeatMap = '" << displayHeatMap << "'\n";
  
      CurvatureMeasure* m = 0;
      QString unit1 = UM_1;
      QString unit2 = UM_2;
      QString unit;
  
      if(displayHeatMap == "curvmax") {
        m = new MaximalCurvature();
        unit = unit1;
      } else if(displayHeatMap == "curvmin") {
        m = new MinimalCurvature();
        unit = unit1;
      } else if(displayHeatMap == "gaussian") {
        m = new GaussianCurvature();
        unit = unit2;
      } else if(displayHeatMap == "rootsumsquare") {
        m = new RootSumSquareCurvature();
        unit = unit2;
      } else if(displayHeatMap == "average") {
        m = new AverageCurvature();
        unit = unit1;
      } else if(displayHeatMap == "signedaverageabs") {
        m = new SignedAverageAbsCurvature();
        unit = unit1;
      } else if(displayHeatMap == "anisotropy") {
        m = new AnisotropyCurvature();
        unit = "";
      } else {
        return setErrorMessage(QString("Error, unknown curvature measure '%1'").arg(displayHeatMap));
      }
  
      // Fill in heat map
      forall(const IntSymTensorPair& p, cellAxis) {
        const SymmetricTensor& tensor = p.second;
        int cell = p.first;
        float curvMax = tensor.evals()[0];
        float curvMin = tensor.evals()[1];
        float value = m->measure(curvMax, curvMin);
  
        labelHeatMap[cell] = value;
        values.push_back(value);
      }
      delete m;
      // Rescale heat map
      sort(values.begin(), values.end());
      int rankmax = values.size() - 1;
      int rankmin = 0;
      if(heatmapPercentile < 100) {
        rankmax = (values.size() * heatmapPercentile) / 100;
        rankmin = values.size() - rankmax;
      }
  
      Information::out << "rankMin = " << rankmin << " - rankMax = " << rankmax << endl;
  
      float maxValue = values[rankmax];
      float minValue = values[rankmin];
  
      Information::out << "minValue = " << minValue << " - maxValue = " << maxValue << endl;
  
      if(symmetricScale) {
        if(displayHeatMap == "anisotropy") {
          if(maxValue > 1)
            maxValue = 2;
          else
            maxValue = 1;
          if(minValue < 1)
            minValue = 0;
          else
            minValue = 1;
        } else {
          if(maxValue <= 0)
            maxValue = 0;
          else if(minValue >= 0)
            minValue = 0;
          else
            minValue = -maxValue;
        }
      }
  
      mesh->heatMapBounds() = Point2f(minValue, maxValue);
      mesh->heatMapUnit() = unit;
      mesh->setShowLabel("Label Heat");
      if(displayCurv != "None")
        mesh->setShowAxis("Cell Axis");
      mesh->updateTriangles();
      Information::out << "Done!" << endl;
    }
  
    return true;
  }
  REGISTER_PROCESS(DisplayTissueCurvature);
}
