/**
 * @brief Backends doing actual diagonalization.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_DIAG_BACKEND_H
#define TBPLAS_DIAG_BACKEND_H

#include <cmath>
#include <cstring>

#include <eigen3/Eigen/Dense>

#include "base/datatypes.h"
#include "config.h"

#ifdef WITH_LAPACKE
#include "lapacke_wrapper.h"
#endif

#ifdef WITH_FEAST
#include "feast_wrapper.h"
#include <vector>
#endif

namespace tbplas::diag {
using base::complex_t;

/**
 * @brief Abstract backend class
 *
 * @tparam model_t datatype of model and overlap
 */
template <typename model_t>
class AbstractBackend {
protected:
    /// @brief pointer to the model for which calculations will be performed
    const model_t* model_ = nullptr;

    /// @brief pointer to the model which holds the overlap
    const model_t* overlap_ = nullptr;

    /// @brief pointer to DiagConfig instance
    const DiagConfig* config_ = nullptr;

public:
    /**
     * @brief Construct a new AbstractBackend object
     *
     * @param[in] model pointer to model
     * @param[in] overlap pointer to overlap
     * @param[in] config pointer to config
     */
    explicit AbstractBackend(
        const model_t* model,
        const model_t* overlap,
        const DiagConfig* config)
        : model_(model)
        , overlap_(overlap)
        , config_(config)
    {
    }

    /**
     * @brief Get the number of orbitals
     *
     * @return number of orbitals
     */
    [[nodiscard]] size_t get_num_orbitals() const
    {
        return model_->get_num_orbitals();
    }

    /**
     * @brief Get the number of states to calculate.
     *
     * @return number of states.
     *
     * @note For backends other than feast, the number of orbitals of the model
     *   is returned. For feast, config.feast_num_states is returned if it is
     *   larger than 0. Otherwise, half of the number of orbitals of the model
     *   is returned.
     */
    [[nodiscard]] virtual size_t get_num_states() const = 0;

    /**
     * @brief Calculate eigenvalues of dense Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    virtual void calc_eigen_values(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues)
        = 0;

    /**
     * @brief Calculate eigenvalues and eigenvectors of dense Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point in eV
     * @param[out] eigenvectors eigenvectors at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    virtual void calc_eigen_all(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues,
        Eigen::MatrixXcd& eigenvectors)
        = 0;

    /**
     * @brief Get the dense overlap matrix for post-processing.
     *
     * @return reference to dense overlap matrix.
     *
     * @note This function must be called after calling calc_eigen_values or calc_eigen_all.
     */
    virtual const Eigen::MatrixXcd& get_sk() = 0;
};

/**
 * @brief Default backend based on Eigen
 *
 * @tparam model_t datatype of model and overlap
 */
template <typename model_t>
class DefaultBackend final : public AbstractBackend<model_t> {

private:
    /// @brief dense Hamiltonian matrix
    Eigen::MatrixXcd hk_;

    /// @brief dense overlap matrix
    Eigen::MatrixXcd sk_;

    /// @brief solver for dense Hamiltonian matrix without overlap
    Eigen::SelfAdjointEigenSolver<Eigen::MatrixXcd> zheev_;

    /// @brief solver for dense Hamiltonian matrix with overlap
    Eigen::GeneralizedSelfAdjointEigenSolver<Eigen::MatrixXcd> zhegv_;

public:
    /**
     * @brief Construct a new DefaultBackend object
     *
     * @param[in] model pointer to model
     * @param[in] overlap pointer to overlap
     * @param[in] config pointer to config
     *
     */
    explicit DefaultBackend(
        const model_t* model,
        const model_t* overlap,
        const DiagConfig* config)
        : AbstractBackend<model_t>(model, overlap, config)
    {
        size_t num_orb = model->get_num_orbitals();
        hk_.resize(num_orb, num_orb);
        if (overlap != nullptr) {
            sk_.resize(num_orb, num_orb);
        }
    }

    /**
     * @brief Get the number of states to calculate.
     *
     * @return number of states.
     */
    [[nodiscard]] size_t get_num_states() const final
    {
        return AbstractBackend<model_t>::model_->get_num_orbitals();
    }

