/* Copyright (c) 2015  Gerald Knizia
 * 
 * This file is part of the IboView program (see: http://www.iboview.org)
 * 
 * IboView 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, version 3.
 * 
 * IboView 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 details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with bfint (LICENSE). If not, see http://www.gnu.org/licenses/
 * 
 * Please see IboView documentation in README.txt for:
 * -- A list of included external software and their licenses. The included
 *    external software's copyright is not touched by this agreement.
 * -- Notes on re-distribution and contributions to/further development of
 *    the IboView software
 */

#ifndef CT8K_BASISSET_H
#define CT8K_BASISSET_H

#include <vector>
#include "Ir.h"
#include "CxPodArray.h"
#include "CtBasisShell.h"

namespace ct {

enum FBasisContext {
   BASIS_Orbital,
   BASIS_JFit,
   BASIS_JkFit,
   BASIS_Mp2Fit,
   BASIS_CcsdFit,
   BASIS_F12RI,
   BASIS_Guess,
   BASIS_IaoFit,
   BASIS_Multipoles
};

std::string BasisContextName( FBasisContext Context );

struct FMatrixView;
struct FAtomSet;

// IR driver-level interface to a basis set. Semi-lightweight.
// supposed to be generated from the program's native basis set
// format (FBasisSet here) on demand for interfacing between IR
// and the program.
struct FRawBasis : public FIntrusivePtrDest
{
   typedef TArray<ir::FRawShell>
      FShellArray;
   FShellArray
      Shells;
   TArray<size_t>
      // offset of Shell[i] with respect to underlying basis in terms of
      // individual functions. Contains one additional element giving the
      // total number of functions.
      ShellOffsets,
      // shells of atom #i start at CenterOffsets[i].
      CenterOffsets,
      // [i]: center index of shell [i].
      ShellCenters;
   size_t
      // largest number of functions in a shell occuring in the basis.
      // May be useful for allocating intermediates.
      nMaxFnPerShell;
   size_t nFn() const { return ShellOffsets.back(); }
   size_t nSh() const { return Shells.size(); }
   size_t nFn(size_t iSh) const { return ShellOffsets[iSh+1]-ShellOffsets[iSh]; }
   size_t iFn(size_t iSh) const { return ShellOffsets[iSh]; }
   // index of first shell on center iCen
   size_t iCenSh(size_t iCen) const { return CenterOffsets[iCen]; }
   size_t nCenSh(size_t iCen) const { return CenterOffsets[iCen+1] - CenterOffsets[iCen]; }
   // index of first function center iCen
   size_t iCenFn(size_t iCen) const { return iFn(iCenSh(iCen)); }
   size_t nCenFn(size_t iCen) const { return iCenFn(iCen+1) - iCenFn(iCen); }
   // total number of centers
   size_t nCen() const { return CenterOffsets.size() - 1; }

   // center index of a given shell.
   size_t iShCen(size_t iSh) const { return ShellCenters[iSh]; }

   // construct derived data (ShellOffsets, CenterOffsets, nMaxFnPerShell) from
   // this->Shells and the shell center indices given as argument. pShCen[iSh]
   // should give the center id for Shells[iSh], and those must be consecutive.
   void Finalize(size_t const *pShCen);
};
typedef boost::intrusive_ptr<FRawBasis>
   FRawBasisPtr;

// program native format of a basis set.
struct FBasisSet : FIntrusivePtrDest
{
   typedef std::vector<FBasisShell>
      FBasisShellArray;
   FBasisShellArray
      Shells;  // Gaussian shells centered on different atoms.
   FBasisContext
      Context; // for informative purposes only.
   std::string
      Name; // for informative purposes only.

   void Print(std::ostream &out) const;
   size_t nFn() const;
   size_t nPrimitiveGtos() const; // for decoration purposes only.
   size_t nFnOfLargestShell() const;
   uint nMaxL() const; // find largest angular momentum in the basis.

   // rotate and/or translate the basis (i.e., the positions of the atoms) in real space. R is a
   // 3x3 matrix and 'd' is a 3-vector such that r' = R (r - d) gives the transformed atomic positions.
   void Transform_3x4(double const *R, double const *d, FMemoryStack &Mem);

   // transform a set of MO coefficients (in C) in such a way that if the molecule and its
   // basis are rotated in real space by the 3x3 matrix R (via FBasisSet Transform_3x4), then
   // this function computes the MO coefficients corresponding to the new alignment.
   void TransformMoCoeffs_3x4(FMatrixView C, double const *R, FMemoryStack &Mem) const;

   // construct from atom set and context
   FBasisSet(FAtomSet const &AtomSet, FBasisContext Context_);
   // construct from explicit list of shells. Name and Context for informative purposes only.
   FBasisSet(FBasisShell const *pShells_, size_t nShells_, FBasisContext Context_, std::string const &Name);
public:
   FRawBasisPtr
      // Only valid if RebuildInterfacingData has been called after the basis
      // was last modified. If constructed regularly, this is done
      // automatically.
      pRawBasis;

   // rebuilds the data for low-level interfaces (i.e., the data in this section)
   // from this->Shells
   void RebuildInterfacingData();

   FRawBasisPtr MakeRawBasis();
private:
   // makes *this a basis set corresponding to the basis descriptions in
   // AtomSet. Attempts to load the neccessary external data. Throws
   // std::runtime_error if failed.
   void LoadFromAtomSet( FAtomSet const &AtomSet, FBasisContext Context );

