/**
 * @brief Classes for storing calculation parameters of TBPM.
 *
 * @author: Yunhai Li
 */

#ifndef TBPLAS_TBPM_CONFIG_H
#define TBPLAS_TBPM_CONFIG_H

#include <set>

#include <eigen3/Eigen/Dense>

#include "base/datatypes.h"

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

/**
 * @brief Class for for storing user-defined calculation parameters of TBPM.
 *
 */
class TBPMConfig {
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::VectorXcd> python_wft_wf0_map_ = Eigen::Map<const Eigen::VectorXcd>(nullptr, 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 Whether to save results to data files, currently affects quasi-eigenstates
    ///   and time-dependent wave functions
    /// @note This flag is for benchmarking purposes, where we need to disable output to avoid
    ///   expensive I/O.
    bool save_data = true;

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

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

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

    /// @brief The spectral width of occupied states.
    /// This is typically defined as the Fermi energy minus the minimum eigenvalue of the Hamiltonian,
    /// extended by a temperature-dependent broadening factor (e.g., +a * k_B * T) to account for states
    /// broadened by temperature. 'a' is a numerical coefficient.
    double occ_width = 0.0; 

    /// @brief The number of electrons.
    int num_electron = 0;

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

    /// @brief Algorithm for TBPM/Haydock/KuboBastin calculations
    /// @note Candidates: cpu/cpu_fast/gpu
    std::string algo = "cpu";

    /// @brief pool size for fast TBPM algorithms
    int pool_size = 4;

    bool opt_pool_size = false;

    int num_internal_runs = 1000; 
    int num_external_runs = 10;

    /// @brief Rescaling factor of Hamiltonian in eV
    /// @note If the factor is negative, then it will be discarded and a reasonable
    /// default value will be used.
    double rescale = -1.0;

    /// @brief Bessel function precision cut-off.
    double bessel_precision = 1.0e-14;

    /// @brief Number of random initial wave functions.
    int num_random_samples = 1;

    /// @brief Number of time steps.
    int num_time_steps = 1024;

    /// @brief Seed for random wave function generation.
    int seed = 1337;

    /// @brief Check the wave function for divergence after this number of steps.
    int wfn_check_time_steps = 128;

    /// @brief Threshold for checking divergence of wave function.
    double wfn_check_threshold = 1.0e-9;

    /// @brief Precision cut-off of Fermi-Dirac distribution.
    double fermi_cheb_precision = 1.0e-12;

    /// @brief Number of Fermi-Dirac distribution FFT steps, must be power of two.
    int num_fermi_fft_steps = 32768;

    /// @brief Whether to calculate LDOS in TBPM DOS calculation
    /// @note For the Haydock method, we always calculate LDOS.
    bool ldos = false;

    /// @brief Site indices for LDOS calculation.
    /// @note There is no limit for the number of sites for TBPM.
    /// However, the Haydock recursion method accepts only one site.
    std::set<size_t> ldos_orbital_indices = { 0 };

    /// @brief infinitesimal in eV
    double ldos_delta = 0.01;

    /// @brief Recursion depth of Haydock method.
    int ldos_recursion_depth = 2000;

    /// @brief FRACTIONAL coordinates of k-points in COLUMN-WISE order
    /// dimension: (3, num_kpt)
    Eigen::Matrix3Xd k_points = Eigen::Matrix3Xd::Zero(3, 1);

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

    /// @brief Scaling factor for Coulomb potential
    double dyn_pol_coulomb_factor = 1.0;

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

    /// @brief Lists of minimum and maximum of energy window for DC conductivity in eV
    /// @note Why not VectorXd? Because the length must be divided by 2!
    std::vector<std::tuple<double, double>> dc_cond_energy_limits { { -0.5, 0.5 } };

    /// @brief Number of time steps for calculating dc correlation.
    int dc_cond_num_time_steps = 1024;

    /// @brief List of chemical potentials to calculate Hall conductivity in eV.
    Eigen::VectorXd dckb_energies = Eigen::VectorXd::LinSpaced(40, -0.2, 0.2);

    /// @brief Number of kernels in Kernel Polynomial Method(KPM).
    /// @note If the model is not large enough, a large dckb_num_kernel will lead to erroneous
    /// results, which have large absolute values.
    int dckb_num_kernel = 2048;

    /// @brief Number of integral steps.
    int dckb_num_integ_steps = 2048;

    /// @brief Component of Hall conductivity.
    /// @note 1 gives XX, 2 gives XY.
    int dckb_component = 1;

    /// @brief Unit of Hall conductivity, should be hbar or h
    std::string dckb_unit = "hbar";

    /// @brief List of energies of quasi-eigenstates in eV.
    Eigen::VectorXd qe_energies { { -0.1 }, { 0.0 }, { 0.1 } };

    /// @brief Fast algorithm used for calculating quasi-eigenstates.
    /// @note Candidates: "cheby_poly" (for few energies) 
    ///       "block_time" (for many energies, large systems)
    std::string qe_method = "cheby_poly";

    /// @brief Fast algorithm used for calculating DC conductivity.
    /// @note Candidates: "cheby_poly" (for few energies) 
    ///       "block_time" (for many energies, large systems)
    std::string dc_method = "cheby_poly";

    /// @brief steps on which time the time-dependent wave function will be saved
    std::set<size_t> wft_time_save = { 0 };

    /// @brief scale factor for the time step with respect to the default value
    double wft_dt_scale = 1.0;

    /// @brief Initial state of time-dependent wave function
    Eigen::VectorXcd wft_wf0;

    /**
     * @brief Set the wft_wf0 attribute
     *
     * @param[in] arg_wft_wf0 pointer to data of wft_wf0
     * @param[in] arg_num_orb size of wft_wf0
     *
     * @note Applied to Python mode only.
     */
    void set_wft_wf0(const complex_t* arg_wft_wf0, const size_t& arg_num_orb);

    /**
     * @brief Get map to the wft_wf0 attribute
     *
     * @return map to the wft_wf0 attribute
     */
    [[nodiscard]] Eigen::Map<const Eigen::VectorXcd> get_wft_wf0() const;
};

/**
 * @brief Class for holding various parameters between TBPMSolver and TBPM functions.
 *
 * @note We set all default values of member variables to 0.0, since it will trigger
 * Inf/NaN or strange results if the developer forgets to set them.
 *
 * @note Member variables of this class may change significantly among different
 * calculations. DO NOT make it a member of TBPMSolver class, because it is bug-prone
 * due to possible incomplete updating of the variables. Instead, generate it from scratch
 * when you need an object of this class.
 */
class TBPMMeta {
public:
    /// @brief copy of 'rescale_' attribute of TBPMSolver
    /// @note This is the ACTUAL rescaling factor used in TBPM calculations. It does not
    /// have to equal to the TBPMConfig::rescale, as the latter will be discarded if
    /// it is negative.
    double rescale = 0.0;

    /// @brief output of TBPMSolver::get_energy_range
    double energy_range = 0.0;

    /// @brief time step in hbar/eV
    double time_step = 0.0;

    /// @brief rescaled 1/kBT and chemical potential
    double beta_re = 0.0;
    double mu_re = 0.0;

    /// @brief List of energies of quasi-eigenstates in eV.
    /// @note This is the ACTUAL quasi-eigenstates energies used in TBPM calculations.
    /// For pure quasi-eigenstates calculation, it is a copy of TBPMConfig::qe_energies.
    /// For DC conductivity calculation, it is generated by TBPMSolver::calc_corr_dc_cond().
    Eigen::VectorXd qe_energies;

    /// @brief rank of this MPI process
    int mpi_rank = 0;

    /// @brief total number of MPI processes
    int mpi_size = 1;

    /// @brief whether this process is the master process (mpi_rank == 0)
    bool is_master = false;

    /// @brief number of samples assigned to this MPI process
    int num_samples_dist = 0;
};
} // namespace tbplas::tbpm

#endif
