//
// 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.
// 
#ifndef KRYLOV_METHODS_HPP
#define KRYLOV_METHODS_HPP

#include <Geometry.hpp>
#include <DistObject.hpp>

namespace mgx
{
  // Stabilized Biconjugate Gradient method using distributed matrix
  template<typename DMatrixT, typename DVectorT>
  int BiCGStab(DMatrixT &A, DVectorT &x, DVectorT &b, double tolerance, uint maxiter=500, double maxnorm=1e100, int info=0)
  {
    // Create distributed objects for temps, GPU storage only
    typename DVectorT::DNhbd &nhbd = x.nhbd();
    typename DVectorT::DVertex xb_o(nhbd), r_o(nhbd), rnxt_o(nhbd), p_o(nhbd), v_o(nhbd);
    typename DVectorT::DVertex s_o(nhbd), t_o(nhbd), vec_o(nhbd), r0_o(nhbd), tx_o(nhbd);
  
    // Create distributed vectors and matrices
    DVectorT xb(xb_o), r(r_o), rnxt(rnxt_o), p(p_o), v(v_o);
    DVectorT s(s_o), t(t_o), vec(vec_o), r0(r0_o), tx(tx_o);
  
    // calculate residual
    r0 = A * x;
    r0 = b - r0;
  
    // first search direction
    p = r0;
    r = r0;
  
    uint iteration = 0, best = 1; 
    double rnorm, minnorm = 0;
  
    do {
      double alpha, alpha_u, alpha_d;
      double beta, beta_u, beta_d;
      double omega, omega_u, omega_d;
  
      // v_j
      v = A * p;
  
      // alpha_j
      alpha_u = r * r0;
      alpha_d = v * r0;
      if(alpha_d == 0)
        alpha = 0;
      else
        alpha = alpha_u / alpha_d;
  
      // s_j
      vec = v * alpha;
      s = r - vec;
  
      // t_j
      t = A * s;
  
      // omega
      omega_u = t * s;
      omega_d = t * t;
      if(omega_d == 0)
        omega = 0;
      else
        omega = omega_u / omega_d;
  
      // x_(j+1)
      vec = p * alpha;
      x += vec;
      vec = s * omega;
      x += vec;
  
      // r_(j+1)
      vec = t * omega;
      rnxt = s - vec;
  
      // beta_j
      beta_u = rnxt * r0;
      beta_d = r * r0;
      if(omega == 0 or beta_d == 0)
        beta = 0;
      else
        beta = (alpha*beta_u) / (omega*beta_d);
  
      // p_(j+1)
      vec = v * omega;
      vec = p - vec;
      vec *= beta;
      p = rnxt + vec;
  
      // change index of residuals
      r = rnxt;
  
      // Grab norm
      rnorm = norm(r);
  
      // Set min norm, 0 if first time through 
      if(iteration == 0 or rnorm <= minnorm) {
        minnorm = rnorm;
        tx = x;
        best = iteration;
      }
      iteration++;
  
      if(info > 1)
        std::cout << iteration << " iterations: " << " norm of residual: " << rnorm << std::endl;
    }
    while(!isNan(rnorm) and rnorm < maxnorm and rnorm/100.0 < minnorm and rnorm > tolerance and iteration <= maxiter);
  
    // Read result vector back from GPU
    x = tx;
  
    if(info > 0)
      std::cout<<"BICGSTAB iterations:"<< iteration - 1 << " best iteration:" << best << " norm:" << minnorm << std::endl;
  
    return iteration - 1;
  };
  
