/**
 * @brief Base class for analytical models.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */
#ifndef TBPLAS_BUILDER_ANALYTICAL_H
#define TBPLAS_BUILDER_ANALYTICAL_H

#include <iostream>
#include <vector>

#include "base.h"
#include "base/consts.h"

namespace tbplas::builder {

/**
 * @brief Base class for analytical models.
 */
class AnalyticalModel {
public:
    /**
     * @brief Construct a new AnalyticalModel object
     *
     * @param[in] num_orb number of orbitals
     * @param[in] num_hop number of hopping terms
     * @param[in] lattice Cartesian coordinates of lattice vectors in NM
     * @note Each row of lattice corresponds to a lattice vector.
     */
    explicit AnalyticalModel(
        orbital_index_t num_orb,
        orbital_index_t num_hop,
        const Eigen::Matrix3d& lattice);

    /**
     * @brief Get the number of orbitals
     *
     * @return number of orbitals
     */
    [[nodiscard]] orbital_index_t get_num_orbitals() const;

    /**
     * @brief Get the number of hopping terms
     *
     * @return  number of hopping terms
     */
    [[nodiscard]] orbital_index_t get_num_hopping() const;

    /**
     * @brief Get lattice vectors
     *
     * @return Cartesian coordinates of lattice vectors in NM
     *
     * @note Each ROW of lattice is a lattice vector.
     */
    [[nodiscard]] const Eigen::Matrix3d& get_lattice() const;

    /**
     * @brief Get lattice area
     *
     * @param[in] direction direction along which the area is evaluated, should be in ("a", "b", "c")
     * @return lattice area in nm^2
     */
    [[nodiscard]] double get_lattice_area(const std::string& direction = "c") const;

    /**
     * @brief Get the lattice volume attribute
     *
     * @return lattice volume in nm^3
     */
    [[nodiscard]] double get_lattice_volume() const;

    /**
     * @brief Get reciprocal lattice vectors
     *
     * @return Cartesian coordinates of reciprocal lattice vectors in 1/NM.
     */
    [[nodiscard]] const Eigen::Matrix3d& get_reciprocal_lattice() const;

    /**
     * @brief Extract hopping data in COO format for Lindhard AC conductivity calculation
     *
     * @param[out] hop_i index i in orbital pairs, i.e, orbital in r0
     * @param[out] hop_j index j in orbital pairs, i.e, orbital in rn
     * @param[out] hop_v hopping energy in eV
     * @param[out] dr_nm hopping distances in nm in COLUMN-WISE order
     *
     * @note Conjugate terms are not included. The caller is responsible for restoring
     * them.
     */
    void build_ham_dr_coo(
        std::vector<orbital_index_t>& hop_i,
        std::vector<orbital_index_t>& hop_j,
        std::vector<complex_t>& hop_v,
        Eigen::Matrix3Xd& dr_nm) const;

    /**
     * @brief Build the attributes of sparse Hamiltonian and current operator
     * in CSR format.
     *
     * @param[out] indptr the 'indptr' attribute (outerIndex of eigen)
     * @param[out] indices the 'indices' attribute (innerIndex of eigen)
     * @param[out] ham non-zero elements of Hamiltonian in eV
     * @param[out] curr_coeff_x coefficients of current operator along along x direction
     * @param[out] curr_coeff_y coefficients of current operator along along x direction
     *
     * @note See eqn. 54 of CPC paper for the formula of current operator.
     */
    void build_ham_curr_csr(
        std::vector<orbital_index_t>& indptr,
        std::vector<orbital_index_t>& indices,
        Eigen::VectorXcd& ham,
        Eigen::VectorXcd& curr_coeff_x,
        Eigen::VectorXcd& curr_coeff_y) const;

    /**
     * @brief Build the attributes of sparse Hamiltonian for given k-point
     * in CSR format.
     *
     * @tparam csr_idx_t datatype of CSR matrix indices
     * @param[out] indptr the 'indptr' attribute (outerIndex of eigen)
     * @param[out] indices the 'indices' attribute (innerIndex of eigen)
     * @param[out] ham non-zero elements of Hamiltonian in eV
     * @param[in] k_point fractional coordinate of k-point
     * @param[in] convention convention of Hamiltonian
     *
     * @note
     * 1. In most cases we need only the Hamiltonian at Gamma point. So we make
     *    it that if convention != 1 or 2, k_point will be discarded and Gamma point
     *    is assumed.
     * 2. The FEAST library requires the indices to be int32_t, while orbital_index_t
     *    defaults to uint64_t. So we make it a template argument.
     */
    template <typename csr_idx_t>
    void build_ham_csr(
        std::vector<csr_idx_t>& indptr,
        std::vector<csr_idx_t>& indices,
        Eigen::VectorXcd& ham,
        const Eigen::Vector3d& k_point = Eigen::Vector3d::Zero(3),
        int convention = -1) const
    {
        std::cerr << "ERROR: AnalyticalModel::build_ham_csr not implemented!" << std::endl;
        std::exit(-1);
    }

    /**
     * @brief Set up dense Hamiltonian matrix for given k-point for diagonalization
     * based calculations.
     *
     * @param[in] k_point fractional coordinate of k-point
     * @param[out] ham incoming Hamiltonian matrix
     * @param[in] convention convention of Hamiltonian
     */
    void build_ham_dense(
        const Eigen::Vector3d& k_point,
        Eigen::MatrixXcd& ham,
        int convention = 1) const;

    /**
     * @brief Build density operator coefficients for given q-point for Lindhard
     * dynamic polarization calculation
     *
     * @param[in] q_point FRACTIONAL coordinate of q-point
     * @param[out] coeff_q coefficients for q point
     * @param[out] coeff_mq coefficients for -q point
     *
     * @note The formula is eqn. 78 of CPC paper.
     */
    void build_density_coeff(
        const Eigen::Vector3d& q_point,
        Eigen::VectorXcd& coeff_q,
        Eigen::VectorXcd& coeff_mq) const;

    /**
     * @brief Set up derivatives of dense Hamiltonian matrix for given k-point.
     * @param[in] k_point fractional coordinate of k-point
     * @param[out] dh_dka dH/dka, not dH/dkx!
     * @param[out] dh_dkb dH/dkb, not dH/dky!
     * @param[out] dh_dkc dH/dkc, not dH/dkz!
     * @param[in] convention convention of Hamiltonian
     */
    void build_ham_der_dense(
        const Eigen::Vector3d& k_point,
        Eigen::MatrixXcd& dh_dka,
        Eigen::MatrixXcd& dh_dkb,
        Eigen::MatrixXcd& dh_dkc,
        int convention = 1) const;

protected:
    /// @brief number or orbitals
    orbital_index_t num_orb_ = 0;

    /// @brief number of hopping terms
    orbital_index_t num_hop_ = 0;

    /// @brief Cartesian coordinates of lattice vectors in NM
    /// @note Each ROW corresponds to one lattice vector.
    Eigen::Matrix3d lattice_ = Eigen::MatrixXd::Identity(3, 3);

    /// @brief Cartesian coordinates of reciprocal lattice vectors in 1/NM
    /// @note Each ROW corresponds to one lattice vector.
    Eigen::Matrix3d reciprocal_lattice_ = tbplas::base::TWOPI * Eigen::MatrixXd::Identity(3, 3);
};
} // namespace tbplas::builder
#endif // TBPLAS_BUILDER_ANALYTICAL_H
