/**
 * @brief Classes for calculating response functions.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_DIAG_LINDHARD_H
#define TBPLAS_DIAG_LINDHARD_H

#include <memory>

#include <eigen3/Eigen/Dense>

#include "base.h"
#include "base/consts.h"
#include "base/datatypes.h"
#include "base/funcs.h"
#include "base/kpoints.h"
#include "base/lattice.h"
#include "builder/base.h"
#include "lindhard_core.h"
#include "results.h"

namespace tbplas::diag {
// External symbols
using tbplas::base::cart2frac;
using tbplas::base::complex_t;
using tbplas::base::EPSILON0;
using tbplas::base::exp_it;
using tbplas::base::frac2cart;
using tbplas::base::gen_kmesh;
using tbplas::base::get_lattice_volume;
using tbplas::base::KB_EVK;
using tbplas::base::NM;
using tbplas::base::TWOPI;
using tbplas::base::UIMAG;

/**
 * @brief Response functions calculator.
 *
 * @tparam model_t datatype of model assigned to this solver
 */
template <typename model_t>
class Lindhard : public DiagSolver<model_t> {
private:
    /**
     * @brief Check for illegal parameters
     *
     */
    void check_config() const
    {
        const DiagConfig& config = DiagSolver<model_t>::config;
        assert((config.dimension == 2) || (config.dimension == 3));
        if (config.dimension == 2) {
            assert(std::get<2>(config.k_grid_size) == 1);
        }
        assert(std::get<0>(config.k_grid_size) > 0);
        assert(std::get<1>(config.k_grid_size) > 0);
        assert(std::get<2>(config.k_grid_size) > 0);
        assert((config.num_spin == 1) || (config.num_spin == 2));
        assert((config.ac_cond_component >= 1) && (config.ac_cond_component <= 9));
    }

    /**
     * @brief Get elementary area/volume in reciprocal space depending on system  dimension.
     *
     * @return  elementary area/volume in 1/nm**2 or 1/nm**3
     */
    [[nodiscard]] double get_dnk() const
    {
        Eigen::Matrix3d g_vec = DiagSolver<model_t>::model_->get_reciprocal_lattice();
        Eigen::Vector3d k_grid = builder::rn2vec(DiagSolver<model_t>::config.k_grid_size);
        for (int i_dim = 0; i_dim < 3; ++i_dim) {
            g_vec.col(i_dim) /= k_grid(i_dim);
        }
        double dnk = 0.0;
        if (DiagSolver<model_t>::config.dimension == 2) {
            dnk = g_vec.row(0).cross(g_vec.row(1)).norm();
        } else {
            dnk = get_lattice_volume(g_vec);
        }
        return dnk;
    }

    /**
     * @brief Get prefactor for dynamic polarizability.
     *
     * @return prefactor for dynamic polarizability
     */
    [[nodiscard]] double get_dyn_pol_factor() const
    {
        double dnk = get_dnk();
        int num_spin = DiagSolver<model_t>::config.num_spin;
        int dimension = DiagSolver<model_t>::config.dimension;
        double factor = num_spin * dnk / pow(base::TWOPI, dimension);
        return factor;
    }

    /**
     * @brief Get Coulomb interaction in momentum space for given q-point.
     *
     * @param[in] q_point CARTESIAN coordinate of q-point in 1/NM
     * @return Coulomb interaction in momentum space in eV
     */
    [[nodiscard]] double get_vq(const Eigen::Vector3d& q_point) const
    {
        double factor = 1.0 / (EPSILON0 * DiagSolver<model_t>::config.back_epsilon);
        double vq = 0.0;
        if (DiagSolver<model_t>::config.dimension == 2) {
            vq = factor * TWOPI / q_point.norm();
        } else {
            vq = factor * 2 * TWOPI / q_point.squaredNorm();
        }
        return vq;
    }

public:
    /**
     * @brief Construct a new Lindhard object
     *
     * @param[in] model model for which calculations will be performed
     */
    explicit Lindhard(const model_t& model)
        : DiagSolver<model_t>(model)
    {
    }

    /**
     * @brief Construct a new Lindhard object with overlap
     *
     * @param[in] model model for which calculations will be performed
     *  @param[in] overlap model which holds the overlap
     */
    explicit Lindhard(const model_t& model, const model_t& overlap)
        : DiagSolver<model_t>(model, overlap)
    {
    }

