/*************************************************************************************

Grid physics library, www.github.com/paboyle/Grid

Source file: ./lib/algorithms/approx/Forecast.h

Copyright (C) 2015

Author: Peter Boyle <paboyle@ph.ed.ac.uk>
Author: paboyle <paboyle@ph.ed.ac.uk>
Author: David Murphy <dmurphy@phys.columbia.edu>

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

See the full license in the file "LICENSE" in the top level distribution directory
*************************************************************************************/
			   /*  END LEGAL */

#ifndef INCLUDED_FORECAST_H
#define INCLUDED_FORECAST_H

NAMESPACE_BEGIN(Grid);

// Abstract base class.
// Takes a matrix (Mat), a source (phi), and a vector of Fields (chi)
// and returns a forecasted solution to the system D*psi = phi (psi).
template<class Matrix, class Field>
class Forecast
{
public:
  virtual Field operator()(Matrix &Mat, const Field& phi, const std::vector<Field>& chi) = 0;
};

// Implementation of Brower et al.'s chronological inverter (arXiv:hep-lat/9509012),
// used to forecast solutions across poles of the EOFA heatbath.
//
// Modified from CPS (cps_pp/src/util/dirac_op/d_op_base/comsrc/minresext.C)
template<class Matrix, class Field>
class ChronoForecast : public Forecast<Matrix,Field>
{
public:
  Field operator()(Matrix &Mat, const Field& phi, const std::vector<Field>& prev_solns)
  {
    int degree = prev_solns.size();
    Field chi(phi); // forecasted solution

    // Trivial cases
    if(degree == 0){ chi = zero; return chi; }
    else if(degree == 1){ return prev_solns[0]; }

    //    RealD dot;
    ComplexD xp;
    Field r(phi); // residual
    Field Mv(phi);
    std::vector<Field> v(prev_solns); // orthonormalized previous solutions
    std::vector<Field> MdagMv(degree,phi);

    // Array to hold the matrix elements
    std::vector<std::vector<ComplexD>> G(degree, std::vector<ComplexD>(degree));

    // Solution and source vectors
    std::vector<ComplexD> a(degree);
    std::vector<ComplexD> b(degree);

    // Orthonormalize the vector basis
    for(int i=0; i<degree; i++){
      v[i] *= 1.0/std::sqrt(norm2(v[i]));
      for(int j=i+1; j<degree; j++){ v[j] -= innerProduct(v[i],v[j]) * v[i]; }
    }

    // Perform sparse matrix multiplication and construct rhs
    for(int i=0; i<degree; i++){
      b[i] = innerProduct(v[i],phi);
      Mat.M(v[i],Mv);
      Mat.Mdag(Mv,MdagMv[i]);
      G[i][i] = innerProduct(v[i],MdagMv[i]);
    }

    // Construct the matrix
    for(int j=0; j<degree; j++){
      for(int k=j+1; k<degree; k++){
	G[j][k] = innerProduct(v[j],MdagMv[k]);
	G[k][j] = conjugate(G[j][k]);
      }}

    // Gauss-Jordan elimination with partial pivoting
    for(int i=0; i<degree; i++){

      // Perform partial pivoting
      int k = i;
      for(int j=i+1; j<degree; j++){ if(abs(G[j][j]) > abs(G[k][k])){ k = j; } }
      if(k != i){
	xp = b[k];
	b[k] = b[i];
	b[i] = xp;
	for(int j=0; j<degree; j++){
	  xp = G[k][j];
	  G[k][j] = G[i][j];
	  G[i][j] = xp;
	}
      }

      // Convert matrix to upper triangular form
      for(int j=i+1; j<degree; j++){
	xp = G[j][i]/G[i][i];
	b[j] -= xp * b[i];
	for(int k=0; k<degree; k++){ G[j][k] -= xp*G[i][k]; }
      }
    }

    // Use Gaussian elimination to solve equations and calculate initial guess
    chi = zero;
    r = phi;
    for(int i=degree-1; i>=0; i--){
      a[i] = 0.0;
      for(int j=i+1; j<degree; j++){ a[i] += G[i][j] * a[j]; }
      a[i] = (b[i]-a[i])/G[i][i];
      chi += a[i]*v[i];
      r -= a[i]*MdagMv[i];
    }

    RealD true_r(0.0);
    ComplexD tmp;
    for(int i=0; i<degree; i++){
      tmp = -b[i];
      for(int j=0; j<degree; j++){ tmp += G[i][j]*a[j]; }
      tmp = conjugate(tmp)*tmp;
      true_r += std::sqrt(tmp.real());
    }

    RealD error = std::sqrt(norm2(r)/norm2(phi));
    std::cout << GridLogMessage << "ChronoForecast: |res|/|src| = " << error << std::endl;

    return chi;
  };
};

NAMESPACE_END(Grid);

#endif