  // Biconjugate Gradient method using distributed matrix
  template<typename DMatrixT, typename DVectorT>
  int BiCG(DMatrixT &A, DMatrixT &AT, DVectorT &x, DVectorT &b, double tolerance, uint maxiter=500, double maxnorm=1e100, int info=0)
  {
    // Create distributed objects for temps, GPU storage only
    typename DVectorT::DNhbd &nhbd = x.nhbd();
    typename DVectorT::DVertex r1_o(nhbd), r2_o(nhbd), p1_o(nhbd), p2_o(nhbd);
    typename DVectorT::DVertex Ap1_o(nhbd), ATp2_o(nhbd), t_o(nhbd);
  
    // Create distributed vectors and matrices
    DVectorT r1(r1_o), r2(r2_o), p1(p1_o), p2(p2_o), Ap1(Ap1_o), ATp2(ATp2_o), t(t_o);
  
    // calculate residual
    x = 0;
    r1 = b * -1.0;
    r2 = r1;
    p1 = b;
    p2 = p1;
    double rr = r1 * r2;
    if(rr < tolerance) {
      std::cout << "Initial residual below tolerance" << std::endl;
      return 0;
    }
  
    uint iteration = 1, best = 1; 
    double minnorm = 0;
    double oldrr;
  
    do {
      Ap1 = A * p1;
      ATp2 = AT * p2; // should be A transpose
  
      double p2Ap1 = p2 * Ap1;
      if(info > 0 and p2Ap1 == 0)
        std::cout << "error" << std::endl;
  
      double alpha = rr / p2Ap1;
      oldrr = rr;
  
      t = p1 * alpha;
      x = x + t;
      t = Ap1 * alpha;
      r1 = r1 + t;
      t = ATp2 * alpha;
      r2 = r2 + t;
      rr = r1 * r2;
  
      double beta = rr/oldrr;
      p1 = p1 * beta;
      p1 = p1 - r1;
      p2 = p2 * beta;
      p2 = p2 - r2;
  
      // Set min norm, 0 if first time through 
      if(minnorm == 0 or rr <= minnorm) {
        minnorm = rr;
        best = iteration;
      }
  
      if(info > 1)
        std::cout << iteration << " iterations: " << " norm of residual: " << rr << std::endl;
  
      iteration++;
    }
    while(!isNan(rr) and rr < maxnorm and rr/100.0 < minnorm and rr > tolerance and iteration <= maxiter);
  
    if(info > 0)
      std::cout<<"BICG iterations:"<< iteration - 1 << " best iteration:" << best << " norm:" << minnorm << std::endl;
  
    return iteration - 1;
  };
  
  // Conjugate Gradient method using distributed matrix
  template<typename DMatrixT, typename DVectorT>
  int CG(DMatrixT &A, DVectorT &x, DVectorT &b, double tolerance, uint maxiter=500, double maxnorm=1e100, int info=0)
  {
    // Create distributed objects for temps, GPU storage only
    typename DVectorT::DNhbd &nhbd = x.nhbd();
    typename DVectorT::DVertex r_o(nhbd), p_o(nhbd), ap_o(nhbd), t_o(nhbd);
  
    // Create distributed vectors and matrices
    DVectorT r(r_o), p(p_o), ap(ap_o);

    // calculate residual
    r = A * x;
    r = b - r;
    p = r;
    double rr = r * r;
    double oldrr = rr;
  
    uint iteration = 1, best = 1; 
    double minnorm = 0;
  
    do {
      ap = A * p;

      double alpha = p * ap;
      alpha = rr / alpha;
      x = p * alpha + x;
      r = ap * -alpha + r;
  
      oldrr = rr;
      rr = r * r;
      double beta = rr/oldrr;
      p = p * beta + r;
  
      // Set min norm, 0 if first time through 
      if(minnorm == 0 or rr <= minnorm) {
        minnorm = rr;
        best = iteration;
      }
  
      if(info > 1)
        std::cout << iteration << " iterations: " << " norm of residual: " << rr << std::endl;
  
      iteration++;
    }
    while(!isNan(rr) and rr < maxnorm and rr/100.0 < minnorm and rr > tolerance and iteration <= maxiter);
  
    if(info > 0)
      std::cout<<"CG iterations:"<< iteration - 1 << " best iteration:" << best << " norm:" << minnorm << std::endl;
  
    return iteration - 1;
  };
  
  // Conjugate Gradient method with Jacobi pre-conditioner using distributed matrix
  template<typename DMatrixT, typename DVectorT>
  int CGPreCond(DMatrixT &A, DVectorT &x, DVectorT &b, double tolerance, uint maxiter=500, double maxnorm=1e100, int info=0)
  {
    // Create distributed objects for temps, GPU storage only
    typename DVectorT::DNhbd &nhbd = x.nhbd();
    typename DVectorT::DVertex r_o(nhbd), p_o(nhbd), ap_o(nhbd), t_o(nhbd), z_o(nhbd);
    typename DMatrixT::DVertex Mv_o(nhbd);
    typename DMatrixT::DEdge Me_o(nhbd);
  
    // Create distributed vectors and matrices
    DVectorT r(r_o), p(p_o), ap(ap_o), z(z_o);
    // Create preconditioner
    DMatrixT M(Mv_o, Me_o);
    jacobiPreCond(M, A);
  
    // calculate residual
    r = A * x;
    r = b - r;
    z = M * r;
    p = z;
    double rz = r * z;
  
    uint iteration = 1, best = 1; 
    double minnorm = 0;
  
    do {
      ap = A * p;
      double alpha = p * ap;
      alpha = rz / alpha;
      x = p * alpha + x;
      r = ap * -alpha + r;
      z = M * r;
  
      double oldrz = rz;
      rz = r * z;
      double beta = rz/oldrz;
      p = p * beta + z;
  
      // Set min norm, 0 if first time through 
      if(minnorm == 0 or rz <= minnorm) {
        minnorm = rz;
        best = iteration;
      }
  
      if(info > 1)
        std::cout << iteration << " iterations: " << " norm of residual: " << rz << std::endl;
  
      iteration++;
    }
    while(!isNan(rz) and rz > tolerance and iteration <= maxiter);
    //while(!isNan(rz) and rz < maxnorm and rz/100.0 < minnorm and rz > tolerance and iteration <= maxiter);
  
    if(info > 0)
      std::cout<<"CG iterations:"<< iteration - 1 << " best iteration:" << best << " norm:" << minnorm << std::endl;
  
    return iteration - 1;
  };

