//
// 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 <Solver.hpp>
#include <KrylovMethods.hpp>

namespace mgx
{
  const int DIM = 2;
  const double DX = .00001;

  void Solver::solve() 
  {
    if(method == Euler)
      solveEuler();
    else if(method == BackwardEuler)
      solveBackwardEuler();
  }
  
  void Solver::solveEuler(){
  
    // update derivs
    forall(const cell &c, C)
      updateDerivatives(c);

    forall(const cell &c, C) {
      // calc new values
      std::vector<double> values = getValues(c);
      std::vector<double> derivs = getDerivatives(c);
  
      for(uint i = 0; i < values.size(); i++)
        values[i] += eulerDt * derivs[i];
  
      setValues(c, values);
    }
  
  }
  
  void Solver::buildJacobian(const cell& c)
  {
      // save initial values
      std::vector<double> values0 = getValues(c);
      // save initial derivatives
      updateDerivatives(c);
      std::vector<double> derivs0 = getDerivatives(c);
  
      // find partials for this cell
      VtxData &vd = c->*cellVars;
      Matrix2d &partials = vd.j;
      for(uint i = 0; i < DIM; i++) {
        std::vector<double> values = values0;
        values[i] += DX;
        setValues(c, values);
        updateDerivatives(c);
        std::vector<double> derivs = getDerivatives(c);

        for(uint j = 0; j < DIM; j++)
          partials[j][i] = (derivs0[j] - derivs[j])/DX * eulerDt;

        partials[i][i] += 1.0;
      }
      // write back old values
      setValues(c, values0);
  
      // handle neighbor cells
      forall(const cell &n, C.neighbors(c)){
        EdgData& ed = C.edge(n,c)->*cellEdges;
        Matrix2d &partials = ed.j;
        std::vector<double> valuesN0 = getValues(n);
  
        for(uint i = 0; i < DIM; i++){
          std::vector<double> values = valuesN0;
          values[i] += DX;
          setValues(n, values);
          updateDerivatives(c);
          std::vector<double> derivs = getDerivatives(c);
          for(uint j = 0; j < DIM; j++)
            partials[j][i] = (derivs0[j] - derivs0[j])/DX * eulerDt;
        }
        // write back old values
        setValues(n, valuesN0);
        //EdgData& ed = C.edge(n,c)->*cellEdges;
        //ed.j = partials;
      }
  }
  
  void Solver::initialize()
  {
  
      // Distributed matrix vectors for FEM solving
      vJ = new DistVertexAttr<DistNhbdC, VtxData, Matrix2d>(*nhbd, cellVars, &VtxData::j);
      eJ = new DistEdgeAttr<DistNhbdC, EdgData, Matrix2d>(*nhbd, cellEdges, &EdgData::j);
      vX = new DistVertexAttr<DistNhbdC, VtxData, Point2d>(*nhbd, cellVars, &VtxData::x);
      vB = new DistVertexAttr<DistNhbdC, VtxData, Point2d>(*nhbd, cellVars, &VtxData::b);
  
      J = new DMatrix2d(*vJ, *eJ);
      X = new DVector2d(*vX);
      B = new DVector2d(*vB);
  
  //std::cout << "CellVars:" << cellVars->size() << endl;
      vJ->setAttrMap(cellVars);
      eJ->setAttrMap(cellEdges);
      vX->setAttrMap(cellVars);
      vB->setAttrMap(cellVars);
  
  }
  
  
  void Solver::solveBackwardEuler()
  {
     initialize();
  
      // Allocate Jacobian, write X and B vectors
      vJ->alloc();
      eJ->alloc();
      vX->alloc();
      vB->alloc();
  
      if(debugOutput) std::cout << "backward euler " << std::endl;
  
      //std::cout << "backward euler " << std::endl;
  
    //std::vector<double> sol;
    std::map<cell, Point2d> backupVal;
    std::map<cell, Point2d> initialVal, initialDeriv;
    std::map<cell, Point2d> bVec;
    std::map<cell, int> cellNr;
  
    int counter = 0;
    forall(const cell &c, C){
      cellNr[c] = counter;
      counter++;
    }
  
    // initialize
    int steps = 0;
    double cgTolerance = 1E-4;
    double cgMaxIter = 20;
    int maxSteps = 10;
    double errTolerance = 1E-4;
  
    // do newton iteration until tolerance reached
    for(int i = 0; i < maxNewtonSteps; ++i) {
      // create jacobian
      forall(const cell &c, C)
        buildJacobian(c);
  
      // create b vector
      forall(const cell &c, C) {
        updateDerivatives(c);
        VtxData& vd = c->*cellVars;
        std::vector<double> derivs = getDerivatives(c);
        for(int j = 0; j < DIM; ++j)
          vd.b[j] = derivs[j] * eulerDt;
      }
  
      // Total iterations of CG
      int totalCGIter = 0;
  
      // Clear x
      *X = 0;
      //std::cout << "now solve " << std::endl;
      // solve for x
  
      nhbd->write();
      J->write();
      B->write();
      X->write();
  
      //totalCGIter += CGPreCond(*J, *X, *B, cgTolerance, cgMaxIter,1E100, 111);
      totalCGIter += BiCGStab(*J, *X, *B, cgTolerance, cgMaxIter,1E100, 0); //111);
      //std::cout << "solved " << totalCGIter << std::endl;
      J->read();
      B->read();
      X->read();
      forall(const cell &c, C) {
        VtxData &p = c->*cellVars;
        //std::cout << "Cell " << c->saveId << "/" << p.x << "//" << p.b << "/" << p.j << std::endl;
      }
  
      // update model values & calc newton delta
      double newtDelta = 0.0;
      forall(const cell &c, C) {
        VtxData &p = c->*cellVars;
        std::vector<double> values = getValues(c);
        for(uint i = 0; i < DIM; i++)
          values[i] += p.x[i];
        newtDelta += norm(p.x);
        //std::cout << "cell " << c->saveId << "/" << values[0] << " " << values[1] << std::endl;
        setValues(c, values);
      }
      //std::cout << "step " << newtDelta << std::endl;
      if(newtDelta <= errTolerance) 
        break;
  
    } // end of newton
    //std::cout << "newton done " << std::endl;
  }
}
