/**
 * @brief Classes for storing calculation parameters of diagonalization.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_DIAG_CONFIG_H
#define TBPLAS_DIAG_CONFIG_H

#include <map>
#include <set>

#include <eigen3/Eigen/Dense>

namespace tbplas::diag {

/**
 * @brief Class for holding parameters of diagonalization-based parameters
 *
 */
class DiagConfig {
private:
    /**
     * @brief Maps to arrays created by python frontend
     *
     * @note Assigning Eigen::Map A to map B is not allowed if B already binds to some address
     * including nullptr. Otherwise, segfault may be raised. You have to destroy the original
     * map B completely and create a new one from scratch. Here are some solutions:
     *
     * 1. C++ placement new syntax:
     *   http://www.eigen.tuxfamily.org/dox-devel/group__TutorialMapClass.html#TutorialMapPlacementNew
     * 2. shared_ptr on Eigen::Map
     *   Remember to set shared_ptr to null before creating the new map. May be slow since the map is
     *   now in the heap.
     * 3. Define your own map class which is memory-safe, and convert to Eigen::Map when needed.
     */
    Eigen::Map<const Eigen::Matrix3Xd> python_kpt_map_ = Eigen::Map<const Eigen::Matrix3Xd>(nullptr, 3, 0);

public:
    /// @brief Whether we are in Python mode (affects access of some arrays)
    bool python_mode = false;

    /// @brief Prefix prepended to output files
    std::string prefix = "sample";

    /// @brief FRACTIONAL coordinates of k-points in COLUMN-WISE order
    /// dimension: (3, num_kpt)
    /// @note Effective in bands, dos, spin texture calculations.
    Eigen::Matrix3Xd k_points = Eigen::Matrix3Xd::Zero(3, 1);

    /// @brief dimension of mesh grid in 1st Brillouin zone
    /// @note
    /// 1. We set the default value to -1 to indicate that it has not been set by the user.
    /// 2. Effective in Berry and Lindhard calculations.
    std::tuple<int, int, int> k_grid_size = { -1, -1, -1 };

    /// @brief convention of Hamiltonian, should be either 1 or 2
    int convention = 1;

    /// @brief dimension of the model, should be 2 or 3
    int dimension = 3;

    /// @brief number of occupied states
    /// @note We set the default value to -1 to indicate that it has not been set by the user.
    int num_occ = -1;

    /// @brief spin degeneracy, should be 1 or 2
    int num_spin = 1;

    /// @brief whether the orbitals are sorted in spin-major order
    bool spin_major = true;

    /// @brief lower bound of the energy grid in eV
    double e_min = -5.0;

    /// @brief upper bound of the energy grid in eV
    double e_max = 5.0;

    /// @brief resolution of energy grid in eV
    double e_step = 0.05;

    /// @brief broadening parameter in eV
    double sigma = 0.05;

    /// @brief basis function to approximate the Delta function
    std::string smear_func = "Gaussian";

    /// @brief Whether to calculate LDOS in DOS calculation
    bool ldos = false;

    /// @brief orbital indices for evaluating local dos, starting from 0
    std::set<size_t> ldos_orbital_indices = { 0 };

    // Z2 specific settings

    /// @brief FRACTIONAL coordinates of the loop along a-axis
    Eigen::VectorXd ka_array = Eigen::VectorXd::LinSpaced(200, -0.5, 0.5);

    /// @brief FRACTIONAL coordinates of the loop along b-axis
    Eigen::VectorXd kb_array = Eigen::VectorXd::LinSpaced(100, 0.0, 0.5);

    /// @brief FRACTIONAL coordinate of the loop along c-axis
    double kc = 0.0;

    // Berry specific settings

    /// @brief Size of extend Brillouin zone with respect to 1st Brillouin zone
    std::tuple<int, int, int> bz_size = { 1, 1, 1 };

    /// @brief Fractional coordinate of shift for k-points in extended Brillouin zone
    Eigen::Vector3d bz_shift = Eigen::Vector3d::Zero();

    /// @brief whether to use analytical derivation of Hamiltonian
    bool ham_deriv_analytical = false;

    /// @brief displacement of k-points in 1/NM for numerical derivation
    double ham_deriv_delta = 1.0e-6;

    // Response functions specific settings