    /**
     * @brief Calculate eigenvalues of dense Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    void calc_eigen_values(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues) final
    {
        const model_t* model = AbstractBackend<model_t>::model_;
        const model_t* overlap = AbstractBackend<model_t>::overlap_;
        model->build_ham_dense(kpt, hk_, convention);
        if (overlap == nullptr) {
            zheev_.compute(hk_, Eigen::EigenvaluesOnly);
            eigenvalues = zheev_.eigenvalues();
        } else {
            overlap->build_ham_dense(kpt, sk_, convention);
            zhegv_.compute(hk_, sk_, Eigen::EigenvaluesOnly);
            eigenvalues = zhegv_.eigenvalues();
        }
    }

    /**
     * @brief Calculate eigenvalues and eigenvectors of dense Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point in eV
     * @param[out] eigenvectors eigenvectors at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    void calc_eigen_all(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues,
        Eigen::MatrixXcd& eigenvectors) final
    {
        const model_t* model = AbstractBackend<model_t>::model_;
        const model_t* overlap = AbstractBackend<model_t>::overlap_;
        model->build_ham_dense(kpt, hk_, convention);
        if (overlap == nullptr) {
            zheev_.compute(hk_);
            eigenvalues = zheev_.eigenvalues();
            eigenvectors = zheev_.eigenvectors();
        } else {
            overlap->build_ham_dense(kpt, sk_, convention);
            zhegv_.compute(hk_, sk_);
            eigenvalues = zhegv_.eigenvalues();
            eigenvectors = zhegv_.eigenvectors();
        }
    }

    /**
     * @brief Get the dense overlap matrix for post-processing.
     *
     * @return reference to dense overlap matrix.
     *
     * @note This function must be called after calling calc_eigen_values or calc_eigen_all.
     */
    const Eigen::MatrixXcd& get_sk() final
    {
        return sk_;
    }
};

#ifdef WITH_LAPACKE
/**
 * @brief Backend based on LAPACKE
 *
 * @tparam model_t datatype of model and overlap
 */
template <typename model_t>
class LAPACKEBackend final : public AbstractBackend<model_t> {

private:
    /// @brief dense Hamiltonian matrix
    Eigen::MatrixXcd hk_;

    /// @brief dense overlap matrix
    Eigen::MatrixXcd sk_;

public:
    /**
     * @brief Construct a new DefaultBackend object
     *
     * @param[in] model pointer to model
     * @param[in] overlap pointer to overlap
     * @param[in] config pointer to config
     *
     */
    explicit LAPACKEBackend(
        const model_t* model,
        const model_t* overlap,
        const DiagConfig* config)
        : AbstractBackend<model_t>(model, overlap, config)
    {
        size_t num_orb = model->get_num_orbitals();
        hk_.resize(num_orb, num_orb);
        if (overlap != nullptr) {
            sk_.resize(num_orb, num_orb);
        }
    }

    /**
     * @brief Get the number of states to calculate.
     *
     * @return number of states.
     */
    [[nodiscard]] size_t get_num_states() const final
    {
        return AbstractBackend<model_t>::model_->get_num_orbitals();
    }

    /**
     * @brief Calculate eigenvalues of dense Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    void calc_eigen_values(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues) final
    {
        // Build Hamiltonian and overlap
        const model_t* model = AbstractBackend<model_t>::model_;
        const model_t* overlap = AbstractBackend<model_t>::overlap_;
        model->build_ham_dense(kpt, hk_, convention);
        if (overlap != nullptr) {
            overlap->build_ham_dense(kpt, sk_, convention);
        }

        // Check if arrays have been allocated
        int num_orb = model->get_num_orbitals();
        if (eigenvalues.size() != num_orb) {
            eigenvalues.resize(num_orb);
        }

        // Call lapacke
        auto* hk_addr = reinterpret_cast<lapack_complex_double*>(hk_.data());
        auto* sk_addr = reinterpret_cast<lapack_complex_double*>(sk_.data());
        if (overlap == nullptr) {
            lapacke_zheev_wrapper(
                num_orb,
                hk_addr,
                eigenvalues.data(),
                nullptr,
                false);
        } else {
            lapacke_zhegv_wrapper(
                num_orb,
                hk_addr,
                sk_addr,
                eigenvalues.data(),
                nullptr,
                false);
        }
    }

    /**
     * @brief Calculate eigenvalues and eigenvectors of dense Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point in eV
     * @param[out] eigenvectors eigenvectors at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    void calc_eigen_all(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues,
        Eigen::MatrixXcd& eigenvectors) final
    {
        // Build Hamiltonian and overlap
        const model_t* model = AbstractBackend<model_t>::model_;
        const model_t* overlap = AbstractBackend<model_t>::overlap_;
        model->build_ham_dense(kpt, hk_, convention);
        if (overlap != nullptr) {
            overlap->build_ham_dense(kpt, sk_, convention);
        }

        // Check if arrays have been allocated
        int num_orb = model->get_num_orbitals();
        if (eigenvalues.size() != num_orb) {
            eigenvalues.resize(num_orb);
        }
        if (eigenvectors.size() != num_orb * num_orb) {
            eigenvectors.resize(num_orb, num_orb);
        }

        // Call lapacke
        auto* hk_addr = reinterpret_cast<lapack_complex_double*>(hk_.data());
        auto* sk_addr = reinterpret_cast<lapack_complex_double*>(sk_.data());
        auto* vec_addr = reinterpret_cast<lapack_complex_double*>(eigenvectors.data());
        if (overlap == nullptr) {
            lapacke_zheev_wrapper(
                num_orb,
                hk_addr,
                eigenvalues.data(),
                vec_addr,
                true);
        } else {
            lapacke_zhegv_wrapper(
                num_orb,
                hk_addr,
                sk_addr,
                eigenvalues.data(),
                vec_addr,
                true);
        }
    }

    /**
     * @brief Get the dense overlap matrix for post-processing.
     *
     * @return reference to dense overlap matrix.
     *
     * @note This function must be called after calling calc_eigen_values or calc_eigen_all.
     */
    const Eigen::MatrixXcd& get_sk() final
    {
        return sk_;
    }
};
#endif // WITH_LAPACKE

