/**
 * @brief Functions and classes for solvers based on TBPM
 *
 * @author: Zewen Wu, Miao Zhang, Yunhai Li
 */

#ifndef TBPLAS_TBPM_SOLVER_H
#define TBPLAS_TBPM_SOLVER_H

#include <algorithm>
#include <cassert>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <memory>
#include <regex>
#include <vector>

#include <eigen3/Eigen/Dense>

#include "base/consts.h"
#include "base/io_text.h"
#include "base/kpoints.h"
#include "base/lattice.h"
#include "base/parallel.h"
#include "config.h"
#include "csr.h"
#include "haydock.h"
#include "kubo_bastin.h"
#include "tbpm.h"
#include "tbpm/backend.h"
#include "tbpm_fast.h"

#ifdef WITH_HDF5
#include "base/io_base.h"
#include "base/io_hdf5.h"
#else
#include "base/io_text.h"
#endif

#ifdef WITH_CUDA
#include "haydock_gpu.h"
#include "kubo_bastin_gpu.h"
#include "tbpm_gpu.h"
#endif

namespace tbplas::tbpm {
// Aliases
using tbplas::base::BaseMPIEnv;
using tbplas::base::RangeInfo;
using tbplas::base::TextFile;

#ifdef WITH_HDF5
using tbplas::base::HDFFile;
using tbplas::base::OptBinMode;
#else
using tbplas::base::TextFile;
#endif

/**
 * @brief Class for performing TBPM calculations
 *
 * @tparam model_t datatype of model assigned to this solver
 */
template <typename model_t>
class TBPMSolver {

public:
    /// @brief configurations controlling the calculations
    TBPMConfig config;

private:
    /// @brief pointer to the model for which calculations will be performed
    const model_t* model_ = nullptr;

    /// @brief 'indptr' attribute of sparse Hamiltonian and hopping vectors
    std::vector<size_t> indptr_;

    /// @brief 'indices' attribute of sparse Hamiltonian and hopping vectors
    std::vector<size_t> indices_;

    /// @brief 'values' attribute of sparse Hamiltonian in rescaling factor
    Eigen::VectorXcd ham_;

    /// @brief 'values' attribute of sparse current operator along x-direction
    /// in eV/hbar * eV * nm
    Eigen::VectorXcd curr_coeff_x_;

    /// @brief 'values' attribute of sparse current operator along y-direction
    /// in eV/hbar * eV * nm
    Eigen::VectorXcd curr_coeff_y_;

    /// @brief Rescaling factor of Hamiltonian in eV
    double rescale_ = 1.0;

    /// @brief MPI environment fo parallel calculation
    base::BaseMPIEnv mpi_env_;

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

    /**
     * @brief Construct a new TBPMSolver object
     *
     * @param[in] model pointer to model for which calculations will be performed
     *
     * @note This constructor is to be called by the Python interface, where pointers
     * are more convenient than references.
     */
    explicit TBPMSolver(const model_t* model)
        : model_(model)
    {
    }

    /**
     * @brief Check if this process is the master process in the communicator
     *
     * @return true if this process is the master process
     * @return false otherwise
     */
    [[nodiscard]] bool is_master() const
    {
        return mpi_env_.isMaster();
    }

    /**
     * @brief Print message on master process to stdout.
     *
     * @param[in] msg message text
     */
    void print_stdout(const std::string& msg) const
    {
        if (is_master()) {
            std::cout << msg << std::endl;
        }
    }

    /**
     * @brief Print message on master process to stderr.
     *
     * @param[in] msg message text
     */
    void print_stderr(const std::string& msg) const
    {
        if (is_master()) {
            std::cerr << msg << std::endl;
        }
    }

