/**
 * @brief Functions and classes for propagation.
 *
 * @author Zewen Wu, Miao Zhang and Yunhai Li
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 *
 * @note
 *   We make the following naming convention for variables, especially loop counters:
 *   1. n, m: Bessel series for given time t
 *   2. k, i_step: time step for given Bessel series order
 *   3. i_e: energy index of quasi-eigenstates
 */

#ifndef TBPLAS_TBPM_PROPAGATION_H
#define TBPLAS_TBPM_PROPAGATION_H

#include <cstdlib>
#include <iostream>
#include <list>
#include <vector>

#include <bessel-library/bessel-library.hpp>

#include <eigen3/Eigen/Dense>

#include "base/consts.h"
#include "chebyshev.h"
#include "config.h"
#include "dense.h"
#include "base/io_text.h"

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


/**
 * @brief Calculates Bessel functions of the first kind, J_n(x), for n = 0, ..., n_trunc.
 *
 * Uses the downward recurrence relation. This method is stable for large n or x.
 * Calculates up to a specified max_order internally, then truncates based on epsilon.
 *
 * @param x The argument for the Bessel function.
 * @param max_order The maximum order n to calculate up to internally before truncation.
 * @param epsilon A precision threshold used for rescaling and truncation.
 * @param[out] Jn A std::vector<double> that will be resized and filled with J_0(x), ..., J_n_trunc(x).
 */
void calculate_bessel_j_series(double x, int max_order, double epsilon, std::vector<double>& Jn)
{
    // Handle the special case where x is close to zero
    // J_0(0) = 1, J_n(0) = 0 for n > 0
    if (std::abs(x) < 1e-10)
    {
        Jn.resize(max_order + 1);
        Jn[0] = 1.0;
        for (int i = 1; i <= max_order; ++i) {
            Jn[i] = 0.0;
        }
        return; 
    }

    // Resize the vector to hold up to max_order terms (indices 0 to max_order)
    Jn.resize(max_order + 1);

    // Initialize the highest two orders for the downward recurrence
    // The exact values don't matter initially as the series will be normalized later.
    Jn[max_order] = 0.0;
    Jn[max_order - 1] = 1.0; // Or a small non-zero value like 1.0e-100

    double abs_x = std::abs(x);

    // Use the downward recurrence relation: J_{n-1} = (2n/|x|) * J_n - J_{n+1}
    // Starting from J_{max_order-2} down to J_0
    for (int i = max_order - 2; i >= 0; --i)
    {
        double temp = 2.0 * (i + 1) / abs_x * Jn[i + 1] - Jn[i + 2];
        Jn[i] = temp;

        // Rescale to avoid potential overflow, based on the Fortran logic
        // If the current term multiplied by epsilon is large, scale down the entire vector.
        if (std::abs(temp) * epsilon > 1.0)
        {
            for (int j = 0; j <= max_order; ++j) {
                Jn[j] *= epsilon;
            }
        }
    }

    // Normalize the Bessel functions using the identity: J_0(x)^2 + 2 * sum(J_n(x)^2 for n=1 to max_order) approx 1
    // This step scales the arbitrary values from the downward recurrence to the correct magnitudes.
    double sum_term = Jn[0] * Jn[0];
    for (int i = 1; i <= max_order; ++i) {
        sum_term += 2.0 * Jn[i] * Jn[i];
    }
    sum_term = std::sqrt(sum_term);


    // Apply normalization and the sign correction for negative x based on J_n(-x) = (-1)^n J_n(x)
    double sign_x = (x < 0.0) ? -1.0 : 1.0;
    for (int i = 0; i <= max_order; ++i) {
        Jn[i] /= sum_term;
        if (sign_x < 0.0 && (i % 2 != 0)) { // Check if x is negative AND order i is odd
            Jn[i] = -Jn[i];
        }
    }

    // This check verifies if the "tail" of the series (at the calculated max_order)
    // has decayed below the precision threshold.
    if (std::abs(Jn[max_order]) >= epsilon) {
        std::cout << "ERROR: Bessel series calculation did not converge sufficiently within max_order."
                  << " |J_" << max_order << "(" << x << ")| = " << std::abs(Jn[max_order])
                  << " >= epsilon = " << epsilon << ". Consider increasing max_order." << std::endl;
        std::exit(-1); 
   }
}

