/**
 * @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"

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

/**
 * @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 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)
{
    int table_size = k_max + 1;
    bes_table = std::vector<Eigen::VectorXcd>(table_size);
    double bessel_precision = config.bessel_precision;

    for (int k = 0; k <= k_max; ++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);
        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;
    }
}

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

        // 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 * std::pow(-UIMAG, n) * bes_vec[n];
            bes_backward_[n] = 2.0 * std::pow(UIMAG, 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, bes_forward_);
    }

    /**
     * @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
     */
    void act(std::vector<Eigen::VectorXcd>& wf_out)
    {
        cheb_eval_wf_multi(h_sparse_, bes_forward_, 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_;

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

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

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

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

    /**
     * @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);
        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) {
                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);
        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 * std::pow(-UIMAG, n);
            complex_t factor_neg = 2.0 * std::pow(UIMAG, 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 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