#ifdef WITH_FEAST
/**
 * @brief Backend based on feast dense solver.
 *
 * @tparam model_t datatype of model
 */
template <typename model_t>
class FeastDenseBackend final : public AbstractBackend<model_t> {
private:
    /// @brief dense Hamiltonian matrix
    Eigen::MatrixXcd hk_;

    /// @brief dense overlap matrix
    Eigen::MatrixXcd sk_;

    /// @brief buffer for eigenvectors, dimension (num_orb, subspace_size), allocated on-the-fly
    Eigen::MatrixXcd eig_vec_;

    /// @brief buffer for eigenvalues, dimension (subspace_size, 1), allocated on-the-fly
    Eigen::VectorXd eig_val_;

    /// @brief buffer for residuals, dimension (subspace_size, 1), allocated on-the-fly
    Eigen::VectorXd residual_;

    /// @brief feast input parameters
    /// @note MKL has its own implementation of FEAST, which is SIGNIFICANTLY different from
    ///   the original version. For example, it uses fpm[128] instead of fpm[64]. We further
    ///   make it fpm[256] to avoid potential segfaults due to lp64/ilp64. In fact, the segfaults
    ///   has cost us two weeks to debug.
    int fpm_[256];

    /**
     * @brief Call FEAST solver to evaluate eigenvalues and eigenvectors
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @return number of eigenstates found
     */
    [[nodiscard]] int calc_eigen_core(
        const Eigen::Vector3d& kpt,
        const int& convention)
    {
        // Aliases
        const model_t* model = AbstractBackend<model_t>::model_;
        const model_t* overlap = AbstractBackend<model_t>::overlap_;
        const DiagConfig* config = AbstractBackend<model_t>::config_;

        // Check if buffers have been allocated
        size_t num_orb = model->get_num_orbitals();
        size_t num_states = get_num_states();
        auto subspace_size = static_cast<size_t>(num_states * config->feast_m0_ratio);
        if (eig_val_.size() != subspace_size) {
            eig_val_.resize(subspace_size);
        }
        if (eig_vec_.size() != num_orb * subspace_size) {
            eig_vec_.resize(num_orb, subspace_size);
        }
        if (residual_.size() != subspace_size) {
            residual_.resize(subspace_size);
        }

        // Set Hamiltonian
        model->build_ham_dense(kpt, hk_, convention);
        if (overlap != nullptr) {
            overlap->build_ham_dense(kpt, sk_, convention);
        }

        // Set feast arguments
        FeastDenseArgs args;
        args.num_orb = num_orb;
        args.hk = reinterpret_cast<double*>(hk_.data());
        if (overlap != nullptr) {
            args.sk = reinterpret_cast<double*>(sk_.data());
        }
        args.fpm = &fpm_[0];
        args.eps_out = config->feast_eps;
        args.e_min = config->feast_e_min;
        args.e_max = config->feast_e_max;
        args.subspace_size = subspace_size;
        args.num_states = num_states;
        args.eig_val = eig_val_.data();
        args.eig_vec = reinterpret_cast<double*>(eig_vec_.data());
        args.residual = residual_.data();

        // Call feast
        if (overlap == nullptr) {
            zfeast_heev_wrapper(&args);
        } else {
            zfeast_hegv_wrapper(&args);
        }

        // Reuse results as initial guess of next run
        if (config->feast_reuse_results) {
            fpm_[4] = 1;
        }

        // Return the number of eigenvalues found
        return args.num_states;
    }

public:
    /**
     * @brief Construct a new FeastDenseBackend object
     *
     * @param[in] model pointer to model
     * @param[in] overlap pointer to overlap
     * @param[in] config pointer to config
     *
     */
    explicit FeastDenseBackend(
        const model_t* model,
        const model_t* overlap,
        const DiagConfig* config)
        : AbstractBackend<model_t>(model, overlap, config)
    {
        // Initialize feast
        feast_init_wrapper(&fpm_[0]);
        auto iter = config->feast_fpm.begin();
        for (; iter != config->feast_fpm.end(); iter++) {
            auto i = iter->first;
            auto j = iter->second;
            assert(0 <= i && i < 64);
            fpm_[i] = j;
        }

        // Allocate arrays
        // Buffers will be allocated on-the-fly.
        size_t num_orb = model->get_num_orbitals();
        hk_.resize(num_orb, num_orb);
        if (overlap != nullptr) {
            sk_.resize(num_orb, num_orb);
        }
    }