    /**
     * @brief Construct a new Lindhard object
     *
     * @param[in] model pointer to model
     * @param[in] overlap pointer to overlap
     *
     * @note This constructor is for Python interface.
     */
    explicit Lindhard(const model_t* model, const model_t* overlap = nullptr)
        : DiagSolver<model_t>(model, overlap)
    {
    }

    /**
     * @brief Convert FRACTIONAL coordinates of k-points to CARTESIAN coordinates in 1/unit.
     *
     * @param[in] frac_coord FRACTIONAL coordinates of k-points, dimension: (3, num_kpt)
     * @param[out] cart_coord CARTESIAN coordinates of k-points in 1/unit, dimension: (3, num_kpt)
     * @param[in] unit scaling factor from unit to NANOMETER, e.g. unit=0.1 for ANGSTROM
     */
    void k_frac2cart(const Eigen::Matrix3Xd& frac_coord, Eigen::Matrix3Xd& cart_coord, double unit = NM) const
    {
        Eigen::Matrix3d recip_lat_vec = DiagSolver<model_t>::model_->get_reciprocal_lattice();
        cart_coord = frac2cart(recip_lat_vec, frac_coord) * unit;
    }

    /**
     * @brief Convert CARTESIAN coordinates of k-points in 1/unit to FRACTIONAL coordinates.
     *
     * @param[in] cart_coord CARTESIAN coordinates in 1/unit, dimension: (3, num_kpt)
     * @param[out] frac_coord FRACTIONAL coordinates of k-points, dimension: (3, num_kpt)
     * @param[in] unit scaling factor from unit to NANOMETER, e.g. unit=0.1 for ANGSTROM
     */
    void k_cart2frac(const Eigen::Matrix3Xd& cart_coord, Eigen::Matrix3Xd& frac_coord, double unit = NM) const
    {
        Eigen::Matrix3d recip_lat_vec = DiagSolver<model_t>::model_->get_reciprocal_lattice();
        Eigen::Matrix3Xd cart_coord_nm = cart_coord / unit;
        frac_coord = cart2frac(recip_lat_vec, cart_coord_nm);
    }

    /**
     * @brief Generate dp_core depending on configurations.
     *
     * @param[in] num_orb number of orbitals
     * @param[in] num_states number of states
     * @param[in] num_qpt number of q-points
     * @return uniq_ptr to dp_core
     */
    [[nodiscard]] std::unique_ptr<AbstractDPCore> make_dp_core(
        const size_t& num_orb,
        const size_t& num_states,
        const size_t& num_qpt) const
    {
        const std::regex pattern_gpu(R"(^gpu$)", std::regex_constants::icase);
        int convention = DiagSolver<model_t>::config.convention;
        if (std::regex_match(DiagSolver<model_t>::config.pp_algo, pattern_gpu)) {
#ifdef WITH_CUDA
            DiagSolver<model_t>::print_stdout("Using DPCoreGPU backend.");
            return std::make_unique<DPCoreGPU>(num_orb, num_states, num_qpt, convention);
#else
            DiagSolver<model_t>::print_stderr("ERROR: CUDA support disabled during compilation!");
            std::exit(-1);
#endif
        }
        DiagSolver<model_t>::print_stdout("Using DefaultDPCore backend.");
        return std::make_unique<DefaultDPCore>(num_orb, num_states, num_qpt, convention);
    }

    /**
     * @brief  Generate ac_core depending on configurations.
     *
     * @param[in] num_orb number of orbitals
     * @param[in] num_states number of states
     * @return uniq_ptr to ac_core
     */
    [[nodiscard]] std::unique_ptr<AbstractACCore> make_ac_core(
        const size_t& num_orb,
        const size_t& num_states)
    {
        const std::regex pattern_gpu(R"(^gpu$)", std::regex_constants::icase);
        if (std::regex_match(DiagSolver<model_t>::config.pp_algo, pattern_gpu)) {
#ifdef WITH_CUDA
            DiagSolver<model_t>::print_stdout("Using ACCoreGPU backend.");
            return std::make_unique<ACCoreGPU>(num_orb, num_states);
#else
            DiagSolver<model_t>::print_stderr("ERROR: CUDA support disabled during compilation!");
            std::exit(-1);
#endif
        }
        DiagSolver<model_t>::print_stdout("Using DefaultACCore backend.");
        return std::make_unique<DefaultACCore>(num_orb, num_states);
    }

    /**
     * @brief Calculate dynamic polarizability for arbitrary q-points.
     *
     * @return dynamic polarizability
     */
    [[nodiscard]] DynPolData calc_dyn_pol()
    {
        check_config();

        // Prepare k-points, q-points and orbital positions
        Eigen::Matrix3Xd k_points_frac = gen_kmesh(DiagSolver<model_t>::config.k_grid_size);
        const Eigen::Matrix3Xd& q_points_cart = DiagSolver<model_t>::config.q_points;
        Eigen::Matrix3Xd q_points_frac;
        k_cart2frac(q_points_cart, q_points_frac);

        // Allocate working arrays
        auto backend = DiagSolver<model_t>::make_backend();
        size_t num_kpt = k_points_frac.cols();
        size_t num_qpt = q_points_frac.cols();
        size_t num_orb = backend->get_num_orbitals();
        size_t num_states = backend->get_num_states();
        Eigen::VectorXd omegas;
        DiagSolver<model_t>::config.gen_energies(omegas);
        Eigen::MatrixXcd dyn_pol = Eigen::MatrixXcd::Zero(omegas.size(), num_qpt);
        auto core = make_dp_core(num_orb, num_states, num_qpt);

        // Other variables
        double mu = DiagSolver<model_t>::config.mu;
        double temperature = DiagSolver<model_t>::config.temperature;
        double beta = 1 / (KB_EVK * temperature);
        double delta = DiagSolver<model_t>::config.delta;
        int convention = DiagSolver<model_t>::config.convention;

        // Evaluate phase in convention II
        if (convention != 1) {
            Eigen::VectorXcd coeff_q;
            Eigen::VectorXcd coeff_mq;
            for (ptrdiff_t iq = 0; iq < num_qpt; iq++) {
                DiagSolver<model_t>::model_->build_density_coeff(q_points_frac.col(iq), coeff_q, coeff_mq);
                core->phase.col(iq) = coeff_q;
            }
        }

        // Loop over k-points
        RangeInfo range = DiagSolver<model_t>::dist_range(num_kpt);
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            // Get eigenvalues and eigenstates at kpt
            Eigen::Vector3d kpt = k_points_frac.col(ik);
            backend->calc_eigen_all(kpt, convention, core->eng_k, core->wfn_k);
            core->f_k = 1.0 / (1.0 + exp(beta * (core->eng_k.array() - mu)));

            // Loop over q-points
            for (ptrdiff_t iq = 0; iq < num_qpt; ++iq) {
                // Get eigenvalues and eigenstates at k+q
                Eigen::Vector3d k_plus_q = kpt + q_points_frac.col(iq);
                backend->calc_eigen_all(k_plus_q, convention, core->eng_kq, core->wfn_kq);
                core->f_kq = 1.0 / (1.0 + exp(beta * (core->eng_kq.array() - mu)));

                // Evaluate delta_eng and prod_df
                core->calc_product(iq);
                core->calc_dyn_pol(omegas, delta, iq, dyn_pol);
            }
        }
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceMatrixXcd(dyn_pol);

        // Normalize
        dyn_pol *= get_dyn_pol_factor();

        // Assemble results and return
        DynPolData data(DiagSolver<model_t>::config.prefix);
        data.energies = std::move(omegas);
        data.dyn_pol = std::move(dyn_pol);
        if (DiagSolver<model_t>::is_master()) {
            data.save("dyn_pol");
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        DiagSolver<model_t>::mpi_env_.barrier();
        return data;
    }

