/**
 * @brief Functions for performing TBPM calculations
 *
 * @author: Zewen Wu, Miao Zhang, Yunhai Li
 */

#ifndef TBPLAS_TBPM_TBPM_H
#define TBPLAS_TBPM_TBPM_H

#include <iostream>
#include <vector>

#include <eigen3/Eigen/Dense>

#include "backend.h"
#include "base/consts.h"
#include "base/utils.h"
#include "config.h"
#include "dense.h"
#include "fermi_dirac.h"
#include "propagation.h"
#include "utils.h"

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

template <typename sparse_t>
class TBPMCPU final : public AbstractTBPM<sparse_t> {
private:
    /**
     * @brief Shared function for calculating quasi-eigenstates (and corr_dos) from initial state
     *
     * @param[in] propagator propagator
     * @param[in] config user-defined calculation configurations
     * @param[in] meta additional meta info.
     * @param[in] calc_corr_dos whether to calculate DOS correlation function
     * @param[in] wf0 initial wave function
     * @param[out] wfq quasi-eigenstates from initial state, dimension (num_eng, num_wf),
     *   MUST BE ALLOCATED in the caller
     * @param[out] corr_dos correlation function of dos, dimension (num_step),
     *   MUST be ALLOCATED in the caller if calc_corr_dos == true
     *
     * @note
     *   1. propagator.act() changes its working arrays, so it is & rather than const &.
     *   2. DO NOT use auto for eigen datatypes (vectors in wfq).
     */
    void tbpm_qe_dos(
        Propagator<sparse_t>& propagator,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        const bool& calc_corr_dos,
        const Eigen::VectorXcd& wf0,
        std::vector<Eigen::VectorXcd>& wfq,
        Eigen::VectorXcd& corr_dos)
    {
        // Initialize wave functions
        size_t n_wf = wf0.size();
        Eigen::VectorXcd wf_t_pos(n_wf);
        Eigen::VectorXcd wf_t_neg(n_wf);
        vec_copy(wf0, wf_t_pos);
        vec_copy(wf0, wf_t_neg);

        // Initialize quasi-eigenstates (and corr_dos)
        int num_steps = config.num_time_steps;
        size_t num_energies = meta.qe_energies.size();
        mat_col_copy(wf0, wfq); // Contribution from t=0.
        if (calc_corr_dos) {
            corr_dos.setZero();
        }

        // Initialize utilities
        ProgressBar prog_bar(num_steps, 64, "Finished DOS/QE timestep");

        // Propagation
        double norm_ref = vec_norm(wf0);
        // i_step from 1 to num_steps, it is different from DC and AC cond.
        for (int i_step = 1; i_step <= num_steps; ++i_step) {
            if (i_step == config.wfn_check_time_steps) {
                check_norm(wf_t_pos, norm_ref, config.wfn_check_threshold);
                check_norm(wf_t_neg, norm_ref, config.wfn_check_threshold);
            }
            propagator.act(wf_t_pos, true);
            propagator.act(wf_t_neg, false);

            // Hanning window
            double window = 0.5 * (1 + std::cos(base::PI * i_step / static_cast<double>(num_steps)));

            // Fourier transform
            for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                double phase = meta.qe_energies[i_e] * i_step * meta.time_step;
                complex_t factor = std::exp(base::UIMAG * phase);
                complex_t f1 = factor * window;
                complex_t f2 = std::conj(factor) * window;
                axpbypz(f1, wf_t_pos, f2, wf_t_neg, wfq[i_e]);
            }

            // DOS correlation function
            // The DOS there is different from the DOS in `tbpm_dos`, 
            // there corr_dos[0] is corr(t=time_step)
            // the DOS in `tbpm_dos` corr_dos[0] is corr(t=0)
            if (calc_corr_dos) {
                corr_dos[i_step - 1] += vec_dot(wf0, wf_t_pos);
            }
            PARA_EXEC(prog_bar.count());
        }

        // Normalization
        for (Eigen::VectorXcd& v : wfq) { // DO NOT use auto for eigen types!
            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_step, 1)
     */
    void tbpm_dos(
        const sparse_t& h_sparse,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::VectorXcd& corr) final
    {
        if (config.ldos) {
            assert(!config.ldos_orbital_indices.empty());
            size_t num_orb = h_sparse.get_dim();
            for (const auto& i : config.ldos_orbital_indices) {
                assert((num_orb > i) && (i >= 0));
            }
            PARA_EXEC(std::cout << "Calculating LDOS correlation function." << std::endl);
        } else {
            PARA_EXEC(std::cout << "Calculating DOS correlation function." << std::endl);
        }

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

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

        // Initialize random generator, propagator and progress bar
        RandomGenerator rand_gen(config.seed);
        Propagator<sparse_t> propagator(h_sparse, config, meta);
        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());