/**
 * @brief Build Bessel vector J_n(t) using the downward recurrence relation.
 *
 * Calculates J_n(time_step) for n = 0, 1, 2, ..., n_trunc, where n_trunc is
 * determined by the precision. Uses a custom implementation based on downward
 * recurrence which is stable for large arguments or orders.
 *
 * @param[in] time_step time step t in J_n(t)
 * @param[in] precision cutoff for truncating J_n at n_trunc. Terms |J_n| <= precision
 * and |J_{n+1}| <= precision indicate truncation point.
 * @param[out] bes_vec Bessel vector for n = 0, 1, 2, ..., n_trunc (resized internally).
 *
 * @note
 * - The downward recurrence method requires calculating up to a sufficiently high
 * max_order. The heuristic used here (`abs(time_step) + buffer`) determines this.
 * The calculate_bessel_j_series function includes an internal check on J_max_order.
 * - This function calls calculate_bessel_j_series to get the full series up to max_order
 * and then performs the truncation based on the precision criterion.
 * - The search for the truncation point starts from int(abs(time_step)),
 * leveraging the property that J_n(x) decays for n > |x|.
 */
inline void build_bessel_vector_new(
    const double& time_step,
    const double& precision,
    Eigen::VectorXcd& bes_vec)
{

    // Determine an appropriate maximum order (M) for the downward recurrence calculation.
    int max_order = static_cast<int>(std::abs(time_step)) + 1000;

    // Use a temporary vector to store the calculated Bessel values up to max_order
    std::vector<double> calculated_Jn;

    // Calculate Bessel functions J_n(time_step) for n = 0 to max_order.
    calculate_bessel_j_series(time_step, max_order, precision, calculated_Jn);

    // Determine the truncation point (n_trunc) based on the precision criterion.
    // We collect terms until two consecutive terms (|J_n| and |J_{n+1}|) are smaller than the precision.
    int n_trunc = max_order; 
    //static_cast<int>(std::abs(time_step)) is a good starting.
    int start_n_search = static_cast<int>(std::abs(time_step));

    for (int n = start_n_search; n < max_order; ++n) // Loop from start_n_search up to max_order - 1
    {
         // The condition checks abs(Jn[n]) and abs(Jn[n+1]).
         // These indices (n and n+1) are valid within calculated_Jn because n < max_order,
         // so n+1 <= max_order.
        if (std::abs(calculated_Jn[n]) <= precision && std::abs(calculated_Jn[n+1]) <= precision)
        {
             n_trunc = n-1;
             break; // Found the truncation point
        }
    }

    // Check for non-convergence
    if (n_trunc == max_order) {
        //Protect output statements from information aliasing
        #pragma omp critical (BesselErrorOutput)
        {
            std::cout << "ERROR: Bessel series not converged\n";
            std::exit(-1);
        }
        
    }

    // Resize the output Eigen vector and populate it with the calculated real values
    // (as complex with zero imaginary part). The size is n_trunc + 1 (indices 0 to n_trunc).
    bes_vec.resize(n_trunc + 1);
    for (int n = 0; n <= n_trunc; ++n) {
        // bes_vec[n] = std::complex<double>(calculated_Jn[n], 0.0);
        bes_vec[n] = calculated_Jn[n];
    }
}


/**
 * @brief Build Bessel vector J_n(t).
 *
 * @param[in] time_step time step t in J_n(t)
 * @param[in] precision cutoff for truncating J_n at n_trunc
 * @param[out] bes_vec Bessel vector for n = 0, 1, 2, ..., n_trunc
 *
 * @note
 * 1. In principle bes_vec should be real. But in most cases it is a column of a
 *    complex matrix or vector<VectorXcd>. So we use VectorXcd instead of VectorXd.
 * 2. std::bessel_cyl is implemented for n<=128, and fails to large time step.
 *    So we have to use 3rd party implementation.
 */
inline void build_bessel_vector(
    const double& time_step,
    const double& precision,
    Eigen::VectorXcd& bes_vec)
{
    // Evaluate Bessel terms up to precision
    // The values may oscillates significantly for large time step, so we need to
    // check both n and n+1 to avoid false convergence.
    std::list<double> bes_list;
    bool converged = false;
    constexpr int n_max = 65536;
    for (int n = 0; n < n_max; ++n) {
        double bes_val = bessel::cyl_j(n, time_step);
        if (std::abs(bes_val) > precision) {
            bes_list.push_back(bes_val);
        } else {
            double bes_val_next = bessel::cyl_j(n + 1, time_step);
            if (std::abs(bes_val_next) > precision) {
                bes_list.push_back(bes_val);
            } else {
                converged = true;
                break;
            }
        }
    }

    // Check for non-convergence
    if (!converged) {
        std::cout << "ERROR: Bessel series not converged\n";
        std::exit(-1);
    }

    // Convert to Eigen vector
    bes_vec = Eigen::VectorXcd(bes_list.size());
    int ptr = 0;
    for (const auto& f : bes_list) {
        bes_vec[ptr] = f;
        ++ptr;
    }
}


