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

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

    Source file: MultiRHSDeflation.h

    Copyright (C) 2023

Author: Peter Boyle <pboyle@bnl.gov>

    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 */
#pragma once

NAMESPACE_BEGIN(Grid);


/* Need helper object for BLAS accelerated mrhs projection

   i) MultiRHS Deflation

   Import Evecs -> nev x vol x internal 
   Import vector of Lattice objects -> nrhs x vol x internal
   => Cij (nrhs x Nev) via GEMM.
   => Guess  (nrhs x vol x internal)  = C x evecs (via GEMM)
   Export

   
   ii) MultiRHS block projection

   Import basis -> nblock x nbasis x  (block x internal) 
   Import vector of fine lattice objects -> nblock x nrhs x (block x internal) 

   => coarse_(nrhs x nbasis )^block = via batched GEMM

   iii)   Alternate interface: 
   Import higher dim Lattice object-> vol x nrhs layout
   
*/
template<class Field>
class MultiRHSDeflation
{
public:

  typedef typename Field::scalar_type   scalar;
  typedef typename Field::scalar_object scalar_object;

  int nev;
  std::vector<RealD> eval;
  GridBase *grid;
  uint64_t vol;
  uint64_t words;
  
  deviceVector<scalar> BLAS_E;      //  nev x vol -- the eigenbasis   (up to a 1/sqrt(lambda))
  deviceVector<scalar> BLAS_R;      // nrhs x vol -- the sources
  deviceVector<scalar> BLAS_G;      // nrhs x vol -- the guess
  deviceVector<scalar> BLAS_C;      // nrhs x nev -- the coefficients 
  
  MultiRHSDeflation(){};
  ~MultiRHSDeflation(){ Deallocate(); };
  
