#ifndef DIAGODAVID_H

#define DIAGODAVID_H

#include "source_base/macros.h" // GetRealType

#include "source_base/module_device/device.h" // base_device

#include "source_base/module_device/memory_op.h"// base_device::memory

#include "source_base/module_container/ATen/kernels/lapack.h" // container::kernels

#include "source_hsolver/diag_comm_info.h"

#include "source_hsolver/kernels/hegvd_op.h"

#include <vector>

#include <functional>

namespace hsolver

{

/**

* @class DiagoDavid

* @brief A class that implements the block-Davidson algorithm for solving generalized eigenvalue problems.

*

* The DiagoDavid class provides methods for performing iterative diagonalization using the Davidson algorithm.

* It supports both real and complex data types and can be executed on different devices (CPU or GPU).

*

* @tparam T The data type of the matrices and arrays (e.g., float, double, std::complex<float>, std::complex<double>).

* @tparam Device The device type (e.g., base_device::DEVICE_CPU or DEVICE_GPU).

*/

template <typename T = std::complex<double>, typename Device = base_device::DEVICE_CPU>

class DiagoDavid

{

private:

// Note GetTypeReal<T>::type will

// return T if T is real type(float, double),

// otherwise return the real type of T(complex<float>, std::complex<double>)

using Real = typename GetTypeReal<T>::type;

public:

/**

* @brief Constructor for the DiagoDavid class.

*

* @param[in] precondition_in Pointer to the preconditioning matrix.

* @param[in] nband_in Number of eigenpairs required(i.e. bands).

* @param[in] dim_in Dimension of the matrix.

* @param[in] david_ndim_in Dimension of the reduced basis set of Davidson.

* `david_ndim_in` * `nband_in` is the maximum allowed size of

* the reduced basis set before \b restart of Davidson.

* @param[in] diag_comm_in Communication information for diagonalization.

*

* @tparam T The data type of the matrices and arrays.

* @tparam Device The device type (base_device::DEVICE_CPU or DEVICE_GPU).

*

* @note Auxiliary memory is allocated in the constructor and deallocated in the destructor.

*/

DiagoDavid(const Real* precondition_in,

const int nband_in,

const int dim_in,

const int david_ndim_in,

const diag_comm_info& diag_comm_in);

/**

* @brief Destructor for the DiagoDavid class.

*

* This destructor releases the dynamically allocated memory used by the class members.

* It deletes the basis, hpsi, spsi, hcc, vcc, lagrange_matrix, and eigenvalue arrays.

*

*/

~DiagoDavid();

// declare type of matrix-blockvector functions.

// the function type is defined as a std::function object.

/**

* @brief A function type representing the HX function.

*

* This function type is used to define a matrix-blockvector operator H.

* For eigenvalue problem HX = λX or generalized eigenvalue problem HX = λSX,

* this function computes the product of the Hamiltonian matrix H and a blockvector X.

*

* Called as follows:

* hpsi(X, HX, ld, nvec) where X and HX are (ld, nvec)-shaped blockvectors.

* Result HX = H * X is stored in HX.

*

* @param[out] X Head address of input blockvector of type `T*`.

* @param[in] HX Head address of output blockvector of type `T*`.

* @param[in] ld Leading dimension of blockvector.

* @param[in] nvec Number of vectors in a block.

*

* @warning X and HX are the exact address to read input X and store output H*X,

* @warning both of size ld * nvec.

*/

using HPsiFunc = std::function<void(T*, T*, const int, const int)>;

/**

* @brief A function type representing the SX function.

*

* nrow is leading dimension of spsi, npw is leading dimension of psi, nbands is number of vecs

*

* This function type is used to define a matrix-blockvector operator S.

* For generalized eigenvalue problem HX = λSX,

* this function computes the product of the overlap matrix S and a blockvector X.

*

* @param[in] X Pointer to the input blockvector.

* @param[out] SX Pointer to the output blockvector.

* @param[in] ld_psi Leading dimension of psi and spsi. Dimension of X&SX: ld * nvec.

* @param[in] nvec Number of vectors.

*/

using SPsiFunc = std::function<void(T*, T*, const int, const int)>;

/**

* @brief Performs iterative diagonalization using the David algorithm.

*

* @warning Please see docs of `HPsiFunc` for more information about the hpsi mat-vec interface.

*

* @tparam T The type of the elements in the matrix.

* @tparam Device The device type (CPU or GPU).

* @param hpsi_func The function object that computes the matrix-blockvector product H * psi.

* @param spsi_func The function object that computes the matrix-blockvector product overlap S * psi.

* @param ld_psi The leading dimension of the psi_in array.

* @param psi_in The input wavefunction.

* @param eigenvalue_in The array to store the eigenvalues.

* @param david_diag_thr The convergence threshold for the diagonalization.

* @param david_maxiter The maximum number of iterations for the diagonalization.

* @param ntry_max The maximum number of attempts for the diagonalization restart.

* @param notconv_max The maximum number of bands unconverged allowed.

* @return The total number of iterations performed during the diagonalization.

*

* @note ntry_max is an empirical parameter that should be specified in external routine, default 5

* notconv_max is determined by the accuracy required for the calculation, default 0

*/

int diag(

const HPsiFunc& hpsi_func, // function void hpsi(T*, T*, const int, const int)

const SPsiFunc& spsi_func, // function void spsi(T*, T*, const int, const int, const int)

const int ld_psi, // Leading dimension of the psi input

T *psi_in, // Pointer to eigenvectors

Real* eigenvalue_in, // Pointer to store the resulting eigenvalues

const std::vector<double>& ethr_band, // Convergence threshold for the Davidson iteration

const int david_maxiter, // Maximum allowed iterations for the Davidson method

const int ntry_max = 5, // Maximum number of diagonalization attempts (5 by default)

const int notconv_max = 0); // Maximum number of allowed non-converged eigenvectors

private:

int test_david = 0;

diag_comm_info diag_comm;

/// number of required eigenpairs

const int nband;

/// dimension of the input matrix to be diagonalized

const int dim;

/// maximum dimension of the reduced basis set

const int nbase_x;

/// dimension of the subspace allowed in Davidson

const int david_ndim = 4;

/// number of unconverged eigenvalues

int notconv = 0;

/// precondition for diag, diagonal approximation of matrix A(i.e. Hamilt)

const Real* precondition = nullptr;

Real* d_precondition = nullptr;

/// eigenvalue results

Real* eigenvalue = nullptr;

T *basis = nullptr; /// pointer to basis set(dim, nbase_x), leading dimension = dim

T* hpsi = nullptr; /// the product of H and psi in the reduced basis set

T* spsi = nullptr; /// the Product of S and psi in the reduced basis set

T* hcc = nullptr; /// Hamiltonian on the reduced basis

T* vcc = nullptr; /// eigenvectors of hc

T* lagrange_matrix = nullptr;

/// device type of psi

Device* ctx = {};

base_device::DEVICE_CPU* cpu_ctx = {};

base_device::AbacusDevice_t device = {};

int diag_once(const HPsiFunc& hpsi_func,

const SPsiFunc& spsi_func,

const int dim,

const int nband,

const int ld_psi,

T *psi_in,

Real* eigenvalue_in,

const std::vector<double>& ethr_band,

const int david_maxiter);

/**

* Calculates the preconditioned gradient of the eigenvectors in Davidson method.

*

* @param hpsi_func The function to calculate the matrix-blockvector product H * psi.

* @param spsi_func The function to calculate the matrix-blockvector product overlap S * psi.

* @param dim The dimension of the blockvector.

* @param nbase The current dimension of the reduced basis.

* @param nbase_x The maximum dimension of the reduced basis set.

* @param notconv The number of unconverged eigenpairs.

* @param hpsi The output array for the Hamiltonian H times blockvector psi.

* @param spsi The output array for the overlap matrix S times blockvector psi.

* @param vcc The input array for the eigenvector coefficients.

* @param unconv The array of indices for the unconverged eigenpairs.

* @param eigenvalue The array of eigenvalues.

*/

void cal_grad(const HPsiFunc& hpsi_func,

const SPsiFunc& spsi_func,

const int& dim,

const int& nbase,

const int nbase_x,

const int& notconv,

T* hpsi,

T* spsi,

const T* vcc,

const int* unconv,

const Real* eigenvalue);

/**

* Calculates the elements of the diagonalization matrix for the DiagoDavid class.

*

* @param dim The dimension of the problem.

* @param nbase The current dimension of the reduced basis.

* @param nbase_x The maximum dimension of the reduced basis set.

* @param notconv The number of newly added basis vectors.

* @param hpsi The output array for the Hamiltonian H times blockvector psi.

* @param spsi The output array for the overlap matrix S times blockvector psi.

* @param hcc Pointer to the array where the calculated Hamiltonian matrix elements will be stored.

*/

void cal_elem(const int& dim,

int& nbase,

const int nbase_x,

const int& notconv,

const T* hpsi,

const T* spsi,

T* hcc);

/**

* Refreshes the diagonalization solver by updating the basis and the reduced Hamiltonian.

*

* @param dim The dimension of the problem.

* @param nband The number of bands.

* @param nbase The number of basis states.

* @param nbase_x The maximum dimension of the reduced basis set.

* @param eigenvalue_in Pointer to the array of eigenvalues.

* @param psi_in Pointer to the array of wavefunctions.

* @param ld_psi The leading dimension of the wavefunction array.

* @param hpsi Pointer to the output array for the updated basis set.

* @param spsi Pointer to the output array for the updated basis set (nband-th column).

* @param hcc Pointer to the output array for the updated reduced Hamiltonian.

* @param vcc Pointer to the output array for the updated eigenvector matrix.

*

*/

void refresh(const int& dim,

const int& nband,

int& nbase,

const int nbase_x,

const Real* eigenvalue,

const T *psi_in,

const int ld_psi,

T* hpsi,

T* spsi,

T* hcc,

T* vcc);

/**

* SchmidtOrth function performs orthogonalization of the starting eigenfunction to those already calculated.

* It takes the dimension of the basis, number of bands, index of the current band, starting eigenfunction psi_m,

* lagrange_m array, mm_size, and mv_size as input parameters.

*

* @param dim The dimension of the basis.

* @param nband The number of bands.

* @param m The index of the current band.

* @param spsi Pointer to the starting eigenfunction psi_m.

* @param lagrange_m Pointer to the lagrange_m array.

* @param mm_size The size of the square matrix for future lagranges.

* @param mv_size The size of the lagrange_m array.

*/

void SchmidtOrth(const int& dim,

const int nband,

const int m,

const T* spsi,

T* lagrange_m,

const int mm_size,

const int mv_size);

/**

* @brief Plans the Schmidt orthogonalization for a given number of bands.

*

* @tparam T The type of the elements in the vectors.

* @tparam Device The device on which the computation will be performed.

* @param nband The number of bands.

* @param pre_matrix_mm_m The vector to store the matrix sizes.

* @param pre_matrix_mv_m The vector to store the number of matrix-vector multiplications.

*/

void planSchmidtOrth(const int nband, std::vector<int>& pre_matrix_mm_m, std::vector<int>& pre_matrix_mv_m);

void diag_zhegvx(const int& nbase,

const int& nband,

const T* hcc,

const int& nbase_x,

Real* eigenvalue,

T* vcc);

/**

* @brief Check the convergence of block eigenvectors in the Davidson iteration.

*

* This function determines whether the block eigenvectors have reached convergence

* during the iterative diagonalization process. Convergence is judged based on

* the number of eigenvectors that have not converged and the maximum allowed

* number of such eigenvectors.

*

* @tparam T The data type for the eigenvalues and eigenvectors (e.g., float, double).

* @tparam Device The device type (e.g., base_device::DEVICE_CPU).

* @param ntry The current number of tries for diagonalization.

* @param notconv The current number of eigenvectors that have not converged.

* @param ntry_max The maximum allowed number of tries for diagonalization.

* @param notconv_max The maximum allowed number of eigenvectors that can fail to converge.

* @return true if the eigenvectors are considered converged or the maximum number

* of tries has been reached, false otherwise.

*

* @note Exits the diagonalization loop if either the convergence criteria

* are met or the maximum number of tries is exceeded.

*/

bool check_block_conv(const int &ntry, const int &notconv, const int &ntry_max, const int &notconv_max);

using resmem_complex_op = base_device::memory::resize_memory_op<T, Device>;

using delmem_complex_op = base_device::memory::delete_memory_op<T, Device>;

using setmem_complex_op = base_device::memory::set_memory_op<T, Device>;

using resmem_var_op = base_device::memory::resize_memory_op<Real, Device>;

using delmem_var_op = base_device::memory::delete_memory_op<Real, Device>;

using setmem_var_op = base_device::memory::set_memory_op<Real, Device>;

using syncmem_var_h2d_op = base_device::memory::synchronize_memory_op<Real, Device, base_device::DEVICE_CPU>;

using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<Real, base_device::DEVICE_CPU, Device>;

using syncmem_complex_op = base_device::memory::synchronize_memory_op<T, Device, Device>;

using castmem_complex_op = base_device::memory::cast_memory_op<std::complex<double>, T, Device, Device>;

using syncmem_h2d_op = base_device::memory::synchronize_memory_op<T, Device, base_device::DEVICE_CPU>;

using syncmem_d2h_op = base_device::memory::synchronize_memory_op<T, base_device::DEVICE_CPU, Device>;

// Note that ct_Device is different from base_device!

using ct_Device = typename ct::PsiToContainer<Device>::type;

// using hpsi_info = typename hamilt::Operator<T, Device>::hpsi_info; // Dependence of hpsi removed

const T *one = nullptr, *zero = nullptr, *neg_one = nullptr;

const T one_ = static_cast<T>(1.0), zero_ = static_cast<T>(0.0), neg_one_ = static_cast<T>(-1.0);

};

} // namespace hsolver

#endif