/**
 * @brief Efficiently calculates (-i)^n, where i is the imaginary unit.
 *
 * This function computes powers of negative i (i.e., std::complex(0.0, -1.0)),
 * which follow a cycle of four values. It uses a precomputed lookup table
 * for optimal performance, avoiding expensive std::pow calls.
 *
 * @param[in] n The exponent. While the type is int, this function is used
 * in contexts where 'n' is guaranteed to be non-negative
 * (e.g., loop indices).
 * @return std::complex<double> The result of (-i)^n.
 *
 * @note The direct indexing `n % 4` relies on the precondition that 'n'
 * is non-negative. If 'n' could be negative, the behavior of '%'
 * for negative dividends would need to be handled to ensure a valid
 * array index [0-3].
 */
inline std::complex<double> power_of_negative_i(int n) {
    // Precomputed values for (-i)^0, (-i)^1, (-i)^2, (-i)^3.
    // These are mathematical constants.
    static const std::complex<double> powers[] = {
        {1.0, 0.0},  // (-i)^0 =  1
        {0.0, -1.0}, // (-i)^1 = -i
        {-1.0, 0.0}, // (-i)^2 = -1
        {0.0, 1.0}   // (-i)^3 =  i
    };
    // For non-negative integers n, n % 4 correctly yields an index in [0, 3].
    return powers[n % 4];
}

/**
 * @brief Efficiently calculates (i)^n, where i is the imaginary unit.
 *
 * This function computes powers of positive i, which follow a cycle of four.
 * It uses a precomputed lookup table for high performance.
 *
 * @param[in] n The exponent. Assumed to be non-negative in the context
 * of its use (e.g., loop indices).
 * @return std::complex<double> The result of (i)^n.
 *
 * @note Relies on 'n' being non-negative for direct `n % 4` indexing.
 */
inline std::complex<double> power_of_positive_i(int n) { // Or size_t if n is always unsigned
    // Precomputed values for i^0, i^1, i^2, i^3.
    static const std::complex<double> powers[] = {
        {1.0, 0.0},  // i^0 =  1
        {0.0, 1.0},  // i^1 =  i
        {-1.0, 0.0}, // i^2 = -1
        {0.0, -1.0}  // i^3 = -i
    };

    return powers[n % 4];
}


/**
 * @brief Build Bessel table bes(n, k) = (-i)^n * J_n(k*dt) for
 * k = 0, 1, 2, ..., k_max and n = 0, 1, 2, ..., num_bes-1
 *
 * @param[in] config user-defined calculation configurations
 * @param[in] meta additional meta info.
 * @param[in] k_max upper bound of k
 * @param[out] bes_table Bessel table, dimension (num_bes, k_max+1)
 *
 * @note Separating definition in cpp file causes compilation errors in bessel-library.
 * So we have to put it in the header with the inline keyword.
 */
inline void build_bessel_table(
    const TBPMConfig& config,
    const TBPMMeta& meta,
    const int& k_max,
    std::vector<Eigen::VectorXcd>& bes_table,
    bool endpoint=true)
{
    int table_size;
    if(endpoint){
        table_size = k_max + 1;
    }else{
        table_size = k_max;
    }
    
    bes_table = std::vector<Eigen::VectorXcd>(table_size);
    double bessel_precision = config.bessel_precision;

    //Moving OMP parallelism to different-time parallelism allows 
    //time-consuming build_bessel_vector to run simultaneously in different threads.
    #pragma omp parallel for default(none) shared(meta, bes_table, table_size, bessel_precision)
    for (int k = 0; k <= table_size-1; ++k) {
        double time_step_scaled = meta.time_step * meta.rescale * k;
        Eigen::VectorXcd& bes_vec = bes_table[k];
        // build_bessel_vector(time_step_scaled, bessel_precision, bes_vec);
        build_bessel_vector_new(time_step_scaled, 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 * power_of_negative_i(n));
        }
        bes_vec[0] *= 0.5;
    }
}

/**
 * @brief Collections of sizes and pointers of arrays in propagator,
 * for GPU programming.
 *
 */
class PropagatorMeta {
public:
    size_t bes_forward_size = 0;
    size_t bes_backward_size = 0;
    size_t t0_size = 0;
    size_t t1_size = 0;
    size_t wf_out_size = 0;
    const complex_t* bes_forward_data = nullptr;
    const complex_t* bes_backward_data = nullptr;

    // Buffers will be modified by GPU, so there are no const qualifiers.
    complex_t* t0_data = nullptr;
    complex_t* t1_data = nullptr;
    complex_t* wf_out_data = nullptr;
};