    /**
     * @brief Get the number of states to calculate.
     *
     * @return number of states.
     */
    [[nodiscard]] size_t get_num_states() const final
    {
        const DiagConfig* config = AbstractBackend<model_t>::config_;
        const model_t* model = AbstractBackend<model_t>::model_;
        size_t num_states = 0;
        if (config->feast_num_states > 0) {
            num_states = config->feast_num_states;
        } else {
            num_states = model->get_num_orbitals() / 2;
        }
        return num_states;
    }

    /**
     * @brief Calculate eigenvalues of dense Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    void calc_eigen_values(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues) final
    {
        // Check if arrays have been allocated
        size_t num_states = get_num_states();
        if (eigenvalues.size() != num_states) {
            eigenvalues.resize(num_states);
        }

        // Set eigenvalues to NaN to mark diverged results.
        eigenvalues.array() = NAN;

        // Call FEAST and copy results
        int num_states_found = calc_eigen_core(kpt, convention);
        int num_states_copy = num_states_found <= num_states ? num_states_found : num_states;
        std::memcpy(eigenvalues.data(), eig_val_.data(), num_states_copy * sizeof(double));
    }

    /**
     * @brief Calculate eigenvalues and eigenvectors of dense Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point in eV
     * @param[out] eigenvectors eigenvectors at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    void calc_eigen_all(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues,
        Eigen::MatrixXcd& eigenvectors) final
    {
        // Check if arrays have been allocated
        const model_t* model = AbstractBackend<model_t>::model_;
        size_t num_orb = AbstractBackend<model_t>::model_->get_num_orbitals();
        size_t num_states = get_num_states();
        if (eigenvalues.size() != num_states) {
            eigenvalues.resize(num_states);
        }
        if (eig_vec_.size() != num_orb * num_states) {
            eig_vec_.resize(num_orb, num_states);
        }

        // Set eigenvalues and eigenstates to NaN to mark diverged results.
        eigenvalues.array() = NAN;
        eigenvectors.array() = NAN;

        // Call FEAST and copy results
        int num_states_found = calc_eigen_core(kpt, convention);
        int num_states_copy = num_states_found <= num_states ? num_states_found : num_states;
        std::memcpy(eigenvalues.data(), eig_val_.data(), num_states_copy * sizeof(double));
        std::memcpy(eigenvectors.data(), eig_vec_.data(), num_states_copy * num_orb * sizeof(complex_t));
    }

    /**
     * @brief Get the dense overlap matrix for post-processing.
     *
     * @return reference to dense overlap matrix.
     *
     * @note This function must be called after calling calc_eigen_values or calc_eigen_all.
     */
    const Eigen::MatrixXcd& get_sk() final
    {
        return sk_;
    }
};

/**
 * @brief Backend based on feast sparse solver.
 *
 * @tparam model_t datatype of model
 */
template <typename model_t>
class FeastCSRBackend final : public AbstractBackend<model_t> {
private:
    /// @brief sparse Hamiltonian
    std::vector<int> indptr_h_;
    std::vector<int> indices_h_;
    Eigen::VectorXcd hk_;

