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

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

    Source file: ./lib/algorithms/LinearOperator.h

    Copyright (C) 2015

Author: Azusa Yamaguchi <ayamaguc@staffmail.ed.ac.uk>
Author: Peter Boyle <paboyle@ph.ed.ac.uk>

    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  GRID_ALGORITHM_LINEAR_OP_H
#define  GRID_ALGORITHM_LINEAR_OP_H

NAMESPACE_BEGIN(Grid);

/////////////////////////////////////////////////////////////////////////////////////////////
// LinearOperators Take a something and return a something.
/////////////////////////////////////////////////////////////////////////////////////////////
//
// Hopefully linearity is satisfied and the AdjOp is indeed the Hermitian conjugateugate (transpose if real):
//SBase
//   i)  F(a x + b y) = aF(x) + b F(y).
//  ii)  <x|Op|y> = <y|AdjOp|x>^\ast
//
// Would be fun to have a test linearity & Herm Conj function!
/////////////////////////////////////////////////////////////////////////////////////////////
template<class Field> class LinearOperatorBase {
public:

  // Support for coarsening to a multigrid
  virtual void OpDiag (const Field &in, Field &out) = 0; // Abstract base
  virtual void OpDir  (const Field &in, Field &out,int dir,int disp) = 0; // Abstract base

  virtual void Op     (const Field &in, Field &out) = 0; // Abstract base
  virtual void AdjOp  (const Field &in, Field &out) = 0; // Abstract base
  virtual void HermOpAndNorm(const Field &in, Field &out,RealD &n1,RealD &n2)=0;
  virtual void HermOp(const Field &in, Field &out)=0;
};


/////////////////////////////////////////////////////////////////////////////////////////////
// By sharing the class for Sparse Matrix across multiple operator wrappers, we can share code
// between RB and non-RB variants. Sparse matrix is like the fermion action def, and then
// the wrappers implement the specialisation of "Op" and "AdjOp" to the cases minimising
// replication of code.
//
// I'm not entirely happy with implementation; to share the Schur code between herm and non-herm
// while still having a "OpAndNorm" in the abstract base I had to implement it in both cases
// with an assert trap in the non-herm. This isn't right; there must be a better C++ way to
// do it, but I fear it required multiple inheritance and mixed in abstract base classes
/////////////////////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////
// Construct herm op from non-herm matrix
////////////////////////////////////////////////////////////////////
template<class Matrix,class Field>
class MdagMLinearOperator : public LinearOperatorBase<Field> {
  Matrix &_Mat;
public:
  MdagMLinearOperator(Matrix &Mat): _Mat(Mat){};

  // Support for coarsening to a multigrid
  void OpDiag (const Field &in, Field &out) {
    _Mat.Mdiag(in,out);
  }
  void OpDir  (const Field &in, Field &out,int dir,int disp) {
    _Mat.Mdir(in,out,dir,disp);
  }
  void Op     (const Field &in, Field &out){
    _Mat.M(in,out);
  }
  void AdjOp     (const Field &in, Field &out){
    _Mat.Mdag(in,out);
  }
  void HermOpAndNorm(const Field &in, Field &out,RealD &n1,RealD &n2){
    _Mat.MdagM(in,out,n1,n2);
  }
  void HermOp(const Field &in, Field &out){
    RealD n1,n2;
    HermOpAndNorm(in,out,n1,n2);
  }
};

////////////////////////////////////////////////////////////////////
// Construct herm op and shift it for mgrid smoother
////////////////////////////////////////////////////////////////////
template<class Matrix,class Field>
class ShiftedMdagMLinearOperator : public LinearOperatorBase<Field> {
  Matrix &_Mat;
  RealD _shift;
public:
  ShiftedMdagMLinearOperator(Matrix &Mat,RealD shift): _Mat(Mat), _shift(shift){};
  // Support for coarsening to a multigrid
  void OpDiag (const Field &in, Field &out) {
    _Mat.Mdiag(in,out);
    assert(0);
  }
  void OpDir  (const Field &in, Field &out,int dir,int disp) {
    _Mat.Mdir(in,out,dir,disp);
    assert(0);
  }
  void Op     (const Field &in, Field &out){
    _Mat.M(in,out);
    assert(0);
  }
  void AdjOp     (const Field &in, Field &out){
    _Mat.Mdag(in,out);
    assert(0);
  }
  void HermOpAndNorm(const Field &in, Field &out,RealD &n1,RealD &n2){
    _Mat.MdagM(in,out,n1,n2);
    out = out + _shift*in;

    ComplexD dot;	
    dot= innerProduct(in,out);
    n1=real(dot);
    n2=norm2(out);
  }
  void HermOp(const Field &in, Field &out){
    RealD n1,n2;
    HermOpAndNorm(in,out,n1,n2);
  }
};