/**
 * @brief Propagator of time-dependent wave function based on Chebyshev expansion.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 */
template <typename sparse_t>
class Propagator {
public:
    /**
     * @brief Construct a new Propagator object
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     */
    explicit Propagator(
        const sparse_t& h_sparse,
        const TBPMConfig& config,
        const TBPMMeta& meta)
        : h_sparse_(&h_sparse)
        , t0_(Eigen::VectorXcd(h_sparse.get_dim()))
        , t1_(Eigen::VectorXcd(h_sparse.get_dim()))
        , wf_out_(Eigen::VectorXcd(h_sparse.get_dim()))
    {
        build_bessel_series(config, meta); // initialize bes_forward_ and bes_backward_
    }

    /**
     * @brief Act on the wave function to propagate it
     *
     * @param[inout] wf_in incoming wave function to propagate
     * @param[in] forward whether to propagate forward
     */
    void act(Eigen::VectorXcd& wf_in, bool forward = true)
    {
        const Eigen::VectorXcd* bes = forward ? &bes_forward_ : &bes_backward_;
        cheb_eval_wf(h_sparse_, *bes, wf_in, t0_, t1_, wf_out_);
        vec_copy(wf_out_, wf_in);
    }

    /**
     * @brief Get meta with sizes and pointers of array attributes.
     *
     * @return Meta info containing the sizes and pointers
     */
    [[nodiscard]] PropagatorMeta get_meta() const
    {
        PropagatorMeta meta;
        meta.bes_forward_size = bes_forward_.size();
        meta.bes_backward_size = bes_backward_.size();
        meta.t0_size = t0_.size();
        meta.t1_size = t1_.size();
        meta.wf_out_size = wf_out_.size();
        meta.bes_forward_data = bes_forward_.data();
        meta.bes_backward_data = bes_backward_.data();
        meta.t0_data = const_cast<complex_t*>(t0_.data());
        meta.t1_data = const_cast<complex_t*>(t1_.data());
        meta.wf_out_data = const_cast<complex_t*>(wf_out_.data());
        return meta;
    }

private:
    /// @brief pointer to scaled sparse Hamiltonian
    const sparse_t* h_sparse_ = nullptr;

    /// @brief Bessel series
    Eigen::VectorXcd bes_forward_;
    Eigen::VectorXcd bes_backward_;

    /// @brief caches for updating wave function
    Eigen::VectorXcd t0_;
    Eigen::VectorXcd t1_;
    Eigen::VectorXcd wf_out_;

    /**
     * @brief Set up the 'bes_*' attribute
     *
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     */
    void build_bessel_series(const TBPMConfig& config, const TBPMMeta& meta)
    {
        double bessel_precision = config.bessel_precision;
        double time_step_scaled = meta.time_step * meta.rescale;
        Eigen::VectorXcd bes_vec;
        // build_bessel_vector(time_step_scaled, bessel_precision, bes_vec);
        build_bessel_vector_new(time_step_scaled, bessel_precision, bes_vec);

        // Populate bessel series
        // Note: we also put the factor of 2 into bessel series!
        size_t num_bes = bes_vec.size();
        bes_forward_ = Eigen::VectorXcd(num_bes);
        bes_backward_ = Eigen::VectorXcd(num_bes);
#pragma omp parallel for simd
        for (int n = 0; n < num_bes; ++n) {
            bes_forward_[n] = 2.0 * power_of_negative_i(n) * bes_vec[n];
            bes_backward_[n] = 2.0 * power_of_positive_i(n) * bes_vec[n];
        }
        // First coefficient has a factor of 1 while others have 2.
        bes_forward_[0] *= 0.5;
        bes_backward_[0] *= 0.5;
    }
};

/**
 * @brief Similar as Propagator, but propagates the wave function by multiple time steps.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 */
template <typename sparse_t>
class MultiPropagator {
public:
    /**
     * @brief Construct a new Multi Propagator object
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     */
    explicit MultiPropagator(const sparse_t& h_sparse, const TBPMConfig& config, const TBPMMeta& meta)
        : h_sparse_(&h_sparse)
        , t0_(Eigen::VectorXcd(h_sparse.get_dim()))
        , t1_(Eigen::VectorXcd(h_sparse.get_dim()))
    {
        build_bessel_table_(config, meta, config.pool_size);
    }