    /// @brief sparse overlap matrix
    std::vector<int> indptr_s_;
    std::vector<int> indices_s_;
    Eigen::VectorXcd sk_;

    /// @brief dense overlap matrix for post-processing
    Eigen::MatrixXcd sk_dense_;

    /// @brief buffer for eigenvectors, dimension (num_orb, subspace_size), allocated on-the-fly
    Eigen::MatrixXcd eig_vec_;

    /// @brief buffer for eigenvalues, dimension (subspace_size, 1), allocated on-the-fly
    Eigen::VectorXd eig_val_;

    /// @brief buffer for residuals, dimension (subspace_size, 1), allocated on-the-fly
    Eigen::VectorXd residual_;

    /// @brief feast input parameters
    /// @note MKL has its own implementation of FEAST, which is SIGNIFICANTLY different from
    ///   the original version. For example, it uses fpm[128] instead of fpm[64]. We further
    ///   make it fpm[256] to avoid potential segfaults due to lp64/ilp64. In fact, the segfaults
    ///   has cost us two weeks to debug.
    int fpm_[256];

    /// @brief Build sparse hk_ or sk_.
    void build_mat_csr(
        const Eigen::Vector3d& kpt,
        const int& convention,
        const bool& is_overlap)
    {
        // Make aliases depending on is_overlap
        const model_t* target = AbstractBackend<model_t>::model_;
        std::vector<int>* indptr = &indptr_h_;
        std::vector<int>* indices = &indices_h_;
        Eigen::VectorXcd* value = &hk_;
        if (is_overlap) {
            target = AbstractBackend<model_t>::overlap_;
            indptr = &indptr_s_;
            indices = &indices_s_;
            value = &sk_;
        }

        target->build_ham_csr(*indptr, *indices, *value, kpt, convention);
        // FORTRAN array index begins from 1!
        for (size_t i = 0; i < indptr->size(); ++i) {
            (*indptr)[i]++;
        }
        for (size_t i = 0; i < indices->size(); ++i) {
            (*indices)[i]++;
        }
    }

    /**
     * @brief Call FEAST solver to evaluate eigenvalues and eigenvectors
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @return number of eigenstates found
     */
    [[nodiscard]] int calc_eigen_core(
        const Eigen::Vector3d& kpt,
        const int& convention)
    {
        // Aliases
        const model_t* model = AbstractBackend<model_t>::model_;
        const model_t* overlap = AbstractBackend<model_t>::overlap_;
        const DiagConfig* config = AbstractBackend<model_t>::config_;

        // Check if buffers have been allocated
        size_t num_orb = model->get_num_orbitals();
        size_t num_states = get_num_states();
        auto subspace_size = static_cast<size_t>(num_states * config->feast_m0_ratio);
        if (eig_val_.size() != subspace_size) {
            eig_val_.resize(subspace_size);
        }
        if (eig_vec_.size() != num_orb * subspace_size) {
            eig_vec_.resize(num_orb, subspace_size);
        }
        if (residual_.size() != subspace_size) {
            residual_.resize(subspace_size);
        }

        // Set Hamiltonian
        build_mat_csr(kpt, convention, false);
        if (overlap != nullptr) {
            build_mat_csr(kpt, convention, true);
        }

        // Set feast args
        FeastCsrArgs args;
        args.num_orb = num_orb;
        args.values_h = reinterpret_cast<double*>(hk_.data());
        args.indptr_h = indptr_h_.data();
        args.indices_h = indices_h_.data();
        if (overlap != nullptr) {
            args.values_s = reinterpret_cast<double*>(sk_.data());
            args.indptr_s = indptr_s_.data();
            args.indices_s = indices_s_.data();
        }
        args.fpm = &fpm_[0];
        args.eps_out = config->feast_eps;
        args.e_min = config->feast_e_min;
        args.e_max = config->feast_e_max;
        args.subspace_size = subspace_size;
        args.num_states = num_states;
        args.eig_val = eig_val_.data();
        args.eig_vec = reinterpret_cast<double*>(eig_vec_.data());
        args.residual = residual_.data();

        // Call feast
        if (overlap == nullptr) {
            zfeast_hcsrev_wrapper(&args);
        } else {
            zfeast_hcsrgv_wrapper(&args);
        }

        // Reuse results as initial guess of next run
        if (config->feast_reuse_results) {
            fpm_[4] = 1;
        }

        // Return the number of eigenvalues found
        return args.num_states;
    }

public:
    /**
     * @brief Construct a new FeastCSRBackend object
     *
     * @param[in] model pointer to model
     * @param[in] overlap pointer to overlap
     * @param[in] config pointer to config
     */
    explicit FeastCSRBackend(
        const model_t* model,
        const model_t* overlap,
        const DiagConfig* config)
        : AbstractBackend<model_t>(model, overlap, config)
    {
        // Initialize feast
        feast_init_wrapper(&fpm_[0]);
        auto iter = config->feast_fpm.begin();
        for (; iter != config->feast_fpm.end(); iter++) {
            auto i = iter->first;
            auto j = iter->second;
            assert(0 <= i && i < 64);
            fpm_[i] = j;
        }

        // Allocate arrays
        // Sparse Hamiltonian, overlap and buffers will be allocated on-the-fly.
        size_t num_orb = model->get_num_orbitals();
        if (overlap != nullptr) {
            sk_dense_.resize(num_orb, num_orb);
        }
    }

