/**
 * @brief Functions for fast TBPM algorithms
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 */

#ifndef TBPLAS_TBPM_TBPM_FAST_H
#define TBPLAS_TBPM_TBPM_FAST_H

#include <iostream>
#include <vector>

#include <eigen3/Eigen/Dense>

#include "backend.h"
#include "chebyshev.h"
#include "config.h"
#include "dense.h"
#include "fermi_dirac.h"
#include "propagation.h"
#include "tbpm/dense.h"
#include "utils.h"

namespace tbplas::tbpm {
using tbplas::base::ProgressBar;
using base::Timer;

template <typename sparse_t>
class TBPMCPUFast final : public AbstractTBPM<sparse_t> {
private:

    /**
     * @brief Shared function for calculating quasi-eigenstates from initial state,
     * using block-time propagation with MultiPropagator.
     *
     * @param[in] propagator MultiPropagator instance, handles propagation of multiple wavefunctions.
     * It changes its working arrays, so it is & rather than const &.
     * @param[in] config user-defined calculation configurations.
     * @param[in] meta additional meta info.
     * @param[in] wf0 initial wave function (state at time t=0 for the current sample).
     * @param[out] wfq quasi-eigenstates accumulated from this initial state, dimension (num_eng, n_wf).
     * MUST BE ALLOCATED and initialized (e.g., with the t=0 contribution) in the caller
     * before passing to this function.
     *
     * @note
     * DO NOT use auto for Eigen datatypes (vectors in wfq) to avoid potential pitfalls.
     */
    void tbpm_qe_block_time(
        MultiPropagator<sparse_t>& propagator,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        const Eigen::VectorXcd& wf0,
        std::vector<Eigen::VectorXcd>& wfq)
    {
        // Initialize wave functions for block propagation
        size_t n_wf = wf0.size();
        int pool_size = config.pool_size;

        // psi1[0] and psi2[0] hold the state at the start of the current pool.
        // psi1[k] and psi2[k] will hold states propagated to k*dt within the pool.
        std::vector<Eigen::VectorXcd> psi1;//positive propagation
        std::vector<Eigen::VectorXcd> psi2;//negative propagation
        mat_alloc(psi1, n_wf, pool_size + 1);
        mat_alloc(psi2, n_wf, pool_size + 1);

        // Copy initial state (t=0) to the starting slot for propagation
        vec_copy(wf0, psi1[0]);
        vec_copy(wf0, psi2[0]);

        // Initialize utilities
        int num_steps = config.num_time_steps;
        size_t num_energies = meta.qe_energies.size();

        ProgressBar prog_bar(num_steps, 64, "Finished timestep");

        // Add the t=0 contribution
        mat_col_copy(wf0, wfq);
        PARA_EXEC(prog_bar.count());

        // Iterate over time in blocks
        // The initial state (j=0) is typically handled before this loop
        // Because there we need consider term `num_steps`, so `i_step` from 1 not 2
        double norm_ref = vec_norm(psi2[0]);
        for (int i_step = 1; i_step <= num_steps; i_step += pool_size) {
            // Perform norm check if this block covers the check time step
            if ((i_step <= config.wfn_check_time_steps) && (i_step + pool_size >= config.wfn_check_time_steps)) {
                check_norm(psi1[0], norm_ref, config.wfn_check_threshold);
                check_norm(psi2[0], norm_ref, config.wfn_check_threshold);
            }

            // Propagate the wave functions for 'pool_size' steps
            // psi1[0] -> psi1[1...pool_size] (forward propagation)
            // psi2[0] -> psi2[1...pool_size] (backward propagation)
            propagator.act(psi1, true);
            propagator.act(psi2, false);

            // Process each time step within the current propagation block
            for (int k = 1; k <= pool_size; ++k) {
                // 'j' is the actual global time step index, 
                // because i_step begins from 1, so `-1` not `-2`.
                int j = i_step + k - 1;

                // Apply Hanning window, calculated for the current global time step 'j'
                double window = 0.5 * (1 + std::cos(base::PI * j / static_cast<double>(num_steps)));

                // Ensure we don't go beyond the total number of time steps
                if (j <= num_steps) {
                    // Perform Fourier transform accumulation for each energy
                    for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                        double phase = meta.qe_energies[i_e] * j * meta.time_step;
                        complex_t factor = std::exp(base::UIMAG * phase);
                        complex_t f1 = factor * window;
                        complex_t f2 = std::conj(factor) * window;
                        // Accumulate weighted contributions into wfq[i_e]
                        axpbypz(f1, psi1[k], f2, psi2[k], wfq[i_e]);
                    }
                    PARA_EXEC(prog_bar.count());
                }
            }
            // Copy the last propagated state from this block (psiX.back() == psiX[pool_size])
            // to the initial slot (psiX[0]) for the next iteration's propagation.
            vec_copy(psi1.back(), psi1[0]);
            vec_copy(psi2.back(), psi2[0]);
        }