    /**
     * @brief Evaluate |psi(t)> for a series of time steps from initial state in wf_out[0]
     *
     * @param[inout] wf_out |psi(t)> for t = 0, 1, 2, ..., pool_size
     * @param[in] forward whether to propagate forward
     */
    void act(std::vector<Eigen::VectorXcd>& wf_out, bool forward = true)
    {
        if(forward){
            cheb_eval_wf_multi(h_sparse_, bes_forward_, t0_, t1_, wf_out);
        } else{
            cheb_eval_wf_multi(h_sparse_, bes_backward_, t0_, t1_, wf_out);
        }
        
    }

private:
    /// @brief pointer to scaled sparse Hamiltonian
    const sparse_t* h_sparse_ = nullptr;

    /// @brief Bessel series
    std::vector<Eigen::VectorXcd> bes_forward_;
    std::vector<Eigen::VectorXcd> bes_backward_;

    /// @brief caches for updating wave function
    Eigen::VectorXcd t0_;
    Eigen::VectorXcd t1_;

    /**
     * @brief Set up the 'bes_*' attribute.
     * 
     * Build Bessel table bes(n, k) = (-i)^n * J_n(k*dt) for
     * k = 0, 1, 2, ..., k_max and n = 0, 1, 2, ..., num_bes-1
     *
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     * @param[in] k_max upper bound of k
     */
    void build_bessel_table_(
        const TBPMConfig& config,
        const TBPMMeta& meta,
        const int& k_max)
    {
        int table_size = k_max + 1;
        // bes_table = std::vector<Eigen::VectorXcd>(table_size);
        bes_backward_ = std::vector<Eigen::VectorXcd>(table_size);
        bes_forward_ = std::vector<Eigen::VectorXcd>(table_size);
        double bessel_precision = config.bessel_precision;

        //Moving OMP parallelism to different-time parallelism allows 
        //time-consuming build_bessel_vector to run simultaneously in different threads.
        #pragma omp parallel for default(none) shared(meta, bes_forward_,bes_backward_, k_max, bessel_precision)
        for (int k = 0; k <= k_max; ++k) {
            double time_step_scaled = meta.time_step * meta.rescale * k;
            Eigen::VectorXcd bes_vec;
            // Eigen::VectorXcd& bes_vec = bes_table[k];
            // build_bessel_vector(time_step_scaled, bessel_precision, bes_vec);
            build_bessel_vector_new(time_step_scaled, bessel_precision, bes_vec);
            size_t num_bes = bes_vec.size();
            bes_forward_[k] = Eigen::VectorXcd(num_bes);
            bes_backward_[k] = Eigen::VectorXcd(num_bes);
    // #pragma omp parallel for simd
            for (int n = 0; n < num_bes; ++n) {
                bes_forward_[k][n] = 2.0 * power_of_negative_i(n) * bes_vec[n];
                bes_backward_[k][n] = 2.0 * power_of_positive_i(n) * bes_vec[n];
            }
            // First coefficient has a factor of 1 while others have 2.
            bes_forward_[k][0] *= 0.5;
            bes_backward_[k][0] *= 0.5;
        }
}
};

/**
 * @brief Correlation function calculator for DOS using fast algorithm.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 */
template <typename sparse_t>
class CDOSCalculator {
public:
    /**
     * @brief Construct a new CDOSCalculator object
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     */
    explicit CDOSCalculator(const sparse_t& h_sparse, const TBPMConfig& config, const TBPMMeta& meta)
        : h_sparse_(&h_sparse)
        , meta_(&meta)
    {
        build_bessel_table(config, meta, config.num_time_steps, bes_table_);
        prefix_ = config.prefix;
    }

    /**
     * @brief Calculate correlation function of DOS.
     *
     * @param[in] wf_in initial wave function
     * @param[out] corr correlation function determined from initial wave function
     */
    void eval(const Eigen::VectorXcd& wf_in, Eigen::VectorXcd& corr)
    {
        size_t num_bes_max = bes_table_.back().size();
        Eigen::VectorXcd cheb_mu(num_bes_max);
        cheb_eval_mu_dos(h_sparse_, meta_, wf_in, cheb_mu);

        //Save cheb_mu
        // base::TextFile data_file(prefix_, false);
        // data_file.save_matrix_complex("cheb_mu", cheb_mu);

        int table_size = bes_table_.size();
        ProgressBar prog_bar(table_size, 64, "Finished timestep");
        corr = Eigen::VectorXcd(table_size);
        for (int k = 0; k < table_size; ++k) {
            corr[k] = vec_dot_no_conj(bes_table_[k], cheb_mu, false);
            if (meta_->is_master) {
                prog_bar.count();
            }
        }
    }

private:
    /// @brief pointer to scaled sparse Hamiltonian
    const sparse_t* h_sparse_ = nullptr;

    /// @brief pointer to meta info
    const TBPMMeta* meta_ = nullptr;