////////////////////////////////////////////////////////////////////
// Wrap an already herm matrix
////////////////////////////////////////////////////////////////////
template<class Matrix,class Field>
class HermitianLinearOperator : public LinearOperatorBase<Field> {
  Matrix &_Mat;
public:
  HermitianLinearOperator(Matrix &Mat): _Mat(Mat){};
  // Support for coarsening to a multigrid
  void OpDiag (const Field &in, Field &out) {
    _Mat.Mdiag(in,out);
  }
  void OpDir  (const Field &in, Field &out,int dir,int disp) {
    _Mat.Mdir(in,out,dir,disp);
  }
  void Op     (const Field &in, Field &out){
    _Mat.M(in,out);
  }
  void AdjOp     (const Field &in, Field &out){
    _Mat.M(in,out);
  }
  void HermOpAndNorm(const Field &in, Field &out,RealD &n1,RealD &n2){
    _Mat.M(in,out);
	
    ComplexD dot= innerProduct(in,out); n1=real(dot);
    n2=norm2(out);
  }
  void HermOp(const Field &in, Field &out){
    _Mat.M(in,out);
  }
};

//////////////////////////////////////////////////////////
// Even Odd Schur decomp operators; there are several
// ways to introduce the even odd checkerboarding
//////////////////////////////////////////////////////////

template<class Field>
class SchurOperatorBase :  public LinearOperatorBase<Field> {
public:
  virtual  RealD Mpc      (const Field &in, Field &out) =0;
  virtual  RealD MpcDag   (const Field &in, Field &out) =0;
  virtual void MpcDagMpc(const Field &in, Field &out,RealD &ni,RealD &no) {
    Field tmp(in._grid);
    ni=Mpc(in,tmp);
    no=MpcDag(tmp,out);
  }
  virtual void HermOpAndNorm(const Field &in, Field &out,RealD &n1,RealD &n2){
    MpcDagMpc(in,out,n1,n2);
  }
  virtual void HermOp(const Field &in, Field &out){
    RealD n1,n2;
    HermOpAndNorm(in,out,n1,n2);
  }
  void Op     (const Field &in, Field &out){
    Mpc(in,out);
  }
  void AdjOp     (const Field &in, Field &out){ 
    MpcDag(in,out);
  }
  // Support for coarsening to a multigrid
  void OpDiag (const Field &in, Field &out) {
    assert(0); // must coarsen the unpreconditioned system
  }
  void OpDir  (const Field &in, Field &out,int dir,int disp) {
    assert(0);
  }
};
template<class Matrix,class Field>
class SchurDiagMooeeOperator :  public SchurOperatorBase<Field> {
protected:
  Matrix &_Mat;
public:
  SchurDiagMooeeOperator (Matrix &Mat): _Mat(Mat){};
  virtual  RealD Mpc      (const Field &in, Field &out) {
    Field tmp(in._grid);
    //	std::cout <<"grid pointers: in._grid="<< in._grid << " out._grid=" << out._grid << "  _Mat.Grid=" << _Mat.Grid() << " _Mat.RedBlackGrid=" << _Mat.RedBlackGrid() << std::endl;

    _Mat.Meooe(in,tmp);
    _Mat.MooeeInv(tmp,out);
    _Mat.Meooe(out,tmp);

    _Mat.Mooee(in,out);
    return axpy_norm(out,-1.0,tmp,out);
  }
  virtual  RealD MpcDag   (const Field &in, Field &out){
    Field tmp(in._grid);

    _Mat.MeooeDag(in,tmp);
    _Mat.MooeeInvDag(tmp,out);
    _Mat.MeooeDag(out,tmp);

    _Mat.MooeeDag(in,out);
    return axpy_norm(out,-1.0,tmp,out);
  }
};
template<class Matrix,class Field>
class SchurDiagOneOperator :  public SchurOperatorBase<Field> {
protected:
  Matrix &_Mat;
public:
  SchurDiagOneOperator (Matrix &Mat): _Mat(Mat){};

