//
// 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 <MeshProcessSignalOrientation.hpp>
#include <Information.hpp>
#include <Progress.hpp>
#include <GraphUtils.hpp>

#include <PCA.hpp>

namespace mgx
{
  bool SignalOrientation::run(Mesh* mesh, float border, float minAreaRatio)
  {
    IntPoint3fAttr Normals;
    IntPoint3fAttr Centers;
    IntFloatMap InsideAreas;
    IntFloatMap InsideSignals;
    IntFloatMap CellAreas;
    IntFloatMap CellSignals;
    IntFloatMap WeightedBorder;
  
    IntPoint3fMap FirstDir;
    std::unordered_map<int, Matrix3f> CorrSignals;
    float totalAvSignal;
    std::unordered_map<int, std::unordered_map<float, Point3f> > CellAnglesSignal;
  
    const vvGraph& S = mesh->graph();

    mesh->updateCentersNormals();
    Normals = mesh->labelNormal(); 
    Centers = mesh->labelCenter(); 
    //Normals = mesh->labelNormalVis(); 
    //Centers = mesh->labelCenterVis(); 
    
    // Mark the vertices within this distance from the border.
    mesh->markBorder(border);
  
    // Find areas, based on labels
    forall(const vertex &v, S) {
      forall(const vertex &n, S.neighbors(v)) {
        const vertex& m = S.nextTo(v, n); 
        if(!S.uniqueTri(v, n, m))
          continue; 
        int label = getLabel(v, n, m);
        if(label <= 0)
          continue;
  
        float area = triangleArea(v->pos, n->pos, m->pos);
        CellAreas[label] += area;
        CellSignals[label] += (v->signal + n->signal + m->signal)/3.0;
        if(n->minb == 0 and v->minb == 0 and m->minb == 0){
         InsideAreas[label] += area; 
         InsideSignals[label] += (v->signal + n->signal + m->signal)/3.0; 
        }
      }
    }

    std::cout << "pol2" << std::endl;
  
    // Normalize the signal / area. This internal signal will be assigned to cell center. 
    forall(const IntFloatPair &p, InsideSignals){ 
      InsideSignals[p.first] /= InsideAreas[p.first]; 
      CellSignals[p.first] /= CellAreas[p.first]; 
      totalAvSignal += CellSignals[p.first];
    }
    totalAvSignal /= InsideSignals.size();
  
    // Clear data in MorphoGraphX
    mesh->clearCellAxis();
  
    // Calculate orientation of vectors center/cell wall, weighted by signal.
    forall(const vertex &v, S) {
      if(v->minb == 0)
        continue;
      forall(const vertex &n, S.neighbors(v)) {
        if(n->minb == 0)
          continue;
        const vertex& m = S.nextTo(v, n); 
        if(m->minb == 0 or !S.uniqueTri(v, n, m))
          continue; 
        int label = getLabel(v, n, m);
        if(label <= 0)
          continue;
        // if the total area of the triangles is too small, don't compute the PO
        if(InsideAreas[label] < CellAreas[label]*minAreaRatio)
          continue;
  
        // the triangle (v,n,m) belongs to the border
        float borderSignal = (v->signal + n->signal + m->signal)/3.0; 
        Point3f borderPosition((v->pos + n->pos + m->pos)/3.0); 
        float borderArea = triangleArea(v->pos, n->pos, m->pos);
        // gradient of signal border vs center
        float diffSignal = borderSignal/borderArea-InsideSignals[label];
        // unit vector giving the direction center-border triangle, within the cell plane
        Point3f direction = (borderPosition - Centers[label]);
        // if border signal is smaller than center signal, inverse direction of unit vector
        if(diffSignal < 0)
          direction = -direction;
        direction = direction - (direction * Normals[label]) * Normals[label];
        direction /= norm(direction);// unit vector 
        // take first unit center-border vector as reference direction 
        if(FirstDir.find(label) == FirstDir.end())
          FirstDir[label] = direction;          

        // compute the angle center-border, with respect to dirX and dirY
        // dirX: first direction found for this label, dirY: normal to dirX
        Point3f dirX = FirstDir[label];
        Point3f dirY = Normals[label] ^ FirstDir[label];
        dirY /= norm(dirY); 
        float angle = atan2(direction * dirY, direction * dirX);// in radian, [-pi,pi]
        angle = 360*(angle + M_PI)/(2*M_PI);// defined for [0, 360] degrees
        CellAnglesSignal[label].insert(std::pair<float, Point3f>(angle, fabs(diffSignal)*direction));
        std::cout << "angle " << angle << "/" << label << std::endl;
      }
    }

    forall(const IntPoint3fPair &p, FirstDir){
      int label = p.first;
      std::unordered_map<float, Point3f> AnglesSignal = CellAnglesSignal[label];
      std::unordered_map<int, Point3f> AnglesSignalBin;
      std::unordered_map<int, Point3f> AnglesWeightBin;
  
      // sum up the signal for all directions within regularly spaced angles
      // angle spacing = 360/nBins
      // vector belongs to bin n if angle between (n*360)/nBins and (n+1)*360/nBins)
      int nBins = 24; 
      float spacing = 360/nBins;
      for(std::unordered_map<float, Point3f>::iterator it=AnglesSignal.begin(); it!=AnglesSignal.end(); ++it){
        float angle = it->first;
        int bin = floor(angle/spacing);
        Point3f vector = it->second;
        if(norm(vector) == 0 )
          continue;
        AnglesSignalBin[bin] += vector;
        AnglesWeightBin[bin] += vector/norm(vector);
      }

      forall(const IntPoint3fPair &p, AnglesSignalBin){
        Point3f signalBin = AnglesSignalBin[p.first]/norm(AnglesWeightBin[p.first]);
        // correlation matrix for the averaged direction * signal
        Matrix3f corrSig;
        for(int i = 0; i < 3; i++)
          for(int j = 0; j < 3; j++)
              corrSig(i,j) = signalBin[i] * signalBin[j];
  
        CorrSignals[label] += corrSig;
      }
  
    }
  
    IntSymTensorAttr &cellAxis = mesh->cellAxis();

    AttrMap<int, SymmetricTensor>& polAttr = mesh->attributes().attrMap<int, SymmetricTensor>("Measure Label Tensor Polarity");
    polAttr.clear();

    forall(const IntFloatPair &p, InsideSignals) {
      Matrix3f corrSignals = CorrSignals[p.first]/pow(totalAvSignal,2);// normalize by signal strength

      // decompose the correlation matrix with pca
      PCA pcaSig;
      if(!pcaSig(corrSignals)) 
        return setErrorMessage("Error performing PCA");
  
       // store the eigen-vectors and values in the cell axis   
      SymmetricTensor& tensor = cellAxis[p.first];
      tensor.ev1() = pcaSig.p1;
      tensor.ev2() = pcaSig.p2;
      tensor.evals() = pcaSig.ev;

      polAttr[p.first] = tensor;
    }
   
    mesh->setCellAxisType("polarization");
    mesh->setCellAxisUnit("a.u.");
    mesh->updateTriangles();
  
    // Run polarization display, with the parameters from the GUI 
    
    
    DisplaySignalOrientation *proc;
    if(!getProcess("Mesh/Cell Axis/Polarization/Display Signal Orientation", proc))  
      throw(QString("MeshProcessHeatMap:: Unable to make DisplaySignalOrientation process")); 
    proc->run();       

    return true;
  }
  REGISTER_PROCESS(SignalOrientation);
  