    /// @brief Chebyshev moment coefficients
    std::vector<Eigen::VectorXcd> bes_table_;

    /// @brief Prefix prepended to output files
    std::string prefix_;
};


/**
 * @brief Correlation function calculator for DC conductivity using fast algorithm.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 */
template <typename sparse_t>
class CDCCalculator {
public:
    /**
     * @brief Construct a new CDCCalculator object
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     */
    explicit CDCCalculator(const sparse_t& h_sparse, const TBPMConfig& config, const TBPMMeta& meta)
        : h_sparse_(&h_sparse)
        , meta_(&meta)
    {
        build_bessel_table(config, meta, config.dc_cond_num_time_steps, bes_table_,false);
    }

    /**
     * @brief Calculate correlation function of DC.
     *
     * @param[in] psi1_x: j_x|psi0>
     * @param[in] psi1_y: j_y|psi0>
     * @param[in] factor:  |<psi0|psi(E)>|
     * @param[in] psi2_x: j_x|psi(E)>
     * @param[in] psi2_y: j_y|psi(E)>
     * @param[out] corr_dc_xx corr correlation function in direction x
     * @param[out] corr_dc_yy corr correlation function in direction y
     */
    void eval(const Eigen::VectorXcd& psi1_x, const Eigen::VectorXcd& psi1_y,
        const Eigen::VectorXd& factor,
        const std::vector<Eigen::VectorXcd>& psi2_x,
        const std::vector<Eigen::VectorXcd>& psi2_y,
        Eigen::MatrixXcd& corr_dc_xx, Eigen::MatrixXcd& corr_dc_yy)
    {
        size_t num_bes_max = bes_table_.back().size();
        size_t num_energies = psi2_x.size();
        int table_size = bes_table_.size();
        std::vector<Eigen::VectorXcd> cheb_mu;

        mat_alloc(cheb_mu, num_bes_max, num_energies);
        mat_init(cheb_mu, 0.);
        cheb_eval_mu_dc(h_sparse_, meta_, psi1_x, psi2_x, cheb_mu);

        ProgressBar prog_bar(table_size, 64, "Finished timestep dc_xx");
        for (int k = 0; k < table_size; ++k) {
            for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                corr_dc_xx(i_e, k) = vec_dot(bes_table_[k], cheb_mu[i_e], false)/ factor[i_e];
            }
            if (meta_->is_master) {
                prog_bar.count();
            }
        }

        mat_init(cheb_mu, 0.);
        cheb_eval_mu_dc(h_sparse_, meta_, psi1_y, psi2_y, cheb_mu);

        ProgressBar prog_bar1(table_size, 64, "Finished timestep dc_yy");
        for (int k = 0; k < table_size; ++k) {
            for (size_t i_e = 0; i_e < num_energies; ++i_e) {
                corr_dc_yy(i_e, k) = vec_dot(bes_table_[k], cheb_mu[i_e], false)/ factor[i_e];
            }
            if (meta_->is_master) {
                prog_bar1.count();
            }
        }
    }

private:
    /// @brief pointer to scaled sparse Hamiltonian
    const sparse_t* h_sparse_ = nullptr;

    /// @brief pointer to meta info
    const TBPMMeta* meta_ = nullptr;

    /// @brief Chebyshev moment coefficients
    std::vector<Eigen::VectorXcd> bes_table_;
};

/**
 * @brief Quasi-eigenstates calculator using fast algorithm.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 */
template <typename sparse_t>
class QECalculator {
public:
    /**
     * @brief Construct a new QECalculator object
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     */
    explicit QECalculator(const sparse_t& h_sparse, const TBPMConfig& config, const TBPMMeta& meta)
        : h_sparse_(&h_sparse)
        , meta_(&meta)
    {
        // build_coeff(config, meta);
        build_coeff_new(config, meta);
    }

    /**
     * @brief Evaluate quasi-eigenstates from initial state.
     *
     * @param[in] wf_in initial state
     * @param[out] wf_qe quasi-eigenstates, dimension (n_wf, num_eng)
     */
    void eval(const Eigen::VectorXcd& wf_in, std::vector<Eigen::VectorXcd>& wf_qe)
    {
        // cheb_eval_qe(h_sparse_, meta_, coeff_, wf_in, wf_qe);
        cheb_eval_qe_new(h_sparse_, meta_, coeff_, wf_in, wf_qe);
        for (Eigen::VectorXcd& v : wf_qe) { // DO NOT use auto for eigen types!
            double norm = vec_norm(v);
            vec_self_div_sca(v, norm);
        }
    }

private:
    /// @brief pointer to scaled sparse Hamiltonian
    const sparse_t* h_sparse_ = nullptr;