    /**
     * @brief Get the number of states to calculate.
     *
     * @return number of states.
     */
    [[nodiscard]] size_t get_num_states() const final
    {
        const DiagConfig* config = AbstractBackend<model_t>::config_;
        const model_t* model = AbstractBackend<model_t>::model_;
        size_t num_states = 0;
        if (config->feast_num_states > 0) {
            num_states = config->feast_num_states;
        } else {
            num_states = model->get_num_orbitals() / 2;
        }
        return num_states;
    }

    /**
     * @brief Calculate eigenvalues of sparse Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    void calc_eigen_values(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues) final
    {
        // Check if arrays have been allocated
        size_t num_states = get_num_states();
        if (eigenvalues.size() != num_states) {
            eigenvalues.resize(num_states);
        }

        // Set eigenvalues to NaN to mark diverged results.
        eigenvalues.array() = NAN;

        // Call FEAST and copy results
        int num_states_found = calc_eigen_core(kpt, convention);
        int num_states_copy = num_states_found <= num_states ? num_states_found : num_states;
        std::memcpy(eigenvalues.data(), eig_val_.data(), num_states_copy * sizeof(double));
    }

    /**
     * @brief Calculate eigenvalues and eigenvectors of sparse Hamiltonian for given k-point
     *
     * @param[in] kpt fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     * @param[out] eigenvalues eigenvalues at k-point in eV
     * @param[out] eigenvectors eigenvectors at k-point
     *
     * @note Some calculations require specific convention, so we cannot use
     * config.convention here as it may be incompatible with the calculation
     * type.
     */
    void calc_eigen_all(
        const Eigen::Vector3d& kpt,
        const int& convention,
        Eigen::VectorXd& eigenvalues,
        Eigen::MatrixXcd& eigenvectors) final
    {
        // Check if arrays have been allocated
        const model_t* model = AbstractBackend<model_t>::model_;
        size_t num_orb = AbstractBackend<model_t>::model_->get_num_orbitals();
        size_t num_states = get_num_states();
        if (eigenvalues.size() != num_states) {
            eigenvalues.resize(num_states);
        }
        if (eig_vec_.size() != num_orb * num_states) {
            eig_vec_.resize(num_orb, num_states);
        }

        // Set eigenvalues and eigenstates to NaN to mark diverged results.
        eigenvalues.array() = NAN;
        eigenvectors.array() = NAN;

        // Call FEAST and copy results
        int num_states_found = calc_eigen_core(kpt, convention);
        int num_states_copy = num_states_found <= num_states ? num_states_found : num_states;
        std::memcpy(eigenvalues.data(), eig_val_.data(), num_states_copy * sizeof(double));
        std::memcpy(eigenvectors.data(), eig_vec_.data(), num_states_copy * num_orb * sizeof(complex_t));
    }

    /**
     * @brief Get the dense overlap matrix for post-processing.
     *
     * @return reference to dense overlap matrix.
     *
     * @note This function must be called after calling calc_eigen_values or calc_eigen_all.
     */
    const Eigen::MatrixXcd& get_sk() final
    {
        // indptr_s and indices_s are for FORTRAN and starts from 1.
        sk_dense_.setZero();
        int num_orb = AbstractBackend<model_t>::model_->get_num_orbitals();
        for (int i = 0; i < num_orb; ++i) {
            for (int j = indptr_s_[i] - 1; j < indptr_s_[i + 1] - 1; ++j) {
                int k = indices_s_[j] - 1;
                sk_dense_(i, k) = sk_[j];
            }
        }
        return sk_dense_;
    }
};

#endif

} // namespace tbplas::diag
#endif