   void MakeAtomOffsets( size_t *&pAtomShellOffsets, size_t *&pAtomBfnOffsets, size_t nAtoms_, FMemoryStack &Mem ) const;
   void MakeShellOffsets( size_t *&pShellOffsets, FMemoryStack &Mem ) const;
   void Finalize();
};

void Trafo3x4(FVector3 &InOut, double const *R, double const *d);


inline std::ostream &operator << ( std::ostream &out, FBasisSet const &BasisSet ) {
   BasisSet.Print(out);
   return out;
}

typedef boost::intrusive_ptr<FBasisSet>
   FBasisSetPtr;

// generalized 2-index integral kernel: effectively a functor which
// can calculate shell doublets for two shell objects. Used for making
// one-electron matrices.
struct FKrn2i : public FIntrusivePtrDest
{
   virtual void EvalInt2e2c(double *pOut, size_t StrideA, size_t StrideC, ir::FRawShell const *pA, ir::FRawShell const *pC, double Prefactor, bool Add, ct::FMemoryStack &Mem) const = 0;
   // Accumulate derivative contraction
   //     Grad[Nuc,xyz] += Prefactor * \sum_{ac} rdm[a,c] d/d[Nuc,xyz] (a|krn|c).
   // Parameters:
   //   - pGrad[3*iCenter + ixyz]: Location in which the output gradient for center iCenter and ixyz=0...2 is stored.
   //   - pRdmAC[StrideA * ia + StrideC * ic]: gradient coefficients for *pA, *pC,
   //     where ia/ic go over the contracted functions of *pA, *pC.
   //   - iCenA, iCenC: center indices in pGrad itself which belong
   // default implementation just crashes and says "can't do this".
   virtual void AccCoreDeriv1d(double *pGrad, double const *pRdmAC, size_t StrideA, size_t StrideC, ir::FRawShell const *pA, size_t iCenterA, ir::FRawShell const *pC, size_t iCenterC, double Prefactor, ct::FMemoryStack &Mem) const;
};
typedef boost::intrusive_ptr<FKrn2i>
   FKrn2iPtr;

void MakeIntMatrix( FMatrixView &Out, FBasisSet const &RowBasis,
   FBasisSet const &ColBasis, FKrn2i const &Krn2i, FMemoryStack &Mem, double Prefactor=1., bool Add=false);
void MakeIntMatrix( FMatrixView &Out, FRawBasis const &RowBasis,
   FRawBasis const &ColBasis, FKrn2i const &Krn2i, FMemoryStack &Mem, double Prefactor=1., bool Add=false);
// increment gradient in pGrad[ixyz + 3 *iCenter] by \sum{a,b} Rdm[a,b] d/d[Eu] Krn2i[a,b].
// If &RowBasis == &ColBasis it is assumed that Rdm is symmetric. If it is not symmetric, pls
// symmetrize first (not checked!).
void AccGradient2ix( double *pGrad, FMatrixView const &Rdm, FRawBasis const &RowBasis,
   FRawBasis const &ColBasis, FKrn2i const &Krn2i, FMemoryStack &Mem, double Prefactor=1.);

// evaluate <a|krn Laplace^LaplaceOrder|b> by directly calling
// a IR integral kernel.
struct FKrn2i_Direct : public FKrn2i
{
   FKrn2i_Direct(ir::FIntegralKernel const *pIrKernel_, uint LaplaceOrder=0, double Prefactor=1.);
   void EvalInt2e2c(double *pOut, size_t StrideA, size_t StrideC, ir::FRawShell const *pA, ir::FRawShell const *pC, double Prefactor, bool Add, ct::FMemoryStack &Mem) const;
//    void EvalInt2e2c1d(double *pOutDa, double *pOutDb, size_t StrideA, size_t StrideC, size_t StrideDeriv, ir::FRawShell const *pA, ir::FRawShell const *pC, double Prefactor, bool Add, ct::FMemoryStack &Mem) const;
   void AccCoreDeriv1d(double *pGrad, double const *pRdmAC, size_t StrideA, size_t StrideC, ir::FRawShell const *pA, size_t iCenterA, ir::FRawShell const *pC, size_t iCenterC, double Prefactor, ct::FMemoryStack &Mem) const;
protected:
   ir::FIntegralKernel const
      *m_pIrKernel;
   uint
      m_LaplaceOrder;
   double
      m_Prefactor;
};

// evaluate <a| v(r) |b> where v(r) is sum of the fields emitted via pIrKernel
// for the sum of the supplied spherical point multipoles.
// Can be used, for example, to make nuclear attraction integrals.
struct FKrn2i_PointMultipoles : public FKrn2i
{
   // each point should provide 2*l+1 coefficients; one for each of its spherical
   // components. Slm order is equal to the order used for basis functions.
   // (in particular, dipole components are x,y,z, not x,z,y.).
   struct FPoint {
      FVector3 vPos;
      uint l;
      ptrdiff_t iCenter;
      // ^- for gradients: identify the atomic center index on which this multipole sits. Set to -1 if None.
      //    note: even if having external lattices, one could just append their gradient array to the QM system gradient array
      //    and then do both in one go.
   };
   FKrn2i_PointMultipoles(ir::FIntegralKernel const *pIrKernel_, FPoint const *pPoints, double const *pCoeffs, size_t nPoints);
   void EvalInt2e2c(double *pOut, size_t StrideA, size_t StrideC, ir::FRawShell const *pA, ir::FRawShell const *pC, double Prefactor, bool Add, ct::FMemoryStack &Mem) const;
protected:
   ir::FIntegralKernel const
      *m_pIrKernel;
   TArray<ir::FRawShell>
      m_PointShells;
   TArray<double>
      m_Data;
   TArray<ptrdiff_t>
      m_PointCenters;
   double
      m_Exp;
   uint
      m_MaxL;
};



} // namespace ct

#endif // CT8K_BASISSET_H