        // Normalization of accumulated quasi-eigenstates
        for (Eigen::VectorXcd& v : wfq) {
            double norm = vec_norm(v);
            vec_self_div_sca(v, norm);
        }
    }

public:
    /**
     * @brief Calculate correlation function for DOS and LDOS
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation configurations
     * @param[in] meta additional meta info.
     * @param[out] corr correlation function of (L)DOS, dimension: (num_steps+1, 1)
     */
    void tbpm_dos(
        const sparse_t& h_sparse,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::VectorXcd& corr) final
    {
        PARA_EXEC(std::cout << "Calculating DOS correlation function." << std::endl);
        PARA_EXEC(std::cout << "Using fast algorithm." << std::endl);
        Timer timer;
        
        // Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
        Eigen::VectorXcd wf0(n_wf);

        // Initialize correlation function
        int num_steps = config.num_time_steps;
        corr = Eigen::VectorXcd::Zero(num_steps + 1);
        timer.tic("init fast");
        // Initialize operators
        CDOSCalculator<sparse_t> cdos_calc(h_sparse, config, meta);

        // Initialize utilities
        RandomGenerator rand_gen(config.seed);
        timer.toc("init fast");
        timer.tic("run");
        // Loop over samples
        int num_samples = meta.num_samples_dist;
        int rank = meta.mpi_rank;
        for (int i_sample = 1; i_sample <= num_samples; ++i_sample) {
            PARA_EXEC(std::cout << "Sample " << i_sample << " of " << num_samples << std::endl);
            int seed_i = config.seed * (i_sample + rank * num_samples);

            if (config.ldos) {
                // Generate a random vector for selected orbitals only
                Eigen::VectorXcd wf_proj(config.ldos_orbital_indices.size());
                rand_gen.random_state(wf_proj, seed_i);

                // Assign wf_proj to wf0
                vec_init(wf0, 0.0);
                size_t counter = 0;
                for (const size_t& i_orb : config.ldos_orbital_indices) {
                    wf0(i_orb) = wf_proj(counter);
                    ++counter;
                }
            } else {
                rand_gen.random_state(wf0, seed_i);
            }
            Eigen::VectorXcd corr_work;
            cdos_calc.eval(wf0, corr_work);
            //TODO. 这儿可以考虑输出corr_work，因为有时候需要分析不同random sample之间的方差来统计误差，
            // 以及gap体系的费米能级通过多个random sample的平均会更靠近gap中心。
            vec_self_add_vec(corr, corr_work);
        }

        // Average over samples
        corr /= static_cast<double>(num_samples);
        timer.toc("run");
        timer.report_total_time();
    }