    /// @brief infinitesimal in eV
    double delta = 0.005;

    /// @brief chemical potential in eV
    double mu = 0.0;

    /// @brief temperature in Kelvin
    double temperature = 300.0;

    /// @brief relative background dielectric constant
    double back_epsilon = 1.0;

    /// @brief CARTESIAN coordinates of q-points in 1/NM in COLUMN-WISE order
    /// dimension: (3, num_qpt)
    Eigen::Matrix3Xd q_points { { 1.0 }, { 0.0 }, { 0.0 } };

    /// @brief component of ac conductivity to evaluate
    /// @note For consistency with TBPMConfig::dckb_component:
    /// 1-XX, 2-XY, 3-XZ, 4-YX, 5-YY, 6-YZ, 7-ZX, 8-ZY, 9-ZZ
    /// DO NOT CHANGE THE ORDER. Otherwise map_ac_cond_component will not work properly.
    int ac_cond_component = 1;

    /// @brief Algorithm settings

    /// @brief Algorithm for diagonalization
    /// @note Candidates:
    /// eigen: CPU solvers from eigen
    /// lapacke: CPU solvers from LAPACKE
    /// feast_dense: CPU solvers from feast for dense matrix
    /// feast_sparse: CPU solvers from feast for sparse matrix
    std::string diag_algo = "eigen";

    /// @brief Algorithm for post-processing
    /// @note Candidates: cpu, gpu
    std::string pp_algo = "cpu";

    /// @brief number of states
    /// @note If set to 0, then it will be discarded and a reasonable default value
    /// will be used.
    size_t feast_num_states = 0;

    /// @brief lower bound of the feast searching range in eV
    double feast_e_min = -5.0;

    /// @brief upper bound of the feast searching range in eV
    double feast_e_max = 5.0;

    /// @brief criterion for convergence
    double feast_eps = 1.0e-5;

    /// @brief ratio of subspace size to number of states (m0/num_states)
    double feast_m0_ratio = 1.0;

    /// @brief whether to reuse results of last run as initial guess
    /// @note
    /// 1. Reusing results should not be enabled from scratch, otherwise there will be
    ///    convergence problems. So we cannot set fpm[4] (fortran fpm(5)) directly.
    /// 2. This parameter sets feast_fpm[4] (fortran fpm(5)) to 1 after the first
    ///    successful diagonalization, which is typically safe.
    bool feast_reuse_results = false;

    /// @brief fpm parameters
    /// @note
    /// 1. 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], and fpm[0]
    ///    cannot be negative. See the following link for more details:
    ///    https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2023-0/extended-eigensolver-input-parameters.html
    ///
    /// 2. Anaconda has MKL shipped with it. So when using Python from Anaconda, it is likely
    ///    that MKL feast will shield the original version. Check that carefully.
    ///
    /// 3. We only disable the verbose output by default (fpm[0] = 0), since carelessly
    ///    changing the parameters may corrupt feast.
    ///
    /// 4. We use C-style fpm with indices starting from 0.
    std::map<int, int> feast_fpm = { { 0, 0 } };

    /**
     * @brief Generate energy grid from e_min, e_max and e_step
     *
     * @param[out] energy energy grid built from e_min, e_max and e_step.
     */
    void gen_energies(Eigen::VectorXd& energy) const;

    /**
     * @brief Remap ac_cond_component to X(0)/Y(1)/Z(2)
     *
     * @param[out] alpha index of component alpha
     * @param[out] beta index of component beta
     *
     * @note DO NOT CHANGE THE ORDER in ac_cond_component above. Otherwise the results will be wrong.
     */
    void map_ac_cond_component(int& alpha, int& beta) const;

    /**
     * @brief Set the k_points attribute
     *
     * @param[in] arg_k_points pointer to external coordinates of k-points in COLUMN-WISE order
     * @param[in] arg_num_kpt number of k-points
     *
     * @note Applied to Python mode only.
     */
    void set_k_points(const double* arg_k_points, const size_t& arg_num_kpt);

    /**
     * @brief Get the map to k_points attribute
     *
     * @return map to k_points
     */
    [[nodiscard]] Eigen::Map<const Eigen::Matrix3Xd> get_k_points() const;
};

} // namespace tbplas::diag

#endif