//
// 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 <iostream>
#include <cstdio>
#include <cmath>

#include <HejnowiczSurface.hpp>
#include <Function.hpp>
#include <Contour.hpp>
#include <Random.hpp>
#include <Information.hpp>
#include <SystemProcessLoad.hpp>

#include <gsl/gsl_vector.h>
#include <gsl/gsl_multiroots.h>

namespace mgx
{
	// for the polar surfaces: 
  const size_t SAMPLES = 10000;	// Size of sampling tables

  // Convert nrs coordinates to Cartesian, 2d not scaled
  Point2d evalPosXY(const Point2d &nrs)
  {
    double a_parm = 1.0;
    return Point2d(
        (a_parm * 2.0/M_PI) * (atan(tan(nrs[1]) * tanh(nrs[0]))),
        (a_parm/M_PI) * (log(pow(cosh(nrs[0]), 2) - pow(sin(nrs[1]), 2)))); 
  }

  // Evaluation function for GSL root search
  int rosenbrock_f(const gsl_vector *x, void *p, gsl_vector *f)
  {
    Point2d pos = *(Point2d *)p;
    Point2d nrs(gsl_vector_get(x, 0), gsl_vector_get(x, 1));
    Point2d y = pos - evalPosXY(nrs);
  
    gsl_vector_set(f, 0, y.x());
    gsl_vector_set(f, 1, y.y());
  
    return GSL_SUCCESS;
  }

  void print_state (size_t iter, gsl_multiroot_fsolver * s)
  {
    Information::out << QString("iter = %1 x = %2 %3 f(x) = %4 %5\n").arg(iter)
        .arg(gsl_vector_get (s->x, 0)) .arg(gsl_vector_get (s->x, 1))
        .arg(gsl_vector_get (s->f, 0)) .arg(gsl_vector_get (s->f, 1)) << endl;
  } 

  Point3d HejnowiczSurface::localToCartesian(const Point2d &local)
  {
    return Point3d(evalPosXY(local) * scale);
  }

  bool evalNRS(Point2d pos, Point2d &nrs)
  {
    bool neg = pos.x() < 0;
    if(nrs.x() < 0)
      nrs.x() = 0;
    if(neg) {
      pos.x() = -pos.x();
      if(nrs.y() < 0)
        nrs.y() = -nrs.y();
    }

    const gsl_multiroot_fsolver_type *T;
    gsl_multiroot_fsolver *s;
  
    int status;
    size_t iter = 0;

    gsl_multiroot_function f = {&rosenbrock_f, 2, &pos};
  
    gsl_vector *x = gsl_vector_alloc (2);
    gsl_vector_set(x, 0, nrs.x());
    gsl_vector_set(x, 1, nrs.y());
  
    T = gsl_multiroot_fsolver_hybrids;
    s = gsl_multiroot_fsolver_alloc(T, 2);
    gsl_multiroot_fsolver_set (s, &f, x);
  
    //print_state (iter, s);
    do 
    {
      iter++;
      status = gsl_multiroot_fsolver_iterate (s);
      //print_state (iter, s);
      if(status)
        break;

      // Clip if anything goes below 0
      if((nrs.x() = gsl_vector_get(s->x, 0)) < 0)
        gsl_vector_set(x, 0, (nrs.x() = 0));
      if((nrs.y() = gsl_vector_get(s->x, 1)) < 0)
        gsl_vector_set(x, 1, (nrs.y() = 0));
  
      status = gsl_multiroot_test_residual(s->f, 1e-12);
    } while (status == GSL_CONTINUE && iter < 100);     
  
    if(status == GSL_SUCCESS) {
      nrs.x() = gsl_vector_get(s->x, 0);
      nrs.y() = gsl_vector_get(s->x, 1);
    } 
    gsl_multiroot_fsolver_free(s);
    gsl_vector_free(x);

    // If negative, flip back
    if(neg and nrs.y() > 0)
      nrs.y() = -nrs.y();
    if(nrs.x() < 0)
      nrs.x() = 0;

    if(status == GSL_SUCCESS)
      return true;
    else
      return false;
  }

  bool HejnowiczSurface::cartesianToLocal(const Point3d &cartesian, Point2d &local)
  {
    // If not found, look around a bit
    Point2d pos = Point2d(cartesian) / scale;
    Point2d nrs(local);
    bool success = evalNRS(pos, nrs);
    if(!success) {
      int maxiter = 100;
      while(!success and maxiter > 0) {
        nrs = local + Point2d(ran(1.0), ran(1.0));
        success = evalNRS(pos, nrs);
        maxiter--;
      }
    }
    if(!success)
      Information::out << "Root search failed, pos:" << cartesian << " nrs:" << local << endl;
    local = nrs;
    return success;
  } 