  virtual  RealD Mpc      (const Field &in, Field &out) {
    Field tmp(in._grid);

    _Mat.Meooe(in,out);
    _Mat.MooeeInv(out,tmp);
    _Mat.Meooe(tmp,out);
    _Mat.MooeeInv(out,tmp);

    return axpy_norm(out,-1.0,tmp,in);
  }
  virtual  RealD MpcDag   (const Field &in, Field &out){
    Field tmp(in._grid);

    _Mat.MooeeInvDag(in,out);
    _Mat.MeooeDag(out,tmp);
    _Mat.MooeeInvDag(tmp,out);
    _Mat.MeooeDag(out,tmp);

    return axpy_norm(out,-1.0,tmp,in);
  }
};
template<class Matrix,class Field>
class SchurDiagTwoOperator :  public SchurOperatorBase<Field> {
protected:
  Matrix &_Mat;
public:
  SchurDiagTwoOperator (Matrix &Mat): _Mat(Mat){};

  virtual  RealD Mpc      (const Field &in, Field &out) {
    Field tmp(in._grid);

    _Mat.MooeeInv(in,out);
    _Mat.Meooe(out,tmp);
    _Mat.MooeeInv(tmp,out);
    _Mat.Meooe(out,tmp);

    return axpy_norm(out,-1.0,tmp,in);
  }
  virtual  RealD MpcDag   (const Field &in, Field &out){
    Field tmp(in._grid);

    _Mat.MeooeDag(in,out);
    _Mat.MooeeInvDag(out,tmp);
    _Mat.MeooeDag(tmp,out);
    _Mat.MooeeInvDag(out,tmp);

    return axpy_norm(out,-1.0,tmp,in);
  }
};
///////////////////////////////////////////////////////////////////////////////////////////////////
// Left  handed Moo^-1 ; (Moo - Moe Mee^-1 Meo) psi = eta  -->  ( 1 - Moo^-1 Moe Mee^-1 Meo ) psi = Moo^-1 eta
// Right handed Moo^-1 ; (Moo - Moe Mee^-1 Meo) Moo^-1 Moo psi = eta  -->  ( 1 - Moe Mee^-1 Meo ) Moo^-1 phi=eta ; psi = Moo^-1 phi
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class Matrix,class Field> using SchurDiagOneRH = SchurDiagTwoOperator<Matrix,Field> ;
template<class Matrix,class Field> using SchurDiagOneLH = SchurDiagOneOperator<Matrix,Field> ;
///////////////////////////////////////////////////////////////////////////////////////////////////
//  Staggered use
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class Matrix,class Field>
class SchurStaggeredOperator :  public SchurOperatorBase<Field> {
protected:
  Matrix &_Mat;
public:
  SchurStaggeredOperator (Matrix &Mat): _Mat(Mat){};
  virtual void HermOpAndNorm(const Field &in, Field &out,RealD &n1,RealD &n2){
    GridLogIterative.TimingMode(1);
    std::cout << GridLogIterative << " HermOpAndNorm "<<std::endl;
    n2 = Mpc(in,out);
    std::cout << GridLogIterative << " HermOpAndNorm.Mpc "<<std::endl;
    ComplexD dot= innerProduct(in,out);
    std::cout << GridLogIterative << " HermOpAndNorm.innerProduct "<<std::endl;
    n1 = real(dot);
  }
  virtual void HermOp(const Field &in, Field &out){
    std::cout << GridLogIterative << " HermOp "<<std::endl;
    Mpc(in,out);
  }
  virtual  RealD Mpc      (const Field &in, Field &out) {
    Field tmp(in._grid);
    Field tmp2(in._grid);

    std::cout << GridLogIterative << " HermOp.Mpc "<<std::endl;
    _Mat.Mooee(in,out);
    _Mat.Mooee(out,tmp);
    std::cout << GridLogIterative << " HermOp.MooeeMooee "<<std::endl;

    _Mat.Meooe(in,out);
    _Mat.Meooe(out,tmp2);
    std::cout << GridLogIterative << " HermOp.MeooeMeooe "<<std::endl;

    RealD nn=axpy_norm(out,-1.0,tmp2,tmp);
    std::cout << GridLogIterative << " HermOp.axpy_norm "<<std::endl;
    return nn;
  }
  virtual  RealD MpcDag   (const Field &in, Field &out){
    return Mpc(in,out);
  }
  virtual void MpcDagMpc(const Field &in, Field &out,RealD &ni,RealD &no) {
    assert(0);// Never need with staggered
  }
};
template<class Matrix,class Field> using SchurStagOperator = SchurStaggeredOperator<Matrix,Field>;