    /// @brief pointer to meta info
    const TBPMMeta* meta_ = nullptr;

    /// @brief Quasi-eigen state coefficients
    std::vector<Eigen::VectorXcd> coeff_;

    /**
     * @brief Set up the 'coeff_' attribute.
     *
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     */
    void build_coeff(const TBPMConfig& config, const TBPMMeta& meta)
    {
        // Build exp(iEt) for different E and t
        int num_steps = config.num_time_steps;
        size_t num_eng = meta.qe_energies.size();
        std::vector<Eigen::VectorXcd> exp_pos;
        std::vector<Eigen::VectorXcd> exp_neg;
        mat_alloc(exp_pos, num_steps + 1, num_eng); // with t = 0 term
        mat_alloc(exp_neg, num_steps + 1, num_eng); // with t = 0 term

        for (size_t i_e = 0; i_e < num_eng; ++i_e) {
            double window = 0.0;
            double phase = 0.0;
            complex_t factor = 0.0;
#pragma omp parallel for private(window, phase, factor)
            for (int i_step = 0; i_step <= num_steps; ++i_step) {
                //Hanning window
                window = 0.5 * (1 + std::cos(base::PI * i_step / static_cast<double>(num_steps)));
                phase = meta.qe_energies[i_e] * i_step * meta.time_step;
                factor = std::exp(base::UIMAG * phase);
                exp_pos[i_e][i_step] = factor * window;
                exp_neg[i_e][i_step] = std::conj(factor) * window;
            }
        }

        // Get the the maximum number of Bessel series from the largest time step
        double time_step_scaled = meta.time_step * meta.rescale;
        Eigen::VectorXcd tmp;
        build_bessel_vector(time_step_scaled * num_steps, config.bessel_precision, tmp);
        // build_bessel_vector_new(time_step_scaled * num_steps, config.bessel_precision, tmp);
        size_t num_bes = tmp.size();

        // Integrate over t
        mat_alloc(coeff_, num_eng, num_bes);
        Eigen::VectorXcd bes_pos(num_steps + 1); // with t = 0 term
        Eigen::VectorXcd bes_neg(num_steps + 1); // with t = 0 term

        for (int n = 0; n < num_bes; ++n) {
            double bes_val = 0.0;
            complex_t factor_pos = 2.0 * power_of_negative_i(n);
            complex_t factor_neg = 2.0 * power_of_positive_i(n);
#pragma omp parallel for private(bes_val)
            for (int i_step = 0; i_step <= num_steps; ++i_step) {
                bes_val = bessel::cyl_j(n, time_step_scaled * i_step);
                bes_pos[i_step] = factor_pos * bes_val;
                bes_neg[i_step] = factor_neg * bes_val;
            }
            if (n == 0) { // This bug wastes me a whole day (yhli).
                vec_self_mul_sca(bes_pos, 0.5);
                vec_self_mul_sca(bes_neg, 0.5);
            }
            for (size_t i_e = 0; i_e < num_eng; ++i_e) {
                coeff_[n][i_e] = vec_dot_no_conj(exp_pos[i_e], bes_pos);
                coeff_[n][i_e] += vec_dot_no_conj(exp_neg[i_e], bes_neg);
                coeff_[n][i_e] -= exp_pos[i_e][0] * bes_pos[0]; // subtract double counted t = 0 term
            }
        }
    }