  bool HejnowiczSurface::initialize(VertexAttr *vData, const QStringList &parms) 
  {
    // Process the parameters
    processParms(parms);

    // Get the attribute maps
    vertexData = vData;

    Information::out << "Initialized Hejnowicz Surface" << endl;

    return true;
  }
  
  // Read in model parameters
  bool HejnowiczSurface::processParms(const QStringList &parms)
  {
     uS = parms[puS].toDouble();
     vR = parms[pvR].toDouble();
     A = parms[pA].toDouble();
     C = parms[pC].toDouble();
     D = parms[pD].toDouble();
     scale = parms[pScale].toDouble();
     length = parms[pLength].toDouble();
     growthScale = parms[pGrowthScale].toDouble();
     growthFuncFile = parms[pGrowthFunc];

     if(parms[pGrowthType] == "Hejnowicz")
       growthType = GROWTH_HEJN;
     else if(parms[pGrowthType] == "GrowthFunc")
       growthType = GROWTH_FUNC;
     else if(parms[pGrowthType] == "GrowthMap")
       growthType = GROWTH_MAP;

     if(growthType == GROWTH_FUNC)
       setupGrowth();

     return true;
  }

  bool HejnowiczSurface::setGrowthMap(std::map<double, double> &growthMap) 
  {
    if(growthType != GROWTH_MAP)
      return false;

    // Clear growth table
    growthTab.clear();
    if(growthMap.size() == 0)
      return false;
    growthTab.resize(SAMPLES + 1);

    // First create a table we can use as a function
    std::map<double, double>::iterator gP = growthMap.begin();
    std::pair<double, double> current = *gP, next = *gP;
    if(++gP != growthMap.end())
      next = *gP;

    // Build function table
    for(uint i = 0; i <= SAMPLES; ++i) {
      double dist = double(i)/SAMPLES * length;
      while(gP != growthMap.end() and dist > next.first) {
        current = next;
        if(++gP != growthMap.end())
          next = *gP;
      }
        
      if(current.second == next.second or dist <= current.first)
        growthTab[i] = current.second;
      else {
        double s = (dist - current.first)/(next.first - current.first);
        growthTab[i] = (1.0 - s) * current.second + s * next.second;
      }
    }

    setupGrowth();

    return true;
  }

  // Function interface into the growth map table, used by integrator.
  double HejnowiczSurface::growthMapFunc(double s)
  {
    if(growthTab.size() == 0)
      return 0;
    if(s <= 0)
      return growthTab[0];
    if(s >= length)
      return growthTab[SAMPLES];

    double dx = 1.0/double(SAMPLES);
    int i = int(s/dx);
    s -= dx * i;

    return (1.0 - s) * growthTab[i] + s * growthTab[i+1];
  }
  
  bool HejnowiczSurface::setupGrowth()
  {
    if(growthType == GROWTH_FUNC)
      growthFunc.setFile(growthFuncFile);

    double dx = 1.0/SAMPLES, sum = 0, y;
    bool yNegative = false;
    
    // Integrate growth function
    velocityCart.resize(SAMPLES + 1);
    for(uint i = 0; i <= SAMPLES; i++) {
      velocityCart[i] = sum;
      if(growthType == GROWTH_FUNC)
        y = growthScale * growthFunc(double(i) * dx);
      else if(growthType == GROWTH_MAP)
        y = growthScale * growthMapFunc(double(i) * dx);
      if(y < 0.0)
        yNegative = true;
      sum += dx * y;
    }

    // Convert to local coords, first find NRS length
    Point2d nrs(.1, 0);
    Point3d cart;
    do {
      nrs.x() *= 2.0;
      cart = localToCartesian(nrs);
    } while(cart.y() < length);

    double upper = nrs.x();
    double lower = upper/2.0;
    // Binary search the NRS coord
    do {
			nrs.x() = (upper + lower)/2.0;
      cart = localToCartesian(nrs);
      if(cart.y() > length)
        upper = nrs.x();
      else
        lower = nrs.x();
    } while(fabs(cart.y() - length) > 1e-6);

    nrsLength = nrs.x();
    velocityNRS.resize(SAMPLES + 1);
    velocityNRS[0] = 0; // first point doesn't move
    for(uint i = 1; i <= SAMPLES; ++i) {
      nrs.x() = double(i)/SAMPLES * nrsLength;
      cart = localToCartesian(nrs);
      double velocity = velocityCart[int(cart.y()/length * SAMPLES)];
      velocityNRS[i] = velocity/length * nrsLength; 
    }

    if(yNegative)
      Information::out << "Warning: Growth function goes negative" << endl;

    Information::out << QString("Growth function NRS length:%1, surface length:%2")
                                        .arg(nrsLength).arg(length) << endl;

    return true;
  }