/////////////////////////////////////////////////////////////
// Base classes for functions of operators
/////////////////////////////////////////////////////////////
template<class Field> class OperatorFunction {
public:
  virtual void operator() (LinearOperatorBase<Field> &Linop, const Field &in, Field &out) = 0;
};

template<class Field> class LinearFunction {
public:
  virtual void operator() (const Field &in, Field &out) = 0;
};

template<class Field> class IdentityLinearFunction : public LinearFunction<Field> {
public:
  void operator() (const Field &in, Field &out){
    out = in;
  };
};


/////////////////////////////////////////////////////////////
// Base classes for Multishift solvers for operators
/////////////////////////////////////////////////////////////
template<class Field> class OperatorMultiFunction {
public:
  virtual void operator() (LinearOperatorBase<Field> &Linop, const Field &in, std::vector<Field> &out) = 0;
};

// FIXME : To think about

// Chroma functionality list defining LinearOperator
/*
  virtual void operator() (T& chi, const T& psi, enum PlusMinus isign) const = 0;
  virtual void operator() (T& chi, const T& psi, enum PlusMinus isign, Real epsilon) const
  virtual const Subset& subset() const = 0;
  virtual unsigned long nFlops() const { return 0; }
  virtual void deriv(P& ds_u, const T& chi, const T& psi, enum PlusMinus isign) const
  class UnprecLinearOperator : public DiffLinearOperator<T,P,Q>
  const Subset& subset() const {return all;}
  };
*/

////////////////////////////////////////////////////////////////////////////////////////////
// Hermitian operator Linear function and operator function
////////////////////////////////////////////////////////////////////////////////////////////
template<class Field>
class HermOpOperatorFunction : public OperatorFunction<Field> {
  void operator() (LinearOperatorBase<Field> &Linop, const Field &in, Field &out) {
    Linop.HermOp(in,out);
  };
};

template<typename Field>
class PlainHermOp : public LinearFunction<Field> {
public:
  LinearOperatorBase<Field> &_Linop;
      
  PlainHermOp(LinearOperatorBase<Field>& linop) : _Linop(linop) 
  {}
      
  void operator()(const Field& in, Field& out) {
    _Linop.HermOp(in,out);
  }
};

template<typename Field>
class FunctionHermOp : public LinearFunction<Field> {
public:
  OperatorFunction<Field>   & _poly;
  LinearOperatorBase<Field> &_Linop;
      
  FunctionHermOp(OperatorFunction<Field> & poly,LinearOperatorBase<Field>& linop) 
    : _poly(poly), _Linop(linop) {};
      
  void operator()(const Field& in, Field& out) {
    _poly(_Linop,in,out);
  }
};

template<class Field>
class Polynomial : public OperatorFunction<Field> {
private:
  std::vector<RealD> Coeffs;
public:
  Polynomial(std::vector<RealD> &_Coeffs) : Coeffs(_Coeffs) { };

  // Implement the required interface
  void operator() (LinearOperatorBase<Field> &Linop, const Field &in, Field &out) {

    Field AtoN(in._grid);
    Field Mtmp(in._grid);
    AtoN = in;
    out = AtoN*Coeffs[0];
    for(int n=1;n<Coeffs.size();n++){
      Mtmp = AtoN;
      Linop.HermOp(Mtmp,AtoN);
      out=out+AtoN*Coeffs[n];
    }
  };
};

NAMESPACE_END(Grid);

#endif