    /**
     * @brief Calculate LDOS using Haydock recursive method.
     *
     * @note
     * 1. Output files:
     *      ${prefix}_ldos.dat (no HDF5)
     *      ${prefix}_tbpm.h5, with datasets 'ldos.energies', 'ldos.ldos'
     *
     * 2. Data format of dat output:
     *      dimension: (num_eng, 2)
     *      columns: energy, ldos
     */
    void calc_ldos_haydock()
    {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        // Prepare Hamiltonian
        rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
        model_->build_ham_csr(indptr_, indices_, ham_);
        rescale_ham();
        sparse_t h_sparse(indptr_, indices_, ham_);

        // Prepare meta info.
        TBPMMeta meta = get_default_meta();

        // Actual calculations
        Eigen::VectorXd energy;
        Eigen::VectorXd ldos;
        auto backend = make_haydock_backend<sparse_t>();
        backend->ldos_haydock(h_sparse, config, meta, energy, ldos);

        // Average over MPI-processes
        mpi_env_.InplaceAllreduceVectorXd(ldos);
        ldos /= static_cast<double>(meta.mpi_size);

        // Save data
        if (is_master()) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_real("ldos.energies", energy);
            data_file.save_matrix_real("ldos.ldos", ldos);
            data_file.close();
#else
            Eigen::MatrixXd data(energy.size(), 2);
            data.col(0) = energy;
            data.col(1) = ldos;
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_real("ldos", data);
#endif
            save_info();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

    /**
     * @brief Calculate correlation function for DOS and LDOS.
     *
     * @note
     * 1. Output files:
     *      ${prefix}_corr_dos.dat (no HDF5)
     *      ${prefix}_tbpm.h5, with datasets 'corr_dos'
     *
     * 2. Data format of dat output:
     *      dimension: (num_step+1, 2)
     *      columns: corr.real, corr.imag
     */
    void calc_corr_dos()
    {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        // Prepare Hamiltonian
        rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
        model_->build_ham_csr(indptr_, indices_, ham_);
        rescale_ham();
        sparse_t h_sparse(indptr_, indices_, ham_);

        // std::map<std::string, double> total_time;
        // total_time = benchmark_ops(h_sparse,1000);
        // std::cout << "amxsy: "<< total_time["amxsy"] << " axpy: "<< total_time["axpy"]<<" amxsy/axpy: "<<total_time["amxsy"]/total_time["axpy"]<<"\n"; 
        if (config.opt_pool_size){
            run_benchmarks_and_analyze(h_sparse, config.num_internal_runs, config.num_external_runs);
        }

        // Prepare meta info.
        TBPMMeta meta = get_default_meta();

        // Calculate correlation function
        Eigen::VectorXcd corr_dos;
        auto backend = make_tbpm_backend<sparse_t>();
        backend->tbpm_dos(h_sparse, config, meta, corr_dos);

        // Average over MPI-processes
        mpi_env_.InplaceAllreduceVectorXcd(corr_dos);
        corr_dos /= static_cast<double>(meta.mpi_size);

        // Output
        if (is_master()) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_complex("corr_dos", corr_dos);
            data_file.close();
#else
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_complex("corr_dos", corr_dos);
#endif
            save_info();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

    /**
     * @brief Calculate correlation function for AC conductivity.
     *
     * @note
     * 1. Output files:
     *      ${prefix}_corr_ac_cond.dat (no HDF5)
     *      ${prefix}_tbpm.h5, with datasets 'corr_ac_cond'
     *
     * 2. Data format of dat output:
     *      dimension: (4, num_step*2)
     *      columns: (t0.real, t0.imag), (t1.real, t1.imag), ...
     *      rows: xx, yx, xy, yy components respectively
     */
    void calc_corr_ac_cond()
    {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        // Prepare Hamiltonian and current operators
        rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
        model_->build_ham_curr_csr(indptr_, indices_, ham_, curr_coeff_x_, curr_coeff_y_);
        rescale_ham();
        sparse_t h_sparse(indptr_, indices_, ham_);
        sparse_t curr_x(indptr_, indices_, curr_coeff_x_);
        sparse_t curr_y(indptr_, indices_, curr_coeff_y_);

        if (config.opt_pool_size){
            run_benchmarks_and_analyze(h_sparse, config.num_internal_runs, config.num_external_runs);
        }
        
        // std::map<std::string, double> total_time;
        // total_time = benchmark_ops(h_sparse,1000);
        // std::cout << "amxsy: "<< total_time["amxsy"] << " axpy: "<< total_time["axpy"]<<" amxsy/axpy: "<<total_time["amxsy"]/total_time["axpy"]<<"\n"; 


        // Meta info
        // WARNING: Missing factor of 0.5 in time step costs me a whole day debugging.
        TBPMMeta meta = get_default_meta();
        meta.time_step *= 0.5;
	// Due to the Fermi-Dirac operator, there are transitions between occupied states above the Fermi level
        // and unoccupied states below it. As a result, the total spectral width of this physical quantity 
        // is broader than that of the Hamiltonian. Since 1/(np.exp(25)+1) = 1.4* 10^{-11}, we extend the spectral
        // width by an additional 50 * K_B Tto alleviate aliasing issues. 
        //NOTE:  if the temperature is very high,  50 * K_B T may exceed the energy range;
        // in that case, it should be replaced with the energy range.
        //meta.time_step = base::TWOPI / (meta.energy_range + 50 * base::KB_EVK * config.temperature);


        // Calculate correlation function.
        Eigen::MatrixXcd corr;
        auto backend = make_tbpm_backend<sparse_t>();
        backend->tbpm_ac_cond(h_sparse, curr_x, curr_y, config, meta, corr);

        // Average over mpi-processes
        mpi_env_.InplaceAllreduceMatrixXcd(corr);
        corr /= static_cast<double>(meta.mpi_size);

        // Save correlation functions
        if (is_master()) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_complex("corr_ac_cond", corr);
            data_file.close();
#else
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_complex("corr_ac_cond", corr);
#endif
            save_info();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

    /**
     * @brief Calculate correlation function for dynamic polarizability.
     *
     * @note
     * 1. Output files:
     *      ${prefix}_corr_dyn_pol.dat (no HDF5)
     *      ${prefix}_tbpm.h5, with datasets 'corr_dyn_pol'
     *
     * 2. Data format of dat output:
     *      dimension: (num_step, num_qpt)
     *      columns: dp_q0, dp_q1, ...
     */
    void calc_corr_dyn_pol()
    {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        // Prepare Hamiltonian
        rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
        model_->build_ham_csr(indptr_, indices_, ham_);
        rescale_ham();
        sparse_t h_sparse(indptr_, indices_, ham_);

        if (config.opt_pool_size){
            run_benchmarks_and_analyze(h_sparse, config.num_internal_runs, config.num_external_runs);
        }

        // std::map<std::string, double> total_time;
        // total_time = benchmark_ops(h_sparse,1000);
        // std::cout << "amxsy: "<< total_time["amxsy"] << " axpy: "<< total_time["axpy"]<<" amxsy/axpy: "<<total_time["amxsy"]/total_time["axpy"]<<"\n"; 


        // Meta info
        TBPMMeta meta = get_default_meta();
        meta.time_step *= 0.5;
	// Due to the Fermi-Dirac operator, there are transitions between occupied states above the Fermi level
        // and unoccupied states below it. As a result, the total spectral width of this physical quantity 
        // is broader than that of the Hamiltonian. Since 1/(np.exp(25)+1) = 1.4* 10^{-11}, we extend the spectral
        // width by an additional 50 * K_B Tto alleviate aliasing issues. 
        //NOTE:  if the temperature is very high,  50 * K_B T may exceed the energy range;
        // in that case, it should be replaced with the energy range.
        //meta.time_step = base::TWOPI / (meta.energy_range + 50 * base::KB_EVK * config.temperature);

        // Convert q-points to from Cartesian in 1/nm to fractional coordinates
        Eigen::Matrix3Xd q_points = base::cart2frac(model_->get_reciprocal_lattice(), config.dyn_pol_q_points);

        // Loop over q-points to calculate correlation function
        Eigen::VectorXcd coeff_q;
        Eigen::VectorXcd coeff_mq;
        Eigen::VectorXd corr(config.num_time_steps);
        size_t num_qpt = q_points.cols();
        Eigen::MatrixXd data(config.num_time_steps, num_qpt);
        auto backend = make_tbpm_backend<sparse_t>();
        // TODO 如果内存足够的时候，其实可以将多个q的演化移到`tbpm_dyn_pol`函数中。
        // 这样的好处是不同的q都可以利用bra{\psi_2(t)}的结果,
        // 这样可以节省很多计算时间，特别是q的个数比较多的时候。
        // 既然DC可以很多能量态同时计算，dyn_pol的内存消耗应该也是可以接受的
        for (size_t i_q = 0; i_q < num_qpt; ++i_q) {
            if (is_master()) {
                std::cout << "q-point " << i_q + 1 << " of " << num_qpt << std::endl;
            }
            model_->build_density_coeff(q_points.col(i_q), coeff_q, coeff_mq);
            backend->tbpm_dyn_pol(h_sparse, coeff_q, coeff_mq, config, meta, corr);
            data.col(i_q) = corr;
        }

        // Average over mpi processes
        mpi_env_.InplaceAllreduceMatrixXd(data);
        data /= static_cast<double>(meta.mpi_size);

        // Save correlation function
        if (is_master()) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_real("corr_dyn_pol", data);
            data_file.close();
#else
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_real("corr_dyn_pol", data);
#endif
            save_info();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

    /**
     * @brief Calculate quasi-eigenstates for given energies.
     *
     * @note
     * 1. Output files:
     *      ${prefix}_quasi_eigen.dat (no HDF5)
     *      ${prefix}_tbpm.h5 (HDF5), with datasets 'quasi_eigen'
     *
     * 2. Data format of dat output:
     *      dimension: (num_orb, num_qe)
     *      columns: qe0, qe1, ...
     */
    void calc_quasi_eigenstates()
    {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        // Prepare Hamiltonian
        rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
        model_->build_ham_csr(indptr_, indices_, ham_);
        rescale_ham();
        sparse_t h_sparse(indptr_, indices_, ham_);

        if (config.opt_pool_size){
            run_benchmarks_and_analyze(h_sparse, config.num_internal_runs, config.num_external_runs);
        }
        // std::map<std::string, double> total_time;
        // total_time = benchmark_ops(h_sparse,1000);
        // std::cout << "amxsy: "<< total_time["amxsy"] << " axpy: "<< total_time["axpy"]<<" amxsy/axpy: "<<total_time["amxsy"]/total_time["axpy"]<<"\n"; 


        // Meta info
        TBPMMeta meta = get_default_meta();

        // Calculate quasi-eigen
        Eigen::MatrixXd wf2_qe;
        auto backend = make_tbpm_backend<sparse_t>();
        backend->tbpm_quasi_eigenstates(h_sparse, config, meta, wf2_qe);

        // Average over mpi-processes
        mpi_env_.InplaceAllreduceMatrixXd(wf2_qe);
        wf2_qe /= static_cast<double>(meta.mpi_size);

        // Output
        if (is_master() && config.save_data) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_real("quasi_eigen", wf2_qe);
            data_file.close();
#else
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_real("quasi_eigen", wf2_qe);
#endif
            save_info();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

    /**
     * @brief Calculate correlation functions for DC conductivity.
     *
     * @note
     * 1. Output files:
     *      ${prefix}_corr_dc_cond_dos.dat, ${prefix}_corr_dc_cond_xx/yy.dat (no HDF5)
     *      ${prefix}_tbpm.h5 (HDF5), with datasets 'corr_dc_cond.dos', 'corr_dc_cond.xx/yy'
     *
     * 2. Data format of dat output of corr_dc_dos:
     *      dimension: (num_step, 2)
     *      columns: corr.real, corr.imag
     *
     *   Data format of dat output of corr_dc_cond_xx:
     *      dimension: (num_energies, 2*num_step_dc)
     *      columns: (t0.real, t0.imag), (t1.real, t1.imag), ...
     *      rows: e0, e1, e2, ..., eN-1
     */
    void calc_corr_dc_cond()
    {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        // Prepare Hamiltonian and current operators
        rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
        model_->build_ham_curr_csr(indptr_, indices_, ham_, curr_coeff_x_, curr_coeff_y_);
        rescale_ham();
        sparse_t h_sparse(indptr_, indices_, ham_);
        sparse_t curr_x(indptr_, indices_, curr_coeff_x_);
        sparse_t curr_y(indptr_, indices_, curr_coeff_y_);

        // std::map<std::string, double> total_time;
        // total_time = benchmark_ops(h_sparse,1000);
        // std::cout << "amxsy: "<< total_time["amxsy"] << " axpy: "<< total_time["axpy"]<<" amxsy/axpy: "<<total_time["amxsy"]/total_time["axpy"]<<"\n"; 

        // NtSplineInterpolator interpolator;
        // double test_t1 = 100; // Within range
        // double nt1 = interpolator.getInterpolatedNt(test_t1);
        // std::cout << "Interpolated N(t=" << test_t1 << "): " << nt1 << std::endl;
        if (config.opt_pool_size){
            run_benchmarks_and_analyze(h_sparse, config.num_internal_runs, config.num_external_runs);
        }

        // Meta info
        TBPMMeta meta = get_default_meta();

        // Full energy grid of DOS
        int num_steps = config.num_time_steps;
        int num_eng = 2 * num_steps;
        double eng_range = get_energy_range();
        Eigen::VectorXd dos_eng(num_eng); // energy grid on [-0.5, 0.5] * eng_range
        for (int i = 0; i < num_eng; ++i) {
            dos_eng[i] = 0.5 * eng_range * (i * 1.0 / num_steps - 1.0);
        }

        // Search for energies that fall within config.dc_cond_energy_limits
        // In most cases the dataset is small so dynamic allocation of std::vector will not be a problem.
        std::vector<double> qe_eng;
        for (const auto& r : config.dc_cond_energy_limits) {
            double e_min = 0.0;
            double e_max = 0.0;
            std::tie(e_min, e_max) = r;
            for (const auto& d : dos_eng) {
                if ((e_min <= d) && (d <= e_max)) {
                    qe_eng.push_back(d);
                }
            }
        }

        // Sort and convert to eigen vector
        std::sort(qe_eng.begin(), qe_eng.end());
        meta.qe_energies = Eigen::VectorXd(qe_eng.size());
        for (size_t i = 0; i < qe_eng.size(); ++i) {
            meta.qe_energies[i] = qe_eng[i];
        }
        std::cout << "num_qe_energies: "<<meta.qe_energies.size()<<"\n";
        // Calculate correlation functions
        Eigen::VectorXcd corr_dos;
        Eigen::MatrixXcd corr_dc_xx;
        Eigen::MatrixXcd corr_dc_yy;
        auto backend = make_tbpm_backend<sparse_t>();
        backend->tbpm_dc_cond(h_sparse, curr_x, curr_y, config, meta,
            corr_dos, corr_dc_xx, corr_dc_yy);

        // Average over mpi processes
        mpi_env_.InplaceAllreduceVectorXcd(corr_dos);
        mpi_env_.InplaceAllreduceMatrixXcd(corr_dc_xx);
        mpi_env_.InplaceAllreduceMatrixXcd(corr_dc_yy);
        corr_dos /= static_cast<double>(meta.mpi_size);
        corr_dc_xx /= static_cast<double>(meta.mpi_size);
        corr_dc_yy /= static_cast<double>(meta.mpi_size);

        // Save data
        if (is_master()) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_complex("corr_dc_cond.dos", corr_dos);
            data_file.save_matrix_complex("corr_dc_cond.xx", corr_dc_xx);
            data_file.save_matrix_complex("corr_dc_cond.yy", corr_dc_yy);
            data_file.close();
#else
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_complex("corr_dc_cond_dos", corr_dos);
            data_file.save_matrix_complex("corr_dc_cond_xx", corr_dc_xx);
            data_file.save_matrix_complex("corr_dc_cond_yy", corr_dc_yy);
#endif
            save_info();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

    /**
     * @brief Calculate Hall conductivity using Kubo-Bastin method.
     *
     * @note
     * 1. Output files:
     *      ${prefix}_hall_cond.dat (no HDF5)
     *      ${prefix}_tbpm.h5, with datasets 'hall_cond.energies', 'hall_cond.hall_cond'
     *
     * 2. Data format of dat output:
     *      dimension: (num_eng, 2)
     *      columns: energies, hall_cond
     */
    void calc_hall_cond()
    {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        // Prepare Hamiltonian and current operators
        rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
        model_->build_ham_curr_csr(indptr_, indices_, ham_, curr_coeff_x_, curr_coeff_y_);
        rescale_ham();
        sparse_t h_sparse(indptr_, indices_, ham_);
        sparse_t curr_x(indptr_, indices_, curr_coeff_x_);
        sparse_t curr_y(indptr_, indices_, curr_coeff_y_);

        // Meta info
        TBPMMeta meta = get_default_meta();

        // Calculate \mu_{mn}
        Eigen::MatrixXcd mu_avg;
        auto backend = make_kubo_bastin_backend<sparse_t>();
        backend->kbdc_mu(h_sparse, curr_x, curr_y, config, meta, mu_avg);

        // Calculate DC conductivity
        Eigen::VectorXd cond;
        backend->cond_from_trace(config, meta, mu_avg, cond);
        double dckb_prefactor = 16 * model_->get_num_orbitals() / (base::PI * std::pow(get_energy_range(), 2));
        if (config.dimension == 2) {
            dckb_prefactor /= model_->get_lattice_area("c");
        } else {
            dckb_prefactor /= model_->get_lattice_volume();
        }
        if (config.dckb_unit == "h") {
            dckb_prefactor *= base::TWOPI;
        }
        cond *= dckb_prefactor;
        cond *= config.num_spin;

        // Average over MPI processes
        // NOTE: It may also be done after calling kbdc_mu.
        mpi_env_.InplaceAllreduceVectorXd(cond);
        cond /= static_cast<double>(meta.mpi_size);

        // Save data
        if (is_master()) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_real("hall_cond.energies", config.dckb_energies);
            data_file.save_matrix_real("hall_cond.hall_cond", cond);
            data_file.close();
#else
            Eigen::MatrixXd data(cond.size(), 2);
            data.col(0) = config.dckb_energies;
            data.col(1) = cond;
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_real("hall_cond", data);
#endif
            save_info();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

    /**
     * @brief Calculate time-dependent wave function.
     *
     * @note
     * 1. Output files:
     *      ${prefix}_wft.dat (no HDF5)
     *      ${prefix}_tbpm.h5, with datasets 'wft'
     *
     * 2. Data format of dat output:
     *      dimension: (num_orb, num_time*2)
     *      columns: (wft0.real, wft0.imag), (wft1.real, wft1.imag), ...
     */
    void calc_wft()
    {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        // Prepare Hamiltonian
        rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
        model_->build_ham_csr(indptr_, indices_, ham_);
        rescale_ham();
        sparse_t h_sparse(indptr_, indices_, ham_);

        // Prepare meta info.
        TBPMMeta meta = get_default_meta();
        meta.time_step *= config.wft_dt_scale;

        // Calculate time-dependent wave function
        Eigen::MatrixXcd wft;
        auto backend = make_tbpm_backend<sparse_t>();
        backend->tbpm_wft(h_sparse, config, meta, wft);

        // Output
        if (is_master() && config.save_data) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_complex("wft", wft);
            data_file.close();
#else
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_complex("wft", wft);
#endif
            save_info();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

    /**
     * @brief Calculate correlation function for band structure.
     *
     * @note
     * 1. Output files:
     *      ${prefix}_k_len.dat, ${prefix}_corr_bands.dat (no HDF5)
     *      ${prefix}_tbpm.h5, with datasets 'corr_bands.k_len' and 'corr_bands.corr_bands'
     *
     * 2. Data format of dat output of k_len:
     *      dimension: (num_kpt, 1)
     *
     *   Data format of dat output of corr_dos:
     *     dimension: (num_step+1, 2*num_kpt)
     *     columns: (corr_k0.real, corr_k0.imag), (corr_k1.real, corr_k1.imag), ...
     */
    void calc_corr_bands()
    {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        int num_kpt = config.k_points.cols();
        Eigen::VectorXd k_len = tbplas::base::gen_kdist(model_->get_lattice(), config.k_points);
        Eigen::MatrixXcd corr_bands(config.num_time_steps + 1, num_kpt);
        auto backend = make_tbpm_backend<sparse_t>();

        RangeInfo range = mpi_env_.distRange(num_kpt, mpi_env_.getSize());
        int num_kpt_range = range.k_max - range.k_min + 1;
        for (int ik = range.k_min; ik <= range.k_max; ++ik) {
            if (is_master()) {
                std::cout << "Kpoint " << ik << " of " << num_kpt_range << std::endl;
            }

            // Prepare Hamiltonian
            rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
            // Must use convention 1 since Python Sample class implements this convention only.
            // Otherwise, the Python interface always produces the Gamma-point Hamiltonian.
            model_->build_ham_csr(indptr_, indices_, ham_, config.k_points.col(ik), 1);
            rescale_ham();
            sparse_t h_sparse(indptr_, indices_, ham_);

            // Prepare meta info.
            TBPMMeta meta = get_default_meta();
            // Jobs are parallelized over k-points, so reset num_samples_dist.
            meta.num_samples_dist = config.num_random_samples;

            // Calculate correlation function
            Eigen::VectorXcd corr_dos;
            backend->tbpm_dos(h_sparse, config, meta, corr_dos);
            corr_bands.col(ik) = corr_dos;
        }

        // Average over MPI processes
        mpi_env_.InplaceAllreduceMatrixXcd(corr_bands);
        corr_bands /= static_cast<double>(mpi_env_.getSize());

        // Output
        if (is_master()) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_real("corr_bands.k_len", k_len);
            data_file.save_matrix_complex("corr_bands.corr_bands", corr_bands);
            data_file.close();
#else
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_real("k_len", k_len);
            data_file.save_matrix_complex("corr_bands", corr_bands);
#endif
            save_info();
        }
        // Master process needs to save data and finishes later than other processes.
        // So we need to put a barrier here.
        mpi_env_.barrier();
    }

    /**
     * @brief Calculate charge density rho(r).
     *
     * @note
     * 1. Output files:
     * ${prefix}_charge_density.dat (no HDF5)
     * ${prefix}_tbpm.h5, with datasets 'charge_density'
     *
     * 2. Data format of dat output:
     * dimension: (n_wf)
     * columns: rho_1, rho_2, ...
     * The charge density is normalized such that its sum over all sites equals `config.num_electron`.
     */
    void calc_charge_density() {
        using sparse_t = CSRMatrixSquare<size_t, Eigen::VectorXcd>;

        // 1. Prepare Hamiltonian
        rescale_ = 1.0; // Reset rescale_ since we rebuild Hamiltonian.
        model_->build_ham_csr(indptr_, indices_, ham_);
        rescale_ham();
        sparse_t h_sparse(indptr_, indices_, ham_);

        // std::map<std::string, double> total_time;
        // total_time = benchmark_ops(h_sparse,1000);
        // std::cout << "amxsy: "<< total_time["amxsy"] << " axpy: "<< total_time["axpy"]<<" amxsy/axpy: "<<total_time["amxsy"]/total_time["axpy"]<<"\n"; 
        if (config.opt_pool_size){
            run_benchmarks_and_analyze(h_sparse, config.num_internal_runs, config.num_external_runs);
        }

        // 2. Prepare Meta info
        TBPMMeta meta = get_default_meta();
        // Set Time step based on Nyquist theorem: 2 pi/(2 occ_width)
        meta.time_step = base::PI / config.occ_width;

        // 3. Calculate charge density
        Eigen::VectorXd charge_density_result; 
        
        auto backend = make_tbpm_backend<sparse_t>();
        backend->tbpm_cd(h_sparse, config, meta, charge_density_result);

        // 4. Average over MPI processes
        mpi_env_.InplaceAllreduceVectorXd(charge_density_result);
        charge_density_result /= static_cast<double>(meta.mpi_size);
        //Normalize to num_electron
        vec_self_mul_sca(charge_density_result, config.num_electron);

        // 5. Save charge density
        if (is_master()) {
#ifdef WITH_HDF5
            HDFFile data_file(config.prefix + "_tbpm", OptBinMode::UPDATE);
            data_file.save_matrix_real("charge_density", charge_density_result);
            data_file.close();
#else
            TextFile data_file(config.prefix, false);
            data_file.save_matrix_real("charge_density", charge_density_result);
#endif
            save_info(); // Assuming this saves general run info
        }

        // 6. MPI barrier to ensure all processes complete before continuing.
        mpi_env_.barrier();
    }

private:
    /**
     * @brief Get the default rescaling factor by summing up the absolute values
     * of each row: rescale = np.max(np.sum(np.abs(ham_dense), axis=1))
     *
     * @return rescaling factor in eV
     *
     * @note It seems that this function cannot be parallelized with OpenMP.
     */
    [[nodiscard]] double get_default_rescale() const
    {
        assert(ham_.size() > 0);
        size_t num_orb = indptr_.size() - 1;
        Eigen::VectorXd ham_sum = Eigen::VectorXd::Zero(num_orb);
        size_t i = 0; // row index
        size_t j = 0; // offset of indices and values
        for (i = 0; i < num_orb; ++i) {
            for (j = indptr_[i]; j < indptr_[i + 1]; ++j) {
                ham_sum[i] += ham_.cwiseAbs()(j);
            }
        }
        return ham_sum.maxCoeff();
    }

    /**
     * @brief Get the energy range in eV
     *
     * @return the energy range
     */
    [[nodiscard]] double get_energy_range() const
    {
        return 2.0 * rescale_;
    }

    /**
     * @brief Get the time step for TBPM calculation by:
     * time_step = 2 * pi / sample.energy_range
     *
     * @return time step in h_bar/eV
     */
    [[nodiscard]] double get_time_step() const
    {
        return base::TWOPI / get_energy_range();
    }

    /**
     * @brief Set up the default TBPMMeta object.
     *
     * @return TBPMMeta object
     *
     * @note This function sets up only the common attributes. Task-specific attributes
     * still need to be set in the caller of this function. In particular, the time step
     * is half of the default value in many kinds of calculations. I have wasted a whole
     * day on the incorrect time step.
     *
     */
    [[nodiscard]] TBPMMeta get_default_meta() const
    {
        TBPMMeta meta;
        meta.rescale = rescale_;
        meta.energy_range = get_energy_range();
        meta.time_step = get_time_step();
        meta.beta_re = 1.0 / (base::KB_EVK * config.temperature) * rescale_;
        meta.mu_re = config.mu / rescale_;
        meta.qe_energies = config.qe_energies;

        // Distribute jobs among MPI processes.
        int num_samples = config.num_random_samples;
        int rank = mpi_env_.getRank();
        int size = mpi_env_.getSize();
        while (num_samples % size != 0) {
            ++num_samples;
        }
        // if (is_master() && num_samples != config.num_random_samples) {
        //     std::cout << "INFO: increasing num_random_samples to "
        //               << num_samples << " for optimal balance.\n";
        // }
        meta.mpi_rank = rank;
        meta.mpi_size = size;
        meta.is_master = is_master();
        meta.num_samples_dist = num_samples / size;
        return meta;
    }

    /**
     * @brief Rescale the Hamiltonian such that it eigenvalues will fall within [-1, 1].
     *
     */
    void rescale_ham()
    {
        assert(ham_.size() > 0);
        double new_factor = config.rescale;
        if (new_factor <= 0.0) {
            new_factor = get_default_rescale();
        }
        vec_self_mul_sca(ham_, rescale_ / new_factor);
        rescale_ = new_factor;
    }

    /**
     * @brief Generate a backend for Haydock calculation.
     *
     * @tparam sparse_t datatype of sparse Hamiltonian
     * @return uniq_ptr to the backend
     */
    template <typename sparse_t>
    std::unique_ptr<AbstractHaydock<sparse_t>> make_haydock_backend() const
    {
        const std::regex pattern_gpu(R"(^gpu$)", std::regex_constants::icase);
        if (std::regex_match(config.algo, pattern_gpu)) {
#ifdef WITH_CUDA
            print_stdout("Using HaydockGPU backend");
            return std::make_unique<HaydockGPU<sparse_t>>();
#else
            print_stderr("ERROR: CUDA support disabled during compilation!");
            std::exit(-1);
#endif
        }
        print_stdout("Using HaydockCPU backend");
        return std::make_unique<HaydockCPU<sparse_t>>();
    }

    /**
     * @brief Generate a backend for Kubo-Bastin calculation
     *
     * @tparam sparse_t datatype of sparse Hamiltonian
     * @return uniq_ptr to the backend
     */
    template <typename sparse_t>
    std::unique_ptr<AbstractKuboBastin<sparse_t>> make_kubo_bastin_backend() const
    {
        const std::regex pattern_gpu(R"(^gpu$)", std::regex_constants::icase);
        if (std::regex_match(config.algo, pattern_gpu)) {
#ifdef WITH_CUDA
            print_stdout("Using KuboBastinGPU backend");
            return std::make_unique<KuboBastinGPU<sparse_t>>();
#else
            print_stderr("ERROR: CUDA support disabled during compilation!");
            std::exit(-1);
#endif
        }
        print_stdout("Using KuboBastinCPU backend");
        return std::make_unique<KuboBastinCPU<sparse_t>>();
    }

    /**
     * @brief Generate a backend for TBPM calculation
     *
     * @tparam sparse_t datatype of sparse Hamiltonian
     * @return uniq_ptr to the backend
     */
    template <typename sparse_t>
    std::unique_ptr<AbstractTBPM<sparse_t>> make_tbpm_backend() const
    {
        const std::regex pattern_gpu(R"(^gpu$)", std::regex_constants::icase);
        const std::regex pattern_cpu_fast(R"(^cpu_fast$)", std::regex_constants::icase);
        if (std::regex_match(config.algo, pattern_gpu)) {
#ifdef WITH_CUDA
            print_stdout("Using TBPMGPU backend");
            return std::make_unique<TBPMGPU<sparse_t>>();
#else
            print_stderr("ERROR: CUDA support disabled during compilation!");
            std::exit(-1);
#endif
        }
        if (std::regex_match(config.algo, pattern_cpu_fast)) {
            print_stdout("Using TBPMCPUFast backend");
            return std::make_unique<TBPMCPUFast<sparse_t>>();
        }
        print_stdout("Using TBPMCPU backend");
        return std::make_unique<TBPMCPU<sparse_t>>();
    }

    /**
     * @brief Save model and config information for postprocessing with Python 'Analyzer' class.
     *
     */
    void save_info() const
    {
        const int prec = 16;
        const int width = 32;
        std::ofstream outfile;

        // model info
        outfile.open(config.prefix + "_info.dat", std::ios::out);
        outfile.setf(std::ios::scientific);
        outfile.precision(prec);
        outfile << "# Model info\n";
        outfile << "energy_range"
                << std::setw(width) << get_energy_range() << "\n";

        outfile << "num_orbitals"
                << std::setw(width) << model_->get_num_orbitals() << "\n";

        outfile << "lattice_area"
                << std::setw(width) << model_->get_lattice_area("c") << "\n";

        outfile << "lattice_volume"
                << std::setw(width) << model_->get_lattice_volume() << "\n";

        // config info
        outfile << "\n# Config info\n";
        outfile << "dimension"
                << std::setw(width) << config.dimension << "\n";

        outfile << "num_time_steps"
                << std::setw(width) << config.num_time_steps << "\n";

        outfile << "num_spin"
                << std::setw(width) << config.num_spin << "\n";

        outfile << "temperature"
                << std::setw(width) << config.temperature << "\n";

        outfile << "dyn_pol_coulomb_factor"
                << std::setw(width) << config.dyn_pol_coulomb_factor << "\n";

        outfile << "dyn_pol_back_epsilon"
                << std::setw(width) << config.dyn_pol_back_epsilon << "\n";

        const Eigen::Matrix3Xd& q_points = config.dyn_pol_q_points;
        int num_qpt = q_points.cols();
        outfile << "begin dyn_pol_q_points"
                << std::setw(width) << num_qpt << std::setw(width) << 3 << "\n";
        for (int iq = 0; iq < num_qpt; ++iq) {
            outfile << std::setw(width) << q_points(0, iq)
                    << std::setw(width) << q_points(1, iq)
                    << std::setw(width) << q_points(2, iq) << "\n";
        }
        outfile << "end dyn_pol_q_points\n";

        outfile << "dc_cond_num_time_steps"
                << std::setw(width) << config.dc_cond_num_time_steps << "\n";

        outfile << "begin dc_cond_energy_limits"
                << std::setw(width) << config.dc_cond_energy_limits.size()
                << std::setw(width) << 2 << "\n";
        for (const auto& t : config.dc_cond_energy_limits) {
            outfile << std::setw(width) << std::get<0>(t)
                    << std::setw(width) << std::get<1>(t) << "\n";
        }
        outfile << "end dc_cond_energy_limits\n";
        outfile.close();
    }
};

} // namespace tbplas::tbpm
#endif // TBPLAS_TBPM_SOLVER_H