  void Deallocate(void)
  {
    nev=0;
    grid=nullptr;
    vol=0;
    words=0;
    BLAS_E.resize(0);
    BLAS_R.resize(0);
    BLAS_C.resize(0);
    BLAS_G.resize(0);
  }
  void Allocate(int _nev,GridBase *_grid)
  {
    nev=_nev;
    grid=_grid;
    vol   = grid->lSites();
    words = sizeof(scalar_object)/sizeof(scalar);
    eval.resize(nev);
    BLAS_E.resize (vol * words * nev );
    std::cout << GridLogMessage << " Allocate for "<<nev<<" eigenvectors and volume "<<vol<<std::endl;
  }
  void ImportEigenVector(Field &evec,RealD &_eval, int ev)
  {
    //    std::cout << " ev " <<ev<<" eval "<<_eval<< std::endl;
    GRID_ASSERT(ev<eval.size());
    eval[ev] = _eval;

    int64_t offset = ev*vol*words;
    autoView(v,evec,AcceleratorRead);
    acceleratorCopyDeviceToDevice(&v[0],&BLAS_E[offset],sizeof(scalar_object)*vol);

  }
  void ImportEigenBasis(std::vector<Field> &evec,std::vector<RealD> &_eval)
  {
    ImportEigenBasis(evec,_eval,0,evec.size());
  }
  // Could use to import a batch of eigenvectors
  void ImportEigenBasis(std::vector<Field> &evec,std::vector<RealD> &_eval, int _ev0, int _nev)
  {
    GRID_ASSERT(_ev0+_nev<=evec.size());

    Allocate(_nev,evec[0].Grid());
    
    // Imports a sub-batch of eigenvectors, _ev0, ..., _ev0+_nev-1
    for(int e=0;e<nev;e++){
      std::cout << "Importing eigenvector "<<e<<" evalue "<<_eval[_ev0+e]<<std::endl;
      ImportEigenVector(evec[_ev0+e],_eval[_ev0+e],e);
    }
  }
  void DeflateSources(std::vector<Field> &source,std::vector<Field> & guess)
  {
    int nrhs = source.size();
    GRID_ASSERT(source.size()==guess.size());
    GRID_ASSERT(grid == guess[0].Grid());
    conformable(guess[0],source[0]);

    int64_t vw = vol * words;

    RealD t0 = usecond();
    BLAS_R.resize(nrhs * vw); // cost free if size doesn't change
    BLAS_G.resize(nrhs * vw); // cost free if size doesn't change
    BLAS_C.resize(nev * nrhs);// cost free if size doesn't change

    /////////////////////////////////////////////
    // Copy in the multi-rhs sources
    /////////////////////////////////////////////
    //    for(int r=0;r<nrhs;r++){
    //      std::cout << " source["<<r<<"] = "<<norm2(source[r])<<std::endl;
    //    }
    for(int r=0;r<nrhs;r++){
      int64_t offset = r*vw;
      autoView(v,source[r],AcceleratorRead);
      acceleratorCopyDeviceToDevice(&v[0],&BLAS_R[offset],sizeof(scalar_object)*vol);
    }

  /*
   * in Fortran column major notation (cuBlas order)
   *
   * Exe = [e1(x)][..][en(x)]
   *
   * Rxr = [r1(x)][..][rm(x)]
   *
   * C_er = E^dag R
   * C_er = C_er / lambda_e 
   * G_xr = Exe Cer
   */
    deviceVector<scalar *> Ed(1);
    deviceVector<scalar *> Rd(1);
    deviceVector<scalar *> Cd(1);
    deviceVector<scalar *> Gd(1);

    scalar * Eh = & BLAS_E[0];
    scalar * Rh = & BLAS_R[0];
    scalar * Ch = & BLAS_C[0];
    scalar * Gh = & BLAS_G[0];

    acceleratorPut(Ed[0],Eh);
    acceleratorPut(Rd[0],Rh);
    acceleratorPut(Cd[0],Ch);
    acceleratorPut(Gd[0],Gh);

    GridBLAS BLAS;

    /////////////////////////////////////////
    // C_er = E^dag R
    /////////////////////////////////////////
    BLAS.gemmBatched(GridBLAS_OP_C,GridBLAS_OP_N, 
    		     nev,nrhs,vw,
		     scalar(1.0),
		     Ed,
		     Rd,
		     scalar(0.0),  // wipe out C
		     Cd);
    BLAS.synchronise();

    GRID_ASSERT(BLAS_C.size()==nev*nrhs);

    std::vector<scalar> HOST_C(BLAS_C.size());      // nrhs . nev -- the coefficients 
    acceleratorCopyFromDevice(&BLAS_C[0],&HOST_C[0],BLAS_C.size()*sizeof(scalar));
    grid->GlobalSumVector(&HOST_C[0],nev*nrhs);
    for(int e=0;e<nev;e++){
      RealD lam(1.0/eval[e]);
      for(int r=0;r<nrhs;r++){
	int off = e+nev*r;
	HOST_C[off]=HOST_C[off] * lam;
	//	std::cout << "C["<<e<<"]["<<r<<"] ="<<HOST_C[off]<< " eval[e] "<<eval[e] <<std::endl;
      }
    }
    acceleratorCopyToDevice(&HOST_C[0],&BLAS_C[0],BLAS_C.size()*sizeof(scalar));

    
    /////////////////////////////////////////
    // Guess G_xr = Exe Cer
    /////////////////////////////////////////
    BLAS.gemmBatched(GridBLAS_OP_N,GridBLAS_OP_N, 
		     vw,nrhs,nev,
		     scalar(1.0),
		     Ed, // x . nev
		     Cd, // nev . nrhs
		     scalar(0.0),
		     Gd);
    BLAS.synchronise();

    ///////////////////////////////////////
    // Copy out the multirhs
    ///////////////////////////////////////
    for(int r=0;r<nrhs;r++){
      int64_t offset = r*vw;
      autoView(v,guess[r],AcceleratorWrite);
      acceleratorCopyDeviceToDevice(&BLAS_G[offset],&v[0],sizeof(scalar_object)*vol);
    }
    RealD t1 = usecond();
    std::cout << GridLogMessage << "MultiRHSDeflation for "<<nrhs<<" sources with "<<nev<<" eigenvectors took " << (t1-t0)/1e3 <<" ms"<<std::endl;
  }
};

NAMESPACE_END(Grid);