    /**
     * @brief Sets up the 'coeff_' attribute using a refactored integration approach.
     *
     * @details
     * The method calculates coefficients `coeff_[i_e][n]`, where 'i_e' is an energy index
     * and 'n' is the Bessel function order. The calculation involves a
     * sum (simulating integration) over time steps `i_step`.
     *
     * The $t=0$ (i_step=0) contribution is handled by initializing all elements of `coeff_` to 0.0 and
     * setting coeff_(:,n=0) to 1.0.
     * Subsequent time steps (t>0, i.e., i_step from 1 to num_steps) provide contributions
     * that are accumulated. Each thread accumulates into a private `coeff_` structure (initialized to zero),
     * which is then reduced (added) into the global `coeff_` at the end of the parallel region.
     *
     * The formulas  are conceptually:
     * coeff_(i_e,n): 
     * For n=0: 1.0 + sum_{t>0} J_0(t)*window(t)*(e^{iEt} + e^{-iEt}),
     *   1.0 is the $t=0$ (i_step=0) contribution
     * For n>0: sum_{t>0} 2*J_n(|t|)*window(t)*((-i)^n*e^{iEt} + (i)^n*e^{-iEt})
     *
     * @param[in] config User-defined calculation parameters
     * @param[in] meta Additional metadata
     */
    void build_coeff_new(const TBPMConfig& config, const TBPMMeta& meta)
    {
        int num_steps = config.num_time_steps;
        size_t num_eng = meta.qe_energies.size();

        // Get the maximum number of Bessel series from the largest time step
        double time_step_scaled = meta.time_step * meta.rescale;
        Eigen::VectorXcd tmp_bessel_max_arg;
        build_bessel_vector_new(time_step_scaled * num_steps,
                                config.bessel_precision, tmp_bessel_max_arg);
        size_t global_num_bes = tmp_bessel_max_arg.size();

        // Allocate the bes_table_ matrix to store all Bessel vectors for each
        // time step.
        std::vector<Eigen::VectorXcd> bes_table_;
        bes_table_ = std::vector<Eigen::VectorXcd>(num_steps+1);
        #pragma omp parallel for simd
        for (int k = 0; k <= num_steps; ++k) {
            double time_step_scaled = meta.time_step * meta.rescale * k;
            Eigen::VectorXcd& bes_vec = bes_table_[k];
            // build_bessel_vector(time_step_scaled, bessel_precision, bes_vec);
            build_bessel_vector_new(time_step_scaled, config.bessel_precision, bes_vec);
        }
        // build_bessel_table(config, meta, num_steps, bes_table_, true);

        // Allocate the coeff_ matrix to store the final coefficients
        mat_alloc(coeff_, global_num_bes, num_eng);
        // Initialize `coeff_` with the t=0 Contribution
        mat_init(coeff_, 0.); // Initialize to zero
        for (size_t i_e = 0; i_e < num_eng; ++i_e)
        {
            coeff_[i_e][0] = 1.0;
        }

        // Second Stage: Parallel computation of coeff_ based on the precomputed
        // bes_vecs
#pragma omp parallel for simd
        for (size_t i_e = 0; i_e < num_eng; ++i_e)
        {
            // Calculate phase factor (e^(i*phase))
            for (int i_step = 1; i_step <= num_steps; ++i_step)
            {
                // Calculate Hanning window for the current time step.
                double window_p =
                    0.5 *
                    (1.0 + std::cos(base::PI * static_cast<double>(i_step) /
                                    static_cast<double>(num_steps)));
                // Calculate phase and the complex exponential factor
                // e^(i*phase).
                double phase_p = meta.qe_energies[i_e] *
                                 static_cast<double>(i_step) * meta.time_step;
                complex_t factor_p_exp_iphase =
                    std::exp(base::UIMAG * phase_p); // e^(i*phase)

                // For n = 0, update the coefficient
                coeff_[i_e][0] +=
                    (std::conj(factor_p_exp_iphase) + factor_p_exp_iphase) *
                    bes_table_[i_step][0] * window_p;

                size_t num_bes_p = bes_table_[i_step].size();

                // For n > 0, update the coefficient
                for (int n = 1; n < num_bes_p; ++n)
                {
                    // Term: 2 * Jn * window * ( (-i)^n * e^(i*phase) + (i)^n *
                    // e^(-i*phase) )
                    complex_t term_component1 =
                        power_of_negative_i(n) * factor_p_exp_iphase;
                    complex_t term_component2 =
                        power_of_positive_i(n) * std::conj(factor_p_exp_iphase);
                    coeff_[i_e][n] += 2.0 * bes_table_[i_step][n] * window_p *
                                      (term_component1 + term_component2);
                }
            } // End of i_step loop
        } // End of i_e loop
    }
};

/**
 * @brief Check the norm of time-dependent wave function during propagation to detect divergence.
 *
 * @tparam vec_t datatype of wave function
 * @param[in] wf_t time-dependent wave function
 * @param[in] norm_ref reference for the norm of wave function
 * @param[in] eps criterion for detecting divergence
 *
 * @note For LDOS calculation the norm of wf_t does not equals 1.0. That's why
 * we need a parameter of norm_ref.
 */
template <typename vec_t>
void check_norm(const Eigen::MatrixBase<vec_t>& wf_t, const double& norm_ref, const double& eps)
{
    double norm = vec_norm(wf_t);
    double dnorm = std::abs(norm - norm_ref);
    if (std::isnan(dnorm) || !std::isfinite(dnorm) || dnorm > eps) {
        std::cout << "==============================================================================\n"
                  << "=                                   ERROR                                    =\n"
                  << "= Wave function norm gets broken during preparation. First of all, check the =\n"
                  << "= Hermitianity of the Hamiltonian. Then try increasing 'rescale' in input.   =\n"
                  << "==============================================================================\n"
                  << std::flush;
        std::exit(-1);
    }
}
} // namespace tbplas::tbpm
#endif