  // Display principal orientations (separatly from each other) and heatmaps of anisotropy etc.
  bool DisplaySignalOrientation::run(Mesh *mesh,const QString displayHeatMap,const QString displayAxis,
         const QColor& qaxisColor1, const QColor& qaxisColor2, float axisLineWidth, float axisLineScale, 
         float axisOffset, float orientationThreshold)                                   
  {
    // if there is no cell axis stored, display error message
    const IntSymTensorAttr& cellAxis = mesh->cellAxis();
    if(cellAxis.size() == 0 or mesh->cellAxisType() != "polarization"){
      setErrorMessage(QString("No orientation axis stored in active mesh!"));
      return false;
    }
    const IntSymTensorAttr& cellAxisBezier = mesh->attributes().attrMap<int, SymmetricTensor>("Cell Axis Bezier");

    // Check cell axis and prepare them for display. Populate cellAxisVis attribute  
    IntMatrix3fAttr& cellAxisVis = mesh->cellAxisVis();
    IntVec3ColorbAttr& cellAxisColor = mesh->cellAxisColor();
    cellAxisVis.clear();
    cellAxisColor.clear();
    mesh->setAxisWidth(axisLineWidth);
    mesh->setAxisOffset(axisOffset);
    Colorb axisColor1(qaxisColor1);
    Colorb axisColor2(qaxisColor2);

    // Set the color
  
    // Update the normals and center of the cell for visualisation.
    if(mesh->labelCenterVis().empty() or mesh->labelNormalVis().empty())
      mesh->updateCentersNormals();
  
    // Check if there is a cell center for each cell axis
    IntPoint3fAttr labelCenterVis = mesh->labelCenterVis();
    int nProblems = 0;
    forall(const IntSymTensorPair &p, cellAxis) {
      int label = p.first;
      if(labelCenterVis.count(label) == 0)
        nProblems++;
    }
    if(nProblems != 0)
      SETSTATUS("Warning: non-existing cell center found for " << nProblems << " cell axis.");
  
    // Scale the cell axis for display
    if(displayAxis != "None") {
      forall(const IntSymTensorPair &p, cellAxis) {
        int cell = p.first;
        SymmetricTensor tensor = p.second;

        if(displayAxis == "BezierX" or displayAxis == "BezierY") {
          if(cellAxisBezier.size() == 0)
            throw(QString("No cell axis bezier. Run 'Cell Axis Custom Directions' first."));
          if(cellAxisBezier.find(cell) == cellAxisBezier.end())
            continue;
          tensor = cellAxisBezier.find(cell)->second;
          if(displayAxis == "BezierX")
            cellAxisVis[cell][0] = tensor.ev1() * tensor.evals()[0] * axisLineScale;
          if(displayAxis == "BezierY")
            cellAxisVis[cell][1] = tensor.ev2() * tensor.evals()[1] * axisLineScale;
        }

        float ratio =  tensor.evals()[0]/tensor.evals()[1]-1;
        float difference =  tensor.evals()[0]-tensor.evals()[1];
        if(displayAxis == "Both") {
          cellAxisVis[cell][0] = tensor.ev1() * tensor.evals()[0] * axisLineScale;
          cellAxisVis[cell][1] = tensor.ev2() * tensor.evals()[1] * axisLineScale;
          cellAxisVis[cell][2] = Point3f(0,0,0);
        } else if(displayAxis == "Ratio") {
          float orientation = ratio;
          cellAxisVis[cell][0] = tensor.ev1() * (orientation > orientationThreshold ? orientation : 0.f) * axisLineScale;
          cellAxisVis[cell][1] = Point3f(0,0,0);
          cellAxisVis[cell][2] = Point3f(0,0,0);
        } else if(displayAxis == "Difference") {
          float orientation = difference;
          cellAxisVis[cell][0] = tensor.ev1() * (orientation > orientationThreshold ? orientation : 0.f) * axisLineScale;
          cellAxisVis[cell][1] = Point3f(0,0,0);
          cellAxisVis[cell][2] = Point3f(0,0,0);
        }
        cellAxisColor[cell][0] = axisColor1;
        cellAxisColor[cell][1] = axisColor2;
      }
    }

    // Heat map of orientation
    if(displayHeatMap != "None") {
      IntFloatAttr labelHeatMap;
      mesh->labelHeat().clear();
      std::vector<float> values;
      values.reserve(cellAxis.size());
      forall(const IntSymTensorPair &p, cellAxis)
      {
        float value;
        int cell = p.first;
        if(displayHeatMap == "Ratio")
          value = p.second.evals()[0] / p.second.evals()[1];
        else if(displayHeatMap == "Difference")
          value = p.second.evals()[0] - p.second.evals()[1];
        else if(displayHeatMap == "BezierX" or displayHeatMap == "BezierY") {
          if(cellAxisBezier.size() == 0)
            throw(QString("No cell axis bezier. Run 'Cell Axis Custom Directions' first."));
          if(cellAxisBezier.find(cell) == cellAxisBezier.end())
            continue;
          SymmetricTensor tensor = cellAxisBezier.find(cell)->second;
          if(displayAxis == "BezierX")
            value = tensor.evals()[0];
          else if(displayAxis == "BezierY")
            value = tensor.evals()[1];
        }

        labelHeatMap[cell] = value;
        values.push_back(value);
      }
      mesh->labelHeat() = labelHeatMap;
      // Rescale heat map
      sort(values.begin(), values.end());
      int rankmax = values.size() - 1;
      int rankmin = 0;
      int heatmapPercentile = 90;// we could make that a parameter
      if(heatmapPercentile < 100) {
        rankmax = (values.size() * heatmapPercentile) / 100;
        rankmin = values.size() - rankmax;
      }
      float maxValue = values[rankmax];
      float minValue = values[rankmin];
  
//      // Adjust heat map if run on parents. TO CHANGE AFTER HEATMAP REFACTORING
//      if(mesh->parents().size() != 0){
//        Information::out << "parent label loaded" << endl;
//        mesh->labelHeat().clear();
//        IntIntAttr & parentMap = mesh->parents();
//        forall(const IntIntPair &p, parentMap)
//          if(p.second > 0 and labelHeatMap.count(p.second) != 0)
//            mesh->labelHeat()[cell] = labelHeatMap[p.second];
//      }
//      else
//        mesh->labelHeat() = labelHeatMap;
  
      mesh->heatMapBounds() = Point2f(minValue, maxValue);
      if(displayHeatMap == "Ratio")
        mesh->heatMapUnit() = QString("Ratio");
      if(displayHeatMap == "Difference")
        mesh->heatMapUnit() = QString("signal/area");
      mesh->setShowLabel("Label Heat");
      mesh->updateTriangles();
    }
  
    return true;
  }
  REGISTER_PROCESS(DisplaySignalOrientation);