    /**
     * @brief Calculate correlation function of AC conductivity.
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] curr_x sparse current operator along x direction
     * @param[in] curr_y sparse current operator along y direction
     * @param[in] config user-defined calculation configurations
     * @param[in] meta additional meta info.
     * @param[out] corr correlation function of AC conductivity, dimension: (4, num_steps)
     */
    void tbpm_ac_cond(
        const sparse_t& h_sparse,
        const sparse_t& curr_x,
        const sparse_t& curr_y,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::MatrixXcd& corr) final
    {
        PARA_EXEC(std::cout << "Calculating AC conductivity correlation function." << std::endl);
        PARA_EXEC(std::cout << "Using fast algorithm with pool size = " << config.pool_size << std::endl);

        // Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
        int pool_size = config.pool_size;
        Eigen::VectorXcd wf0(n_wf);
        std::vector<Eigen::VectorXcd> psi1_x;
        std::vector<Eigen::VectorXcd> psi1_y;
        std::vector<Eigen::VectorXcd> psi2;
        mat_alloc(psi1_x, n_wf, pool_size + 1); // Additional 1 for t = 0 element
        mat_alloc(psi1_y, n_wf, pool_size + 1);
        mat_alloc(psi2, n_wf, pool_size + 1);

        // Initialize correlation function
        int num_steps = config.num_time_steps;
        corr = Eigen::MatrixXcd::Zero(4, num_steps);

        // Initialize operators
        // Current operators are bare sparse matrices as function arguments.
        MultiPropagator<sparse_t> propagator(h_sparse, config, meta);
        FermiDiracOperator<sparse_t> fermi_dirac(h_sparse, config, meta);

        // Initialize utilities
        RandomGenerator rand_gen(config.seed);
        ProgressBar prog_bar(num_steps, 64, "Finished timestep");

        // Loop over samples
        int num_samples = meta.num_samples_dist;
        int rank = meta.mpi_rank;
        for (int i_sample = 1; i_sample <= num_samples; ++i_sample) {
            PARA_EXEC(std::cout << "Sample " << i_sample << " of " << num_samples << std::endl);
            int seed_i = config.seed * (i_sample + rank * num_samples);
            PARA_EXEC(prog_bar.reset());

            // Make random state and psi1, psi2
            // See eqn. 53 of CPC paper.
            rand_gen.random_state(wf0, seed_i);
            curr_x.mv(wf0, psi1_x[0]); // j_beta|psi> for beta = x
            curr_y.mv(wf0, psi1_y[0]); // j_beta|psi> for beta = y

            fermi_dirac.act(psi1_x[0], true); // [1-f(H)]j_beta|psi> for beta = x
            fermi_dirac.act(psi1_y[0], true); // [1-f(H)]j_beta|psi> for beta = y
            fermi_dirac.act_out(wf0, psi2[0], false); // f(H)|psi>

            // Get initial correlation function in all directions
            corr(0, 0) += curr_x.xmy(psi2[0], psi1_x[0]); // <psi|f(H)j_alpha[1-f(H)]j_beta|psi> for xx
            corr(1, 0) += curr_y.xmy(psi2[0], psi1_x[0]); // yx
            corr(2, 0) += curr_x.xmy(psi2[0], psi1_y[0]); // xy
            corr(3, 0) += curr_y.xmy(psi2[0], psi1_y[0]); // yy
            PARA_EXEC(prog_bar.count());

            // Iterate over time
            // Norm check is performed on psi2. It can also be done for psi1_x and psi1_y.
            double norm_ref = vec_norm(psi2[0]);
            for (int i_step = 2; i_step <= num_steps; i_step += pool_size) {
                if ((i_step <= config.wfn_check_time_steps) && (i_step + pool_size >= config.wfn_check_time_steps)) {
                    check_norm(psi2[0], norm_ref, config.wfn_check_threshold);
                }
                // TODO. 如果要进一步节约内存可以这样做，这儿的x和y方向是独立的，psi1_x和psi1_y只保留一个就行，
                // 先传播psi2, 然后传播x的，获得x对应的所有的corr(0 and 1, j)之后再计算y的。
                // y方向的不需要保存pool_size+1个向量，而只需要保存一个向量psi1_y[0]作为下次传播的初始起点即可。
                // 
                propagator.act(psi1_x);
                propagator.act(psi1_y);
                propagator.act(psi2);

                // Loop over wave functions at t/k = 1, 2, ..., pool_size
                // e.g, for i_step = 2 and k = 1, j = 2 + 1 - 2 = 1
                for (int k = 1; k <= pool_size; ++k) {
                    int j = i_step + k - 2;
                    if (j < num_steps) { // NOT '=' because the maximum of num_step is num_steps - 1
                        corr(0, j) += curr_x.xmy(psi2[k], psi1_x[k]);
                        corr(1, j) += curr_y.xmy(psi2[k], psi1_x[k]);
                        corr(2, j) += curr_x.xmy(psi2[k], psi1_y[k]);
                        corr(3, j) += curr_y.xmy(psi2[k], psi1_y[k]);
                        PARA_EXEC(prog_bar.count());
                    }
                }

                // Copy last state as the next initial state
                vec_copy(psi1_x.back(), psi1_x[0]);
                vec_copy(psi1_y.back(), psi1_y[0]);
                vec_copy(psi2.back(), psi2[0]);
            }
        }

        // Average over samples
        corr /= static_cast<double>(num_samples);
    }
    /**
     * @brief Calculate correlation function at given q-point
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] den_coeff_q coefficients for density operator at q
     * @param[in] den_coeff_mq coefficients for density operator at -q
     * @param[in] config user-defined calculation configurations
     * @param[in] meta additional meta info.
     * @param[out] corr correlation function of dynamic polarizability, dimension: (num_step, 1)
     */
    void tbpm_dyn_pol(
        const sparse_t& h_sparse,
        const Eigen::VectorXcd& den_coeff_q,
        const Eigen::VectorXcd& den_coeff_mq,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::VectorXd& corr) final
    {
        PARA_EXEC(std::cout << "Calculating dynamic polarizability correlation function." << std::endl);
        PARA_EXEC(std::cout << "Using fast algorithm with pool size = " << config.pool_size << std::endl);

        // Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
        int pool_size = config.pool_size;
        Eigen::VectorXcd wf0(n_wf);
        std::vector<Eigen::VectorXcd> psi1;
        std::vector<Eigen::VectorXcd> psi2;
        mat_alloc(psi1, n_wf, pool_size + 1);
        mat_alloc(psi2, n_wf, pool_size + 1);

        // Initialize correlation function
        int num_steps = config.num_time_steps;
        corr = Eigen::VectorXd::Zero(num_steps);

        // Initialize operators
        // Current operators are bare sparse matrices as function arguments.
        MultiPropagator<sparse_t> propagator(h_sparse, config, meta);
        FermiDiracOperator<sparse_t> fermi_dirac(h_sparse, config, meta);

        // Initialize utilities
        RandomGenerator rand_gen(config.seed);
        ProgressBar prog_bar(num_steps, 64, "Finished timestep");

        int num_samples = meta.num_samples_dist;
        int rank = meta.mpi_rank;
        for (int i_sample = 1; i_sample <= num_samples; ++i_sample) {
            PARA_EXEC(std::cout << "Sample " << i_sample << " of " << num_samples << std::endl);
            int seed_i = config.seed * (i_sample + rank * num_samples);
            PARA_EXEC(prog_bar.reset());

            // Make random state and psi1, psi2
            // See eqn.79-80 of CPC paper.
            rand_gen.random_state(wf0, seed_i);
            vec_mul_vec(den_coeff_mq, wf0, psi1[0]); // psi1 = rho(-q) * psi0
            fermi_dirac.act(psi1[0], true); // psi1 = (1-nF) * psi1
            fermi_dirac.act_out(wf0, psi2[0], false); // psi2 = nF * psi0

            // Initial value of correlation function
            // See eqn.77 of CPC paper.
            corr[0] += xmy_diag(psi2[0], den_coeff_q, psi1[0]).imag(); // corr = Im[<psi2|rho(q)|psi1>]
            PARA_EXEC(prog_bar.count());

            // Iterate over time
            // Norm check is performed on psi2. It can also be done for psi1_x and psi1_y.
            double norm_ref = vec_norm(psi2[0]);
            // NOTE: i_step is from 2 to  num_steps, i.e. the maximum of num_step
            // is num_steps - 1 NOT num_steps_dc. This is different from DOS.
            for (int i_step = 2; i_step <= num_steps; i_step += pool_size) {
                if ((i_step <= config.wfn_check_time_steps) && (i_step + pool_size >= config.wfn_check_time_steps)) {
                    check_norm(psi2[0], norm_ref, config.wfn_check_threshold);
                }
                propagator.act(psi1);
                propagator.act(psi2);

                for (int k = 1; k <= pool_size; ++k) {
                    int j = i_step + k - 2;
                    if (j < num_steps) {// NOT '=' because the maximum of num_step is num_steps - 1
                        corr[j] += xmy_diag(psi2[k], den_coeff_q, psi1[k]).imag();
                        PARA_EXEC(prog_bar.count());
                    }
                }

                vec_copy(psi1.back(), psi1[0]);
                vec_copy(psi2.back(), psi2[0]);
            }
        }

        // Average over samples
        corr /= static_cast<double>(num_samples);
    }
    /**
     * @brief Calculate squared modulus of quasi-eigenstates for given energies.
     *
     * This function provides different implementations for calculating quasi-eigenstates
     * based on the `config.qe_method` setting:
     *
     * - **"cheby_poly" (Chebyshev Polynomial Method):** This is typically faster when you
     * need to calculate quasi-eigenstates for a *small number* of energies. It often
     * leverages a QECalculator based on Chebyshev polynomial expansion.
     *
     * - **"block_time" (Block-Time Propagation Method):** This method is often faster for
     * systems where you need to calculate quasi-eigenstates for a *large number* of energies,
     * especially for very large systems. It utilizes a MultiPropagator to process
     * a block of time steps.
     *
     *
     * @param[in] h_sparse scaled sparse Hamiltonian.
     * @param[in] config user-defined calculation configurations, including 'qe_method' to select the algorithm.
     * @param[in] meta additional meta info.
     * @param[out] wf2_qe squared modulus of quasi-eigenstates, dimension: (n_wf, num_eng).
     */
    void tbpm_quasi_eigenstates(
        const sparse_t& h_sparse,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::MatrixXd& wf2_qe) final
    {
        PARA_EXEC(std::cout << "Calculating quasi-eigenstates." << std::endl);

        // Initialize common variables needed regardless of the chosen method
        size_t n_wf = h_sparse.get_dim();
        size_t num_energies = meta.qe_energies.size();
        Eigen::VectorXcd wf0(n_wf); // Represents the initial random state for each sample

        // Initialize the final result matrix to zeros. This matrix accumulates results
        // from all samples and is the ultimate output of this function.
        wf2_qe = Eigen::MatrixXd::Zero(n_wf, num_energies);

        // This vector will temporarily hold the accumulated quasi-eigenstates for a *single sample*.
        // Its content is reset for each new sample.
        std::vector<Eigen::VectorXcd> wf_qe_sample;
        mat_alloc(wf_qe_sample, n_wf, num_energies);

        // Initialize random number generator for producing different initial states
        RandomGenerator rand_gen(config.seed);

        // Factors for averaging the squared modulus of quasi-eigenstates over multiple samples
        int num_samples = meta.num_samples_dist;
        int rank = meta.mpi_rank;
        double avg_factor = 1.0 / static_cast<double>(num_samples);
        
        // Loop over different random initial states (samples) to average results
        for (int i_sample = 1; i_sample <= num_samples; ++i_sample) {
            PARA_EXEC(std::cout << "Sample " << i_sample << " of " << num_samples << std::endl);
            int seed_i = config.seed * (i_sample + rank * num_samples);

            // Generate a new random initial state (wf0) for the current sample
            rand_gen.random_state(wf0, seed_i);


            // --- Algorithm Selection based on config.qe_method ---
            if (config.qe_method == "block_time") {
                // --- Block-Time Propagation Method (optimized for few energies, large systems) ---
                PARA_EXEC(std::cout << "Using Block-Time Propagation method with pool size = " << config.pool_size << std::endl);

                // Create the MultiPropagator instance, which will be passed to tbpm_qe_block_time
                MultiPropagator<sparse_t> propagator(h_sparse, config, meta);
                tbpm_qe_block_time(propagator, config, meta, wf0, wf_qe_sample);

            } else if (config.qe_method == "cheby_poly") {
                // --- Chebyshev Polynomial Method (default, for many energies) ---
                PARA_EXEC(std::cout << "Using QECalculator" << std::endl);

                QECalculator<sparse_t> qe_calc(h_sparse, config, meta);
                qe_calc.eval(wf0, wf_qe_sample);

            } else {
                // Handle unsupported method gracefully
                PARA_EXEC(std::cerr << "Error: Unsupported quasi-eigenstate method: " << config.qe_method << std::endl);
                std::exit(-1); 
            }

            // After calculating wf_qe_sample for the current sample,
            // evaluate its squared modulus and add it to the running average.
            for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                wf2_qe.col(i_e) += wf_qe_sample[i_e].cwiseAbs2() * avg_factor;
            }
        }
    }

    /**
     * @brief Calculate correlation functions of DC conductivity.
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] curr_x sparse current operator along x direction
     * @param[in] curr_y sparse current operator along y direction
     * @param[in] config user-defined calculation configurations
     * @param[in] meta additional meta info.
     * @param[out] corr_dos correlation function of DOS, dimension (num_steps_dos, 1)
     * @param[out] corr_dc_xx correlation function of xx component of DC conductivity,
     *   dimension (num_energies, num_steps_dc)
     * @param[out] corr_dc_yy correlation function of yy component of DC conductivity,
     *   dimension (num_energies, num_steps_dc)
     *
     * @note This function is twice slower than the default implementation, although it is named with
     *   'fast'. For now it is kept for testing purposes. DO NOT use it for production calculations.
     */
    void tbpm_dc_cond(
        const sparse_t& h_sparse,
        const sparse_t& curr_x,
        const sparse_t& curr_y,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::VectorXcd& corr_dos,
        Eigen::MatrixXcd& corr_dc_xx,
        Eigen::MatrixXcd& corr_dc_yy) final
    {
        PARA_EXEC(std::cout << "Calculating DC conductivity correlation function." << std::endl);
        PARA_EXEC(std::cout << "Using fast algorithm with pool size = " << config.pool_size << std::endl);

        // Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
        int pool_size = config.pool_size;
        Eigen::VectorXcd wf0(n_wf);

        // Initialize quasi-eigenstates
        size_t num_energies = meta.qe_energies.size();
        std::vector<Eigen::VectorXcd> wf_qe; // quasi-eigenstates and j_x|psi(E)>
        std::vector<Eigen::VectorXcd> psi2_y; // j_x|psi(E)>
        Eigen::VectorXcd buffer(n_wf);
        mat_alloc(wf_qe, n_wf, num_energies);
        mat_alloc(psi2_y, n_wf, num_energies);

        // Initialize correlation functions
        int num_steps_dos = config.num_time_steps;
        int num_steps_dc = config.dc_cond_num_time_steps;
        Eigen::VectorXcd corr_dos_work;
        corr_dos = Eigen::VectorXcd::Zero(num_steps_dos);
        corr_dc_xx = Eigen::MatrixXcd::Zero(num_energies, num_steps_dc);
        corr_dc_yy = Eigen::MatrixXcd::Zero(num_energies, num_steps_dc);

        // Initialize operators and utilities
        CDOSCalculator<sparse_t> dos_calc(h_sparse, config, meta);
        // QECalculator<sparse_t> qe_calc(h_sparse, config, meta);
        MultiPropagator<sparse_t> propagator(h_sparse, config, meta);
        RandomGenerator rand_gen(config.seed);
        ProgressBar prog_bar(num_steps_dc, 64, "Finished DC timestep");

        // Main loop
        int num_samples = meta.num_samples_dist;
        int rank = meta.mpi_rank;
        for (int i_sample = 1; i_sample <= num_samples; ++i_sample) {
            PARA_EXEC(std::cout << "Sample " << i_sample << " of " << num_samples << std::endl);
            int seed_i = config.seed * (i_sample + rank * num_samples);

            //Get corr_dos for this initial condition 
            rand_gen.random_state(wf0, seed_i);
            PARA_EXEC(std::cout << "Calculating corr_dos." << std::endl);
            dos_calc.eval(wf0, corr_dos_work);
            // corr_dos_work from dos_calc.eval() has length of num_steps_dos+1
            // The first term (t=0) should be dropped.
            #pragma omp parallel for simd
            for (int i = 0; i < num_steps_dos; ++i) {
                corr_dos[i] += corr_dos_work[i + 1];
            }

            //Calculate Quasi-Eigenstates (QE)
            PARA_EXEC(std::cout << "Calculating qe." << std::endl);
            if (config.qe_method == "block_time") {
                // Block-Time Propagation Method for QE calculation
                MultiPropagator<sparse_t> qe_propagator(h_sparse, config, meta);
                tbpm_qe_block_time(qe_propagator, config, meta, wf0, wf_qe);

            } else if (config.qe_method == "cheby_poly") {
                // Chebyshev Polynomial Method for QE calculation
                QECalculator<sparse_t> qe_calc(h_sparse, config, meta);
                qe_calc.eval(wf0, wf_qe); 

            } else {
                PARA_EXEC(std::cerr << "Error: Unsupported quasi-eigenstate method: " << config.qe_method << std::endl);
                std::exit(-1);
            }   
              
            
            // Calculate |<psi0|psi(E)>|
            PARA_EXEC(std::cout << "Calculating corr_dc." << std::endl);
            Eigen::VectorXd factor(num_energies);
            for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                factor[i_e] = std::abs(vec_dot(wf0, wf_qe[i_e]));
            }

            if (config.dc_method == "block_time") {

                PARA_EXEC(prog_bar.reset());
                std::vector<Eigen::VectorXcd> psi0_x; // j_x|psi0> in eqn.62 of CPC paper
                std::vector<Eigen::VectorXcd> psi0_y; // j_y|psi0>
                mat_alloc(psi0_x, n_wf, pool_size + 1);
                mat_alloc(psi0_y, n_wf, pool_size + 1);

                // Apply current operators
                curr_x.mv(wf0, psi0_x[0]);
                curr_y.mv(wf0, psi0_y[0]);
                for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                    curr_y.mv(wf_qe[i_e], psi2_y[i_e]);
                    curr_x.mv(wf_qe[i_e], buffer);
                    vec_copy(buffer, wf_qe[i_e]);
                }
                
                // Initial corr_dc_*
                const std::vector<Eigen::VectorXcd>& psi2_x = wf_qe;
                for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                    corr_dc_xx(i_e, 0) += vec_dot(psi0_x[0], psi2_x[i_e]) / factor[i_e];
                    corr_dc_yy(i_e, 0) += vec_dot(psi0_y[0], psi2_y[i_e]) / factor[i_e];
                }
                PARA_EXEC(prog_bar.count());

                // Propagation
                for (int i_step = 2; i_step <= num_steps_dc; i_step += pool_size) {
                    // TODO. 如果要进一步节约内存可以这样做，这儿的x和y方向是独立的，psi0_x和psi0_y只保留一个就行，
                    // 先传播x的，获得x对应的corr之后再计算y的。y方向的不需要保存pool_size+1个向量，而只需要保存一个
                    // 向量psi0_y[0]作为下次传播的初始起点即可。
                    propagator.act(psi0_x);
                    propagator.act(psi0_y);
                    for (int k = 1; k <= pool_size; ++k) {
                        int j = i_step + k - 2;
                        if (j < num_steps_dc) { // NOT '=' because the maximum of num_step is  num_steps_dc - 1
                            for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                                corr_dc_xx(i_e, j) += vec_dot(psi0_x[k], psi2_x[i_e]) / factor[i_e];
                                corr_dc_yy(i_e, j) += vec_dot(psi0_y[k], psi2_y[i_e]) / factor[i_e];
                            }
                            PARA_EXEC(prog_bar.count());
                        }
                    }
                    vec_copy(psi0_x.back(), psi0_x[0]);
                    vec_copy(psi0_y.back(), psi0_y[0]);
                }
            } else if (config.dc_method == "cheby_poly") {
                Eigen::VectorXcd psi0_x(n_wf);// j_x|psi0>
                Eigen::VectorXcd psi0_y(n_wf);// j_y|psi0>
                // Apply current operators
                curr_x.mv(wf0, psi0_x);
                curr_y.mv(wf0, psi0_y);
                for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                    curr_y.mv(wf_qe[i_e], psi2_y[i_e]);
                    curr_x.mv(wf_qe[i_e], buffer);
                    vec_copy(buffer, wf_qe[i_e]);
                }
                // Chebyshev Polynomial Method for DC calculation
                CDCCalculator<sparse_t> cdc_calc(h_sparse, config, meta);
                cdc_calc.eval(psi0_x, psi0_y, factor, wf_qe, psi2_y, corr_dc_xx, corr_dc_yy); 

            } else {
                PARA_EXEC(std::cerr << "Error: Unsupported DC conductivity method: " << config.dc_method << std::endl);
                std::exit(-1);
            }  
        }

        // Average over samples
        corr_dos /= static_cast<double>(num_samples);
        corr_dc_xx /= static_cast<double>(num_samples);
        corr_dc_yy /= static_cast<double>(num_samples);
    }

    /**
     * @brief Calculate time-dependent wave function.
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation configurations
     * @param[in] meta additional meta info.
     * @param[out] wft time-dependent wave functions, dimension (n_wf, num_save)
     */
    void tbpm_wft(
        const sparse_t& h_sparse,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::MatrixXcd& wft) final
    {
        PARA_EXEC(std::cout << "Calculating time-dependent wave function." << std::endl);

        // Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
        size_t num_save = config.wft_time_save.size();
        wft = Eigen::MatrixXcd::Zero(n_wf, num_save);

        // Put wf0 in the 0-th position in wft_buffer
        // wf0 must be copy of wft_wf0 since we are going to normalize it
        Eigen::VectorXcd wf0 = config.get_wft_wf0();
        vec_self_div_sca(wf0, vec_norm(wf0));
        std::vector<Eigen::VectorXcd> wft_buffer;
        mat_alloc(wft_buffer, n_wf, num_save + 1);
        vec_copy(wf0, wft_buffer[0]);
        Eigen::VectorXcd t0(n_wf);
        Eigen::VectorXcd t1(n_wf);

        // Build bessel table for selected time steps
        std::vector<size_t> k_vec;
        k_vec.push_back(0);
        for (size_t k = 0; k <= config.num_time_steps; ++k) {
            if (config.wft_time_save.find(k) != config.wft_time_save.end()) {
                k_vec.push_back(k);
            }
        }
        std::vector<Eigen::VectorXcd> bes_table;
        for (const auto& k : k_vec) {
            double time_step_scaled = meta.time_step * meta.rescale * k;
            Eigen::VectorXcd bes_vec;
            build_bessel_vector(time_step_scaled, config.bessel_precision, bes_vec);
            size_t num_bes = bes_vec.size();
#pragma omp parallel for simd
            for (int n = 0; n < num_bes; ++n) {
                bes_vec[n] *= 2.0 * std::pow(-UIMAG, n);
            }
            bes_vec[0] *= 0.5;
            bes_table.push_back(bes_vec);
        }

        // Evaluate time-dependent wave functions
        // We call functions in chebyshev.h without bothering OOP.
        cheb_eval_wf_multi(&h_sparse, bes_table, t0, t1, wft_buffer);

        // Copy wft_buffer to wft
        // NOTE: If wave function at t = 0 is to be saved, we have to use the initial state rather than
        // the one from Chebyshev expansion with t = 0, since the latter is wrong.
        for (int i_frame = 0; i_frame < num_save; ++i_frame) {
            if (k_vec[i_frame + 1] == 0) {
                wft.col(i_frame) = wf0;
            } else {
                wft.col(i_frame) = wft_buffer[i_frame + 1];
            }
        }
    }

    /**
     * @brief Calculate charge density rho(r) via Block Time Evolution.
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation configurations
     * @param[in] meta additional meta info.
     * @param[out] charge_density output charge density vector, dimension: (n_wf).
     *      The elements of this vector are normalized such that their sum over all sites is 1.
     */
    void tbpm_cd(
        const sparse_t& h_sparse,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::VectorXd& charge_density) final
    {
        PARA_EXEC(std::cout << "Calculating charge density." << std::endl);
         PARA_EXEC(std::cout << "Using Block-Time Propagation method with pool size = " << config.pool_size << std::endl);

        // 1. Initialize wave functions  for block propagation
        size_t n_wf = h_sparse.get_dim();
        int pool_size = config.pool_size;
       
        Eigen::VectorXcd psi(n_wf);           // Initial random state |psi>
        std::vector<Eigen::VectorXcd> wf_t_pos; // //positive propagation
        std::vector<Eigen::VectorXcd> wf_t_neg; // negative propagation
        mat_alloc(wf_t_pos, n_wf, pool_size + 1);
        mat_alloc(wf_t_neg, n_wf, pool_size + 1);

        // 2. Initialize charge density vector
        charge_density = Eigen::VectorXd::Zero(n_wf);
        Eigen::VectorXd buffer(n_wf); // Temporary vector to store charge density

        // 3. Initialize operators and utilities
        FermiDiracOperator<sparse_t> fermi_dirac(h_sparse, config, meta);
        
        MultiPropagator<sparse_t> propagator(h_sparse, config, meta);
        RandomGenerator rand_gen(config.seed);
        ProgressBar prog_bar(config.num_time_steps, 64, "Finished timestep for CD"); // Assuming config.num_time_steps for N_Dt

        // Loop over samples
        int num_samples = meta.num_samples_dist;
        int rank = meta.mpi_rank;
        int num_steps = config.num_time_steps;

        for (int i_sample = 1; i_sample <= num_samples; ++i_sample) {
            PARA_EXEC(std::cout << "Sample " << i_sample << " of " << num_samples << std::endl);
            int seed_i = config.seed * (i_sample + rank * num_samples);
            PARA_EXEC(prog_bar.reset());

            // a. Generate initial random state |psi>
            rand_gen.random_state(psi, seed_i);

            // b. Apply Fermi-Dirac operator: |psi^CD(t_0)> <- f^{1/2}(H) |psi>
            fermi_dirac.act_sqrt_fd(psi, wf_t_pos[0]);
            vec_copy(wf_t_pos[0], wf_t_neg[0]); 

            // c. Initial accumulation for t = 0
            // rho(r) <- |<r|psi^CD(t_0)>|^2
            vec_init(buffer,0.0);
            vec_abs2_add(wf_t_pos[0],buffer);

            // Store the initial norm if you plan to do norm checks during evolution
            double norm_ref = tbplas::tbpm::vec_norm(wf_t_pos[0]);
            // std::cout<<"pos and neg: 0 "<<wf_t_pos[0][100]<<"  "<<wf_t_neg[0][100]<<std::endl;

            // d. Iterate over time in blocks
            for (int i_step = 2; i_step <= num_steps; i_step += pool_size) {
                if (i_step == config.wfn_check_time_steps) {
                    tbplas::tbpm::check_norm(wf_t_pos[0], norm_ref, config.wfn_check_threshold);
                }
                // Propagate the wave functions for 'pool_size' steps
                // wf_t_pos[0] -> wf_t_pos[1...pool_size] (forward propagation)
                // wf_t_neg[0] -> wf_t_neg[1...pool_size] (backward propagation)
                propagator.act(wf_t_pos, true); 
                propagator.act(wf_t_neg, false); 
                


                // e. Accumulate charge density for current time step
                 // Process each time step within the current propagation block
                for (int k = 1; k <= pool_size; ++k) {
                    // 'j' is the actual global time step index
                    int j = i_step + k - 2;
                    // Ensure we don't go beyond the total number of time steps
                    if (j <= num_steps) {
                        // rho(r) <- rho(r) + |<r|psi^CD(t_j)>|^2 + |<r|psi^CD(-t_j)>|^2
                        // std::cout<<j<<" pos and neg: 0 "<<wf_t_pos[k][100]<<"  "<<wf_t_neg[k][100]<<std::endl;
                        vec_abs2_add(wf_t_pos[k],buffer);
                        vec_abs2_add(wf_t_neg[k],buffer);
                        PARA_EXEC(prog_bar.count());
                    }
                }
                // Copy last state as the next initial state
                vec_copy(wf_t_pos.back(), wf_t_pos[0]);
                vec_copy(wf_t_neg.back(), wf_t_neg[0]);
            }
            vec_self_div_sca(buffer, 2*num_steps + 1);
            // Calculate squared norm and normalize
            double norm = buffer.sum();
            std::cout << "Sample " << i_sample << " norm: " << norm << std::endl;
            vec_self_div_sca(buffer, norm);
            vec_self_add_vec(charge_density, buffer);

        }

        charge_density /= static_cast<double>(num_samples);
    }
};

} // namespace tbplas::tbpm

#endif