            // Randomize wf0 and copy to wf_t
            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);
            }
            vec_copy(wf0, wf_t);

            // Update correlation function
            corr[0] += 1.0; // At t = 0 wf_t == wf0, so their inner product is always 1.
            double norm_ref = 1.0; // As wf_t has been normalized for both DOS and LDOS, its norm is always 1.

            // Propagate
            for (int i_step = 1; i_step <= num_steps; ++i_step) {
                if (i_step == config.wfn_check_time_steps) {
                    check_norm(wf_t, norm_ref, config.wfn_check_threshold);
                }
                propagator.act(wf_t, true);
                corr[i_step] += vec_dot(wf0, wf_t);
                PARA_EXEC(prog_bar.count());
            }
        }

        // Average over samples
        corr /= static_cast<double>(num_samples);
    }

    /**
     * @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);

        // Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
        Eigen::VectorXcd wf0(n_wf);
        Eigen::VectorXcd psi1_x(n_wf);
        Eigen::VectorXcd psi1_y(n_wf);
        Eigen::VectorXcd psi2(n_wf);

        // 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.
        Propagator<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); // j_beta|psi> for beta = x
            curr_y.mv(wf0, psi1_y); // j_beta|psi> for beta = y
            fermi_dirac.act(psi1_x, true); // [1-f(H)]j_beta|psi> for beta = x
            fermi_dirac.act(psi1_y, true); // [1-f(H)]j_beta|psi> for beta = y
            fermi_dirac.act_out(wf0, psi2, false); // f(H)|psi>

            // Get initial correlation function in all directions
            corr(0, 0) += curr_x.xmy(psi2, psi1_x); // <psi|f(H)j_alpha[1-f(H)]j_beta|psi> for xx
            corr(1, 0) += curr_y.xmy(psi2, psi1_x); // yx
            corr(2, 0) += curr_x.xmy(psi2, psi1_y); // xy
            corr(3, 0) += curr_y.xmy(psi2, psi1_y); // 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);
            // NOTE: i_step is from 2 to  num_steps, i.e. the maximum of num_step
            // is num_steps - 1 NOT num_steps. This is different from DOS.
            for (int i_step = 2; i_step <= num_steps; ++i_step) {
                if (i_step == config.wfn_check_time_steps) {
                    check_norm(psi2, norm_ref, config.wfn_check_threshold);
                }
                propagator.act(psi1_x, true);
                propagator.act(psi1_y, true);
                propagator.act(psi2, true);

                // Get correlation function in all directions
                // Formula exactly the same as initial values
                int j = i_step - 1;
                corr(0, j) += curr_x.xmy(psi2, psi1_x);
                corr(1, j) += curr_y.xmy(psi2, psi1_x);
                corr(2, j) += curr_x.xmy(psi2, psi1_y);
                corr(3, j) += curr_y.xmy(psi2, psi1_y);
                PARA_EXEC(prog_bar.count());
            }
        }

        // 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);

        // Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
        Eigen::VectorXcd wf0(n_wf);
        Eigen::VectorXcd psi1(n_wf);
        Eigen::VectorXcd psi2(n_wf);

        // 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.
        Propagator<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); // psi1 = rho(-q) * psi0
            fermi_dirac.act(psi1, true); // psi1 = (1-nF) * psi1
            fermi_dirac.act_out(wf0, psi2, false); // psi2 = nF * psi0

            // Initial value of correlation function
            // See eqn.77 of CPC paper.
            corr[0] += xmy_diag(psi2, den_coeff_q, psi1).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);
            for (int i_step = 2; i_step <= num_steps; ++i_step) {
                if (i_step == config.wfn_check_time_steps) {
                    check_norm(psi2, norm_ref, config.wfn_check_threshold);
                }
                propagator.act(psi1, true);
                propagator.act(psi2, true);
                corr[i_step - 1] += xmy_diag(psi2, den_coeff_q, psi1).imag();
                PARA_EXEC(prog_bar.count());
            }
        }

        // Average over samples
        corr /= static_cast<double>(num_samples);
    }

    /**
     * @brief Calculate squared modulus of quasi-eigenstates for given energies.
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation configurations
     * @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 wave functions
        size_t n_wf = h_sparse.get_dim();
        Eigen::VectorXcd wf0(n_wf);

        // Initialize quasi-eigenstates
        size_t num_energies = meta.qe_energies.size();
        wf2_qe = Eigen::MatrixXd::Zero(n_wf, num_energies);
        std::vector<Eigen::VectorXcd> wfq;
        mat_alloc(wfq, n_wf, num_energies);
        Eigen::VectorXcd corr_dos; // just for filling arguments of tbpm_qe_core

        // Initialize operators
        Propagator<sparse_t> propagator(h_sparse, config, meta);

        // Initialize utilities
        RandomGenerator rand_gen(config.seed);

        int num_samples = meta.num_samples_dist;
        int rank = meta.mpi_rank;
        double avg_factor = 1.0 / static_cast<double>(num_samples);
        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 quasi-eigenstates for this initial condition
            rand_gen.random_state(wf0, seed_i);
            tbpm_qe_dos(propagator, config, meta, false, wf0,
                wfq, corr_dos);

            // Evaluate abs2 of qe and average over samples
            for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                wf2_qe.col(i_e) += wfq[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)
     */
    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);

        // Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
        Eigen::VectorXcd wf0(n_wf);
        Eigen::VectorXcd psi0_x(n_wf); // j_x|psi0> in eqn.62 of CPC paper
        Eigen::VectorXcd psi0_y(n_wf); // j_y|psi0>

        // 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(num_steps_dos);
        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
        Propagator<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 quasi-eigenstates and corr_dos for this initial condition
            rand_gen.random_state(wf0, seed_i);
            PARA_EXEC(std::cout << "Calculating corr_dos and qe." << std::endl);
            tbpm_qe_dos(propagator, config, meta, true, wf0,
                wf_qe, corr_dos_work);
            vec_self_add_vec(corr_dos, corr_dos_work);

            PARA_EXEC(prog_bar.reset());
            // 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]));
            }

            // 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]);
            }

            // 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, psi2_x[i_e]) / factor[i_e];
                corr_dc_yy(i_e, 0) += vec_dot(psi0_y, psi2_y[i_e]) / factor[i_e];
            }
            
            PARA_EXEC(prog_bar.count());

            // Propagation
            // NOTE: i_step is from 2 to  num_steps_dc, i.e. the maximum of num_step
            // is num_steps_dc - 1 NOT num_steps_dc. This is different from DOS.
            for (int i_step = 2; i_step <= num_steps_dc; ++i_step) {
                propagator.act(psi0_x, true); // exp(-iHt)j_x|psi0>
                propagator.act(psi0_y, true); // exp(-iHt)j_y|psi0>
                int j = i_step - 1;
                for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                    corr_dc_xx(i_e, j) += vec_dot(psi0_x, psi2_x[i_e]) / factor[i_e];
                    corr_dc_yy(i_e, j) += vec_dot(psi0_y, psi2_y[i_e]) / factor[i_e];
                }
                PARA_EXEC(prog_bar.count());
            }
        }

        // 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();
        // 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));
        wft = Eigen::MatrixXcd::Zero(n_wf, num_save);

        // Initialize propagator and progress bar
        int num_steps = config.num_time_steps;
        Propagator<sparse_t> propagator(h_sparse, config, meta);
        ProgressBar prog_bar(num_steps, 64, "Finished timestep");

        // Check if we need to save the initial state
        size_t counter = 0;
        if (config.wft_time_save.find(0) != config.wft_time_save.end()) {
            wft.col(counter) = wf0;
            counter++;
        }

        // Propagate
        double norm_ref = vec_norm(wf0);
        for (int i_step = 1; i_step <= num_steps; ++i_step) {
            if (i_step == config.wfn_check_time_steps) {
                check_norm(wf0, norm_ref, config.wfn_check_threshold);
            }
            propagator.act(wf0);
            PARA_EXEC(prog_bar.count());
            if (config.wft_time_save.find(i_step) != config.wft_time_save.end()) {
                wft.col(counter) = wf0;
                counter++;
            }
        }
    }

    /**
     * @brief Calculate charge density rho(r) via Single-step 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);


        // 1. Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
       
        Eigen::VectorXcd psi(n_wf);           // Initial random state |psi>
        Eigen::VectorXcd wf_t_pos(n_wf);      // |psi^CD(t)>            
        Eigen::VectorXcd wf_t_neg(n_wf);      // |psi^CD(-t)>

        // 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);
        
        Propagator<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;

        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);
            vec_copy(wf_t_pos, wf_t_neg); 

            // 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,buffer);

            // Store the initial norm if you plan to do norm checks during evolution
            double norm_ref = tbplas::tbpm::vec_norm(wf_t_pos);
            // std::cout<<"pos and neg: 0 "<<wf_t_pos[100]<<"  "<<wf_t_neg[100]<<std::endl;
            // d. Loop for time evolution
            // For j <- 1 to N_Dt
            for (int j = 1; j <= config.num_time_steps; ++j) {
                if (j == config.wfn_check_time_steps) {
                    tbplas::tbpm::check_norm(wf_t_pos, norm_ref, config.wfn_check_threshold);
                }
                // t_j <- j * Delta_t is handled internally by propagator
                // Evolve |psi^CD(t_{j-1})> -> |psi^CD(t_j)>
                propagator.act(wf_t_pos, true); 
                // Evolve |psi^CD(-t_{j-1})> -> |psi^CD(-t_j)>
                propagator.act(wf_t_neg, false); 
                // std::cout<<j<<" pos and neg: 0 "<<wf_t_pos[100]<<"  "<<wf_t_neg[100]<<std::endl;
                // e. Accumulate charge density for current time step
                // rho(r) <- rho(r) + |<r|psi^CD(t_j)>|^2 + |<r|psi^CD(-t_j)>|^2
                vec_abs2_add(wf_t_pos,buffer);
                vec_abs2_add(wf_t_neg,buffer);
                PARA_EXEC(prog_bar.count());
            }
            vec_self_div_sca(buffer, 2*config.num_time_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