/**
 * @brief Core classes for Lindhard, for runtime switching of backends.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */
#ifndef TBPLAS_DIAG_LINDHARD_CORE_H
#define TBPLAS_DIAG_LINDHARD_CORE_H

#include <eigen3/Eigen/Dense>

#include "base/datatypes.h"

#ifdef WITH_CUDA
#include "cuda/cudaConfig.h"
#endif

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

/**
 * @brief Core part of dynamic polarizability evaluation
 *
 */
class AbstractDPCore {
public:
    // energies at k and k+q in eV, dimension: (num_states)
    Eigen::VectorXd eng_k, eng_kq;

    // Wave functions at k and k+q in COLUMN-WISE order, dimension: (num_orb, num_states)
    Eigen::MatrixXcd wfn_k, wfn_kq;

    // Occupation numbers at k and k+q, dimension: (num_states)
    Eigen::VectorXd f_k, f_kq;

    // exp_iz(q_dot_orb_pos) in COLUMN-WISE order, dimension: (num_orb, num_qpt)
    Eigen::MatrixXcd phase;

    // energy differences in eV, dimension: (num_states, num_states)
    Eigen::MatrixXd delta_eng;

    // |<k+q|phase|k>|^2 * (f_k - f_kq), dimension: (num_states, num_states)
    Eigen::MatrixXcd product;

    // convention of Hamiltonian
    int convention = 1;

    /**
     * @brief Construct a new AbstractDPCore object
     *
     * @param[in] num_orb number of orbitals of the model
     * @param[in] num_states number of states to calculate
     * @param[in] num_qpt number of q-points
     * @param[in] conv convention of hamiltonian
     */
    explicit AbstractDPCore(
        const size_t& num_orb,
        const size_t& num_states,
        const size_t& num_qpt,
        const int& conv);

    /**
     * @brief Calculate energy differences and product matrix
     *
     * @param[in] iq index of q-point
     */
    virtual void calc_product(const int& iq) = 0;

    /**
     * @brief Calculate dynamic polarizability from energy differences and product matrix
     *
     * @param[in] omegas energies in ev, dimension: (num_eng)
     * @param[in] delta infinitesimal in eV
     * @param[in] iq index of q-point
     * @param[inout] dyn_pol dynamic polarizability, dimension: (num_eng, num_qpt)
     */
    virtual void calc_dyn_pol(
        const Eigen::VectorXd& omegas,
        const double& delta,
        const int& iq,
        Eigen::MatrixXcd& dyn_pol) const
        = 0;
};

/**
 * @brief Core part of AC conductivity evaluation
 *
 */
class AbstractACCore {
public:
    // energies at k in eV, dimension: (num_states)
    Eigen::VectorXd eng_k;

    // Wave functions at k in COLUMN-WISE order, dimension: (num_orb, num_states)
    Eigen::MatrixXcd wfn_k;

    // Occupation numbers at k and k+q, dimension: (num_states)
    Eigen::VectorXd f_k;

    // energy differences in eV, dimension: (num_states, num_states)
    Eigen::MatrixXd delta_eng;

    // <va><vb>(f_k-f_kq), dimension: (num_states, num_states)
    Eigen::MatrixXcd product;

    /// @brief index i in orbital pairs, i.e, orbital in r0
    std::vector<size_t> hop_i;

    /// @brief index j in orbital pairs, i.e, orbital in rn
    std::vector<size_t> hop_j;

    /// @brief hopping energy in eV
    std::vector<complex_t> hop_v;

    /// @brief hopping distances in nm in COLUMN-WISE order
    Eigen::Matrix3Xd dr_nm;

    /// @brief velocity matrices
    Eigen::MatrixXcd v_alpha;
    Eigen::MatrixXcd v_beta;

    /**
     * @brief Construct a new AbstractACCore object
     *
     * @param[in] num_orb number of orbitals of the model
     * @param[in] num_states number of states to calculate
     */
    explicit AbstractACCore(
        const size_t& num_orb,
        const size_t& num_states);

    /**
     * @brief Calculate energy differences and product matrix
     *
     * @param[in] k_points_cart Cartesian coordinate of k-point in 1/NM in COLUMN-WISE order, dimension: (3, num_kpt)
     * @param[in] ik index of current k-point in k_points_cart
     * @param[in] comp_alpha index of component of v_alpha
     * @param[in] comp_beta index of component of v_beta
     *
     * @note Passing k_points_cart.col(ik) as argument is slow.
     */
    virtual void calc_product(
        const Eigen::Matrix3Xd& k_points_cart,
        const int& ik,
        const int& comp_alpha,
        const int& comp_beta)
        = 0;

    /**
     * @brief Calculate ac conductivity from energy differences and product matrix
     *
     * @param[in] omegas energies in ev, dimension: (num_eng)
     * @param[in] delta infinitesimal in eV
     * @param[inout] ac_cond AC conductivity, dimension (num_eng)
     */
    virtual void calc_ac_cond(
        const Eigen::VectorXd& omegas,
        const double& delta,
        Eigen::VectorXcd& ac_cond) const
        = 0;
};

class DefaultDPCore final : public AbstractDPCore {
public:
    explicit DefaultDPCore(
        const size_t& num_orb,
        const size_t& num_states,
        const size_t& num_qpt,
        const int& conv)
        : AbstractDPCore(num_orb, num_states, num_qpt, conv) {};

    void calc_product(const int& iq) final;

    void calc_dyn_pol(
        const Eigen::VectorXd& omegas,
        const double& delta,
        const int& iq,
        Eigen::MatrixXcd& dyn_pol) const final;
};

class DefaultACCore final : public AbstractACCore {
public:
    explicit DefaultACCore(
        const size_t& num_orb,
        const size_t& num_states)
        : AbstractACCore(num_orb, num_states) {};

    void calc_product(
        const Eigen::Matrix3Xd& k_points_cart,
        const int& ik,
        const int& comp_alpha,
        const int& comp_beta) final;

    void calc_ac_cond(
        const Eigen::VectorXd& omegas,
        const double& delta,
        Eigen::VectorXcd& ac_cond) const final;
};

#ifdef WITH_CUDA
/// @todo Implement the methods
class DPCoreGPU final : public AbstractDPCore {
private:
    CudaConfig cuda_;

public:
    explicit DPCoreGPU(
        const size_t& num_orb,
        const size_t& num_states,
        const size_t& num_qpt,
        const int& conv)
        : AbstractDPCore(num_orb, num_states, num_qpt, conv) {};

    void calc_product(const int& iq) final;

    void calc_dyn_pol(
        const Eigen::VectorXd& omegas,
        const double& delta,
        const int& iq,
        Eigen::MatrixXcd& dyn_pol) const final;
};

class ACCoreGPU final : public AbstractACCore {
private:
    CudaConfig cuda_;

public:
    explicit ACCoreGPU(
        const size_t& num_orb,
        const size_t& num_states)
        : AbstractACCore(num_orb, num_states) {};

    void calc_product(
        const Eigen::Matrix3Xd& k_points_cart,
        const int& ik,
        const int& comp_alpha,
        const int& comp_beta) final;

    void calc_ac_cond(
        const Eigen::VectorXd& omegas,
        const double& delta,
        Eigen::VectorXcd& ac_cond) const final;
};
#endif // WITH_CUDA
} // namespace tbplas::diag
#endif