    /**
     * @brief Calculate relative dielectric function from dynamic polarizability
     *
     * @param[inout] data data holding dynamic polarizability and dielectric function
     * @return dielectric function
     */
    [[nodiscard]] EpsilonData calc_epsilon(DynPolData& data) const
    {
        size_t num_omega = data.energies.size();
        size_t num_qpt = data.dyn_pol.cols();
        assert(num_qpt > 0);
        EpsilonData eps_data(DiagSolver<model_t>::config.prefix);
        eps_data.energies = data.energies;
        eps_data.epsilon = Eigen::MatrixXcd::Zero(num_omega, num_qpt);
        for (ptrdiff_t iq = 0; iq < num_qpt; ++iq) {
            double vq = get_vq(DiagSolver<model_t>::config.q_points.col(iq));
            eps_data.epsilon.col(iq) = 1.0 - vq * data.dyn_pol.col(iq).array();
        }
        if (DiagSolver<model_t>::is_master()) {
            eps_data.save("epsilon");
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need a barrier here.
        DiagSolver<model_t>::mpi_env_.barrier();
        return eps_data;
    }

    /**
     * @brief Calculate dynamic polarizability for arbitrary q-points.
     *
     * @return dynamic polarizability
     */
    [[nodiscard]] ACCondData calc_ac_cond()
    {
        check_config();

        // Prepare k-points
        Eigen::Matrix3Xd k_points_frac = gen_kmesh(DiagSolver<model_t>::config.k_grid_size);
        Eigen::Matrix3Xd k_points_cart;
        k_frac2cart(k_points_frac, k_points_cart);

        // Allocate working arrays
        auto backend = DiagSolver<model_t>::make_backend();
        size_t num_kpt = k_points_frac.cols();
        size_t num_orb = backend->get_num_orbitals();
        size_t num_states = backend->get_num_states();
        Eigen::VectorXd omegas;
        DiagSolver<model_t>::config.gen_energies(omegas);
        Eigen::VectorXcd ac_cond = Eigen::VectorXcd::Zero(omegas.size());
        auto core = make_ac_core(num_orb, num_states);
        DiagSolver<model_t>::model_->build_ham_dr_coo(core->hop_i, core->hop_j, core->hop_v, core->dr_nm);

        // Other variables
        double mu = DiagSolver<model_t>::config.mu;
        double temperature = DiagSolver<model_t>::config.temperature;
        double beta = 1 / (KB_EVK * temperature);
        double delta = DiagSolver<model_t>::config.delta;
        constexpr int convention = 1; // AC cond. should always use convention 1.
        int comp_alpha = 0;
        int comp_beta = 0;
        DiagSolver<model_t>::config.map_ac_cond_component(comp_alpha, comp_beta);

        // Loop over k-points
        RangeInfo range = DiagSolver<model_t>::dist_range(num_kpt);
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            // Get eigenvalues and eigenstates at kpt
            Eigen::Vector3d kpt = k_points_frac.col(ik);
            backend->calc_eigen_all(kpt, convention, core->eng_k, core->wfn_k);
            core->f_k = 1.0 / (1.0 + exp(beta * (core->eng_k.array() - mu)));

            // Evaluate intermediate results
            core->calc_product(k_points_cart, ik, comp_alpha, comp_beta);

            // Evaluate ac conductivity
            core->calc_ac_cond(omegas, delta, ac_cond);
        }
        DiagSolver<model_t>::mpi_env_.InplaceAllreduceVectorXcd(ac_cond);

        // Multiply prefactor
        int dimension = DiagSolver<model_t>::config.dimension;
        int num_spin = DiagSolver<model_t>::config.num_spin;
        complex_t prefactor;
        if (dimension == 2) {
            double area = DiagSolver<model_t>::model_->get_lattice_area("c");
            prefactor = static_cast<double>(num_spin) * UIMAG / (area * num_kpt);
        } else {
            double volume = DiagSolver<model_t>::model_->get_lattice_volume();
            prefactor = static_cast<double>(num_spin) * UIMAG / (volume * num_kpt);
        }
        ac_cond *= prefactor;

        // Assemble results and return
        ACCondData data(DiagSolver<model_t>::config.prefix);
        data.energies = std::move(omegas);
        data.ac_cond = std::move(ac_cond);
        if (DiagSolver<model_t>::is_master()) {
            data.save("ac_cond");
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        DiagSolver<model_t>::mpi_env_.barrier();
        return data;
    }

    /**
     * @brief Calculate dielectric function from AC conductivity for q = 0.
     *
     * @param[in] ac_cond AC conductivity
     * @return relative dielectric function
     */
    [[nodiscard]] EpsilonData calc_epsilon_q0(const ACCondData& ac_cond)
    {
        assert(DiagSolver<model_t>::config.dimension == 3);
        double prefactor = 2 * tbplas::base::TWOPI / (EPSILON0 * DiagSolver<model_t>::config.back_epsilon);
        EpsilonData data(DiagSolver<model_t>::config.prefix);
        data.energies = ac_cond.energies;
        data.epsilon = 1.0 + UIMAG * prefactor * ac_cond.ac_cond.array() / ac_cond.energies.array();
        if (DiagSolver<model_t>::is_master()) {
            data.save("epsilon_q0");
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        DiagSolver<model_t>::mpi_env_.barrier();
        return data;
    }
};
} // namespace tbplas::diag
#endif