  // Point inversion. Convert 3d point to nrs coordinates.
  bool HejnowiczSurface::setPoint(vertex vp, vertex vsp, Point3d cp)
  {
    VertexData &p = vp->*vertexData;
    VertexData &sp = vsp->*vertexData;
    
    // Look for the NRS coordinates, set starting point
    p.nrs = sp.nrs;
    return cartesianToLocal(cp, p.nrs);
  }

  // Grow Hejnowicz surface point 
  bool HejnowiczSurface::growPoint(vertex vp, double dt, double /*time*/)
  {
    VertexData &p = vp->*vertexData;

    // Growth function of map (morphogens)
    if(growthType == GROWTH_FUNC or growthType == GROWTH_MAP) {
      // If velocity table empty just return
      if(velocityNRS.size() == 0)
        return false;

      // For now only grow points > uS
      if (p.nrs[0] > uS) {
        int i = p.nrs[0]/nrsLength * SAMPLES;
        // If over surface length, growth stops
        if(i >= SAMPLES)
          i = SAMPLES - 1;
        p.nrs[0] += velocityNRS[i] * dt; 
      }
    } else {
      // Model 2 of Hejnowicz & 1993
      if (p.nrs[0] <= uS)
        p.nrs[0] += A * p.nrs.x() * dt;
      else
        p.nrs[0] += (C * (p.nrs.x() - uS) + A * uS) * dt; 
    }

    // v is the same for both
    if(fabs(p.nrs[1]) > vR)
      p.nrs[1] += -D * sin(4 * p.nrs[1]) * dt; 

    vp->pos = Point3d(localToCartesian(p.nrs));

    return true;
  }

  // Create an initial cell at the origin, size and walls ignored for now
  bool HejnowiczSurface::initialCell(CellTissue &T, 
         const Point2d &top, const Point2d &bottom, const Point2d &middle)
  {
    VtxVec poly;
    // The center
    vertex c;
    poly.push_back(c);

    // The junctions
    Point2d cnrs;
    vertex v1;
    (v1->*vertexData).nrs = bottom;
    cnrs += bottom;
    poly.push_back(v1);
    vertex v2;
    (v2->*vertexData).nrs = middle;
    cnrs += middle;
    poly.push_back(v2);
    vertex v3;
    (v3->*vertexData).nrs = Point2d(bottom.x(), -bottom.y());
    cnrs += Point2d(bottom.x(), -bottom.y());
    poly.push_back(v3);
    vertex v4;
    (v4->*vertexData).nrs = Point2d(top.x(), -top.y());
    cnrs += Point2d(top.x(), -top.y());
    poly.push_back(v4);
    vertex v5;
    (v5->*vertexData).nrs = top;
    cnrs += top;
    poly.push_back(v5);

    forall(const vertex &v, poly) {
      updatePos(v);
      updateNormal(v);
    }

    cnrs /= 5.0;
    (c->*vertexData).nrs = cnrs;
    updatePos(c);
    updateNormal(c);
     
    return T.addCell(poly);
  }

  // Updates xyz coordinates from nrs (Hejnowicz) coordinates
  bool HejnowiczSurface::updatePos(vertex vp)
  {
    VertexData &p = vp->*vertexData;
    vp->pos = Point3d(localToCartesian(p.nrs));

    return true;
	}

  bool HejnowiczSurface::updateNormal(vertex vp)
  {
		// For now only 2d
    vp->nrml = Point3d(0,0,1);

    return true;
  }

  // Update edge data, s indicates distance along edge
  // Called when a vertex is inserted in an edge
  bool HejnowiczSurface::Subdivide::updateEdgeData(vertex l, vertex v, vertex r, double s)
  {
    // If no surface just average
    bool result = true;
    if(surface) {
      result = surface->setPoint(v, l, v->pos);
      if(!result)
        result = surface->setPoint(v, r, v->pos);
      if(result) {
        surface->updatePos(v);
        return true;
      }
    }
    // If no surface or setPoint fails, average the cylindrical coordinates 
    VertexData &L = l->*vertexData;
    VertexData &V = v->*vertexData;
    VertexData &R = r->*vertexData;
    V.nrs = (1 - s) * L.nrs + s * R.nrs;

    if(surface) {
      result = surface->setPoint(v, r, v->pos);
      if(result)
        surface->updatePos(v);
      return result;
    }
    return true; // No surface, should we really return false?
  }
}