  bool SignalOrientation3D::run(Mesh* mesh)
  {
    //IntPoint3fAttr Normals;
    //IntPoint3fAttr Centers;
    IntFloatMap CellAreas;
    IntFloatMap CellSignals;
    IntFloatMap WeightedBorder;
  
    IntPoint3fMap FirstDir;
    std::unordered_map<int, Matrix3f> CorrSignals;
    float totalAvSignal;
    std::unordered_map<int, std::unordered_map<std::pair<float, float>, Point3f> > CellAnglesSignal;
  
    Point3f xDir(1.,0.,0.);
    Point3f yDir(0.,1.,0.);
    Point3f zDir(0.,0.,1.);

    const vvGraph& S = mesh->graph();

    AttrMap<int, Point3d>& centroids = mesh->attributes().attrMap<int, Point3d>("Measure Label Vector CellCentroids");

    IntPoint3fAttr &Centers = mesh->labelCenterVis();
    Centers.clear();

    forall(auto p, centroids){
      Centers[p.first] = Point3f(p.second);
    }
  
    // // Find areas, based on labels
    forall(const vertex &v, S) {
      forall(const vertex &n, S.neighbors(v)) {
        const vertex& m = S.nextTo(v, n); 
        if(!S.uniqueTri(v, n, m))
          continue; 
        int label = getLabel(v, n, m);
        if(label <= 0)
          continue;
  
        float area = triangleArea(v->pos, n->pos, m->pos);
        CellAreas[label] += area;
        CellSignals[label] += (v->signal + n->signal + m->signal)/3.0;
      }
    }

  
    // Normalize the signal / area. This internal signal will be assigned to cell center. 
    forall(const IntFloatPair &p, CellSignals){ 
      CellSignals[p.first] /= CellAreas[p.first]; 
      totalAvSignal += CellSignals[p.first];
    }
    totalAvSignal /= CellSignals.size();
  
    // Clear data in MorphoGraphX
    mesh->clearCellAxis();
  
    // Calculate orientation of vectors center/cell wall, weighted by signal.
    forall(const vertex &v, S) {
      forall(const vertex &n, S.neighbors(v)) {
        const vertex& m = S.nextTo(v, n); 
        if(!S.uniqueTri(v, n, m))
          continue; 
        int label = getLabel(v, n, m);
        if(label <= 0)
          continue;
        // if the total area of the triangles is too small, don't compute the PO
        //if(InsideAreas[label] < CellAreas[label]*minAreaRatio)
        //  continue;
  
        // the triangle (v,n,m) belongs to the border
        float borderSignal = (v->signal + n->signal + m->signal)/3.0; 
        Point3f borderPosition((v->pos + n->pos + m->pos)/3.0); 
        float borderArea = triangleArea(v->pos, n->pos, m->pos);
        // gradient of signal border vs center
        float diffSignal = borderSignal/borderArea-CellSignals[label];
        // unit vector giving the direction center-border triangle, within the cell plane
        Point3f direction = (borderPosition - Centers[label]);
        // if border signal is smaller than center signal, inverse direction of unit vector
        if(diffSignal < 0)
          direction = -direction;
        //direction = direction - (direction * Normals[label]) * Normals[label];
        direction /= norm(direction);// unit vector 
        // take first unit center-border vector as reference direction 
        if(FirstDir.find(label) == FirstDir.end())
          FirstDir[label] = direction;          

        // compute the angle center-border, with respect to dirX and dirY
        // dirX: first direction found for this label, dirY: normal to dirX
        Point3f dirX = xDir;// FirstDir[label];
        Point3f dirY = yDir; //Normals[label] ^ FirstDir[label];
        Point3f dirZ = zDir;
        // dirY /= norm(dirY); 
        float angle = atan2(direction * dirY, direction * dirX);// in radian, [-pi,pi]
        angle = 360*(angle + M_PI)/(2*M_PI);// defined for [0, 360] degrees
        float angle2 = atan2(direction * dirZ, direction * dirY);// in radian, [-pi,pi]
        angle2 = 360*(angle2 + M_PI)/(2*M_PI);// defined for [0, 360] degrees
        std::pair<float, float> angleP = std::make_pair(angle,angle2);
        std::pair<std::pair<float, float>, Point3f> signalMapP = std::make_pair(angleP,fabs(diffSignal)*direction);
        CellAnglesSignal[label].insert(signalMapP);



      }
    }

    forall(auto p, CellAnglesSignal){
      int label = p.first;
      std::unordered_map<std::pair<float, float>, Point3f> AnglesSignal = CellAnglesSignal[label];
      std::unordered_map<std::pair<int, int>, Point3f> AnglesSignalBin;
      std::unordered_map<std::pair<int, int>, Point3f> AnglesWeightBin;
  
      // sum up the signal for all directions within regularly spaced angles
      // angle spacing = 360/nBins
      // vector belongs to bin n if angle between (n*360)/nBins and (n+1)*360/nBins)
      int nBins = 24; 
      float spacing = 360/nBins;
      for(std::unordered_map<std::pair<float, float>, Point3f>::iterator it=AnglesSignal.begin(); it!=AnglesSignal.end(); ++it){
        float angle1 = it->first.first;
        float angle2 = it->first.second;
        int bin1 = floor(angle1/spacing);
        int bin2 = floor(angle2/spacing);
        Point3f vector = it->second;
        if(norm(vector) == 0 )
          continue;
        std::pair<int, int> binP = std::make_pair(bin1, bin2);
        AnglesSignalBin[binP] += vector;
        AnglesWeightBin[binP] += vector/norm(vector);
      }

      forall(auto p, AnglesSignalBin){
        Point3f signalBin = AnglesSignalBin[p.first]/norm(AnglesWeightBin[p.first]);
        // correlation matrix for the averaged direction * signal
        Matrix3f corrSig;
        for(int i = 0; i < 3; i++)
          for(int j = 0; j < 3; j++)
            corrSig(i,j) = signalBin[i] * signalBin[j];
  
        CorrSignals[label] += corrSig;
      }
  
    }
  
    AttrMap<int, SymmetricTensor>& polAttr = mesh->attributes().attrMap<int, SymmetricTensor>("Measure Label Tensor Polarity3D");
    polAttr.clear();

    IntSymTensorAttr &cellAxis = mesh->cellAxis();
    forall(const IntFloatPair &p, CellSignals) {
      Matrix3f corrSignals = CorrSignals[p.first]/pow(totalAvSignal,2);// normalize by signal strength

      // decompose the correlation matrix with pca
      PCA pcaSig;
      if(!pcaSig(corrSignals)) 
        return setErrorMessage("Error performing PCA");
  
      // store the eigen-vectors and values in the cell axis   
      SymmetricTensor& tensor = cellAxis[p.first];
      tensor.ev1() = pcaSig.p1;
      tensor.ev2() = pcaSig.p2;
      tensor.evals() = pcaSig.ev;

      polAttr[p.first] = tensor;
    }
   
    mesh->setCellAxisType("polarization");
    mesh->setCellAxisUnit("a.u.");
    mesh->updateTriangles();
  
    // Run polarization display, with the parameters from the GUI 
    
    
    DisplaySignalOrientation3D *proc;
    if(!getProcess("Mesh/Cell Axis 3D/Polarization/Display Signal Orientation", proc))  
      throw(QString("MeshProcessHeatMap:: Unable to make DisplaySignalOrientation process")); 
    proc->run();      

    return true;
  }
  REGISTER_PROCESS(SignalOrientation3D);


  // Display principal orientations (separatly from each other) and heatmaps of anisotropy etc.
  bool DisplaySignalOrientation3D::run(Mesh *mesh,const QString displayHeatMap,const QString displayAxis,
         const QColor& qaxisColor1, const QColor& qaxisColor2, const QColor& qaxisColor3, float axisLineWidth, float axisLineScale, 
         float orientationThreshold)                                   
  {
    // if there is no cell axis stored, display error message
    const IntSymTensorAttr& cellAxis = mesh->cellAxis();
    if(cellAxis.size() == 0 or mesh->cellAxisType() != "polarization"){
      setErrorMessage(QString("No orientation axis stored in active mesh!"));
      return false;
    }
    const IntSymTensorAttr& cellAxisBezier = mesh->attributes().attrMap<int, SymmetricTensor>("Measure Label Tensor CustomDirections");

    AttrMap<int, Point3d>& centroids = mesh->attributes().attrMap<int, Point3d>("Measure Label Vector CellCentroids");

    IntPoint3fAttr &centers = mesh->labelCenterVis();
    centers.clear();

    forall(auto p, centroids){
      centers[p.first] = Point3f(p.second);
    }

    // Check cell axis and prepare them for display. Populate cellAxisVis attribute  
    IntMatrix3fAttr& cellAxisVis = mesh->cellAxisVis();
    IntVec3ColorbAttr& cellAxisColor = mesh->cellAxisColor();
    cellAxisVis.clear();
    cellAxisColor.clear();
    mesh->setAxisWidth(axisLineWidth);
    mesh->setAxisOffset(0.0);
    Colorb axisColor1(qaxisColor1);
    Colorb axisColor2(qaxisColor2);
    Colorb axisColor3(qaxisColor3);

    // Set the color
  
    // // Update the normals and center of the cell for visualisation.
    // if(mesh->labelCenterVis().empty() or mesh->labelNormalVis().empty())
    //   mesh->updateCentersNormals();
  
    // // Check if there is a cell center for each cell axis
    // IntPoint3fAttr labelCenterVis = mesh->labelCenterVis();
    // int nProblems = 0;
    // forall(const IntSymTensorPair &p, cellAxis) {
    //   int label = p.first;
    //   if(labelCenterVis.count(label) == 0)
    //     nProblems++;
    // }
    // if(nProblems != 0)
    //   SETSTATUS("Warning: non-existing cell center found for " << nProblems << " cell axis.");
  
    // Scale the cell axis for display
    if(displayAxis != "None"){
      forall(const IntSymTensorPair &p, cellAxis) {
        int cell = p.first;
        SymmetricTensor tensor = p.second;

        if(displayAxis == "BezierX" or displayAxis == "BezierY"){
          if(cellAxisBezier.size() == 0)
            throw(QString("No cell axis bezier. Run 'Cell Axis Custom Directions' first."));
          if(cellAxisBezier.find(cell) == cellAxisBezier.end())
            continue;
          tensor = cellAxisBezier.find(cell)->second;
          if(displayAxis == "BezierX")
            cellAxisVis[cell][0] = tensor.ev1() * tensor.evals()[0] * axisLineScale;
          if(displayAxis == "BezierY")
            cellAxisVis[cell][1] = tensor.ev2() * tensor.evals()[1] * axisLineScale;
        }

        float ratio =  tensor.evals()[0]/tensor.evals()[1]-1;
        float difference =  tensor.evals()[0]-tensor.evals()[1];
        if(displayAxis == "All"){
          cellAxisVis[cell][0] = tensor.ev1() * tensor.evals()[0] * axisLineScale;
          cellAxisVis[cell][1] = tensor.ev2() * tensor.evals()[1] * axisLineScale;
          cellAxisVis[cell][2] = (tensor.ev1() ^ tensor.ev2()) * tensor.evals()[2] * axisLineScale;
        } else if(displayAxis == "Both"){
          cellAxisVis[cell][0] = tensor.ev1() * tensor.evals()[0] * axisLineScale;
          cellAxisVis[cell][1] = tensor.ev2() * tensor.evals()[1] * axisLineScale;
          cellAxisVis[cell][2] = Point3f(0,0,0);
        }
        if(displayAxis == "Ratio"){
          float orientation = ratio;
          cellAxisVis[cell][0] = tensor.ev1() * (orientation > orientationThreshold ? orientation : 0.f) * axisLineScale;
          cellAxisVis[cell][1] = Point3f(0,0,0);
          cellAxisVis[cell][2] = Point3f(0,0,0);
        }
        if(displayAxis == "Difference"){
          float orientation = difference;
          cellAxisVis[cell][0] = tensor.ev1() * (orientation > orientationThreshold ? orientation : 0.f) * axisLineScale;
          cellAxisVis[cell][1] = Point3f(0,0,0);
          cellAxisVis[cell][2] = Point3f(0,0,0);
        }
        cellAxisColor[cell][0] = axisColor1;
        cellAxisColor[cell][1] = axisColor2;
        cellAxisColor[cell][2] = axisColor3;
      }
    }

    // Heat map of orientation
    if(displayHeatMap != "None"){
      IntFloatAttr labelHeatMap;
      mesh->labelHeat().clear();
      std::vector<float> values;
      values.reserve(cellAxis.size());
      forall(const IntSymTensorPair &p, cellAxis)
      {
        float value;
        int cell = p.first;
        if(displayHeatMap == "Ratio")
          value = (p.second.evals()[0] - 0.5*p.second.evals()[1] - 0.5*p.second.evals()[2]) / (p.second.evals()[0] + p.second.evals()[1] + p.second.evals()[2]);
        if(displayHeatMap == "Difference")
          value = p.second.evals()[0] - 0.5*p.second.evals()[1] - 0.5*p.second.evals()[2];

        if(displayHeatMap == "BezierX" or displayHeatMap == "BezierY"){
          if(cellAxisBezier.size() == 0)
            throw(QString("No cell axis bezier. Run 'Cell Axis Custom Directions' first."));
          if(cellAxisBezier.find(cell) == cellAxisBezier.end())
            continue;
          SymmetricTensor tensor = cellAxisBezier.find(cell)->second;
          if(displayAxis == "BezierX")
            value = tensor.evals()[0];
          if(displayAxis == "BezierY")
            value = tensor.evals()[1];
        }

        labelHeatMap[cell] = value;
        values.push_back(value);
      }
      mesh->labelHeat() = labelHeatMap;
      // Rescale heat map
      sort(values.begin(), values.end());
      int rankmax = values.size() - 1;
      int rankmin = 0;
      int heatmapPercentile = 90;// we could make that a parameter
      if(heatmapPercentile < 100) {
        rankmax = (values.size() * heatmapPercentile) / 100;
        rankmin = values.size() - rankmax;
      }
      float maxValue = values[rankmax];
      float minValue = values[rankmin];
  
//      // Adjust heat map if run on parents. TO CHANGE AFTER HEATMAP REFACTORING
//      if(mesh->parents().size() != 0){
//        Information::out << "parent label loaded" << endl;
//        mesh->labelHeat().clear();
//        IntIntAttr & parentMap = mesh->parents();
//        forall(const IntIntPair &p, parentMap)
//          if(p.second > 0 and labelHeatMap.count(p.second) != 0)
//            mesh->labelHeat()[cell] = labelHeatMap[p.second];
//      }
//      else
//        mesh->labelHeat() = labelHeatMap;
  
      mesh->heatMapBounds() = Point2f(minValue, maxValue);
      if(displayHeatMap == "Ratio")
        mesh->heatMapUnit() = QString("Ratio");
      if(displayHeatMap == "Difference")
        mesh->heatMapUnit() = QString("signal/area");
      mesh->setShowLabel("Label Heat");
      mesh->updateTriangles();
    }
  
    return true;
  }
  REGISTER_PROCESS(DisplaySignalOrientation3D);
}