  // Preconditioned Stabilized Biconjugate Gradient method using distributed matrix
  // RSS Does not seem to work, where is the bug?
  template<typename DMatrixT, typename DVectorT>
  int BiCGStabPreCond(DMatrixT &A, DVectorT &x, DVectorT &b, double tolerance, uint maxiter=500, double maxnorm=1e100, int info=0)
  {
    // Create distributed objects for temps, GPU storage only
    typename DVectorT::DNhbd &nhbd = x.nhbd();
    typename DVectorT::DVertex xb_o(nhbd), r_o(nhbd), rnxt_o(nhbd), p_o(nhbd), v_o(nhbd), pHat_o(nhbd), sHat_o(nhbd);
    typename DVectorT::DVertex s_o(nhbd), t_o(nhbd), vec_o(nhbd), r0_o(nhbd), tx_o(nhbd);
  
    // Create distributed vectors and matrices
    DVectorT xb(xb_o), r(r_o), rnxt(rnxt_o), p(p_o), v(v_o), pHat(pHat_o), sHat(sHat_o);;
    DVectorT s(s_o), t(t_o), vec(vec_o), r0(r0_o), tx(tx_o);

    // Create preconditioner
    typename DMatrixT::DVertex Mv_o(nhbd);
    typename DMatrixT::DEdge Me_o(nhbd);
  
    DMatrixT M(Mv_o, Me_o);
    jacobiPreCond(M, A); 

    // calculate residual
    r0 = A * x;
    r0 = b - r0;
  
    // first search direction
    p = r0;
    r = r0;
  
    uint iteration = 0, best = 1; 
    double rnorm, minnorm = 0;
  
    do {
      double alpha, alpha_u, alpha_d;
      double beta, beta_u, beta_d;
      double omega, omega_u, omega_d;

      pHat = M * p;
  
      // v_j
      v = A * pHat;
  
      // alpha_j
      alpha_u = r * r0;
      alpha_d = v * r0;
      if(alpha_d == 0)
        alpha = 0;
      else
        alpha = alpha_u / alpha_d;
  
      // s_j
      vec = v * alpha;
      s = r - vec;
      sHat = M * s;
  
      // t_j
      t = A * sHat;
  
      // omega
      omega_u = t * s;
      omega_d = t * t;
      if(omega_d == 0)
        omega = 0;
      else
        omega = omega_u / omega_d;
  
      // x_(j+1)
      vec = pHat * alpha;
      x += vec;
      vec = sHat * omega;
      x += vec;
  
      // r_(j+1)
      vec = t * omega;
      rnxt = s - vec;
  
      // beta_j
      beta_u = rnxt * r0;
      beta_d = r * r0;
      if(omega == 0 or beta_d == 0)
        beta = 0;
      else
        beta = (alpha*beta_u) / (omega*beta_d);
  
      // p_(j+1)
      vec = v * omega;
      vec = p - vec;
      vec *= beta;
      p = rnxt + vec;
  
      // change index of residuals
      r = rnxt;
  
      // Grab norm
      rnorm = norm(r);
  
      // Set min norm, 0 if first time through 
      if(iteration == 0 or rnorm <= minnorm) {
        minnorm = rnorm;
        tx = x;
        best = iteration;
      }
      iteration++;
  
      if(info > 1)
        std::cout << iteration << " iterations: " << " norm of residual: " << rnorm << std::endl;
    }
    while(!isNan(rnorm) and rnorm < maxnorm and rnorm > tolerance and iteration <= maxiter);
    //while(!isNan(rnorm) and rnorm < maxnorm and rnorm/100.0 < minnorm and rnorm > tolerance and iteration <= maxiter);
  
    // Read result vector back from GPU
    x = tx;
  
    if(info > 0)
      std::cout<<"BICGSTAB iterations:"<< iteration - 1 << " best iteration:" << best << " norm:" << minnorm << std::endl;
  
    return iteration - 1;
  };
}
#endif
