/**
 * @brief Classes and functions for Chebyshev expansion.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @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_CHEBYSHEV_H
#define TBPLAS_TBPM_CHEBYSHEV_H

#include <vector>

#include <eigen3/Eigen/Dense>

#include "base/datatypes.h"
#include "base/utils.h"
#include "config.h"
#include "dense.h"

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

/**
 * @brief Evaluate |wf_out> = \sum_{i}c_{i}*T_{i}(H)|wf_in>
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 * @tparam coeff_t datatype of coefficients
 * @param[in] h_sparse rescaled sparse Hamiltonian
 * @param[in] coeff Chebyshev expansion coefficients
 * @param[in] wf_in incoming wave function
 * @param[in] t0 buffer at step i-1, must be pre-allocated
 * @param[in] t1 buffer at step i, must be pre-allocated
 * @param[out] wf_out the result, must be pre-allocated
 *
 * @note
 * 1. Why use pointer instead of reference? It is because that this function will
 *    be called by the 'act' method of 'Propagator' and 'FermiDirac', which handle
 *    the Hamiltonian with pointers rather than references. For consistence, we also
 *    use pointers here.
 * 2. Why pass t0 and t1 as arguments rather than making them local variables?
 *    It is because that this function will be called thousands of time, and
 *    re-allocating memory is slow.
 */
template <typename sparse_t, typename coeff_t>
void cheb_eval_wf(
    const sparse_t* h_sparse,
    const Eigen::MatrixBase<coeff_t>& coeff,
    const Eigen::VectorXcd& wf_in,
    Eigen::VectorXcd& t0,
    Eigen::VectorXcd& t1,
    Eigen::VectorXcd& wf_out)
{
    size_t dim = h_sparse->get_dim();
    assert(dim == wf_in.size());
    assert(dim == t0.size());
    assert(dim == t1.size());
    assert(dim == wf_out.size());

    // Initialize pointers
    Eigen::VectorXcd* p0 = &t0;
    Eigen::VectorXcd* p1 = &t1;
    Eigen::VectorXcd* p2 = nullptr;

    // Initial condition
    // t0 = H * wf_in
    // t1 = 2 * H * t0 - wf_in
    // wf_out = coeff[0] * wf_in + coeff[1] * t0 + coeff[2] * t1
    h_sparse->mv(wf_in, t0);
    h_sparse->amxsy_out(2.0, t0, wf_in, t1);
    axpbypcz_out(coeff[0], wf_in, coeff[1], t0, coeff[2], t1, wf_out);

    // Continue the series expansion
    // p2 = 2 * H * p1 - p0
    // wf_out += coeff[i] * p2
    size_t num_series = coeff.size();
    for (size_t n = 3; n < num_series; ++n) {
        p2 = p0;
        h_sparse->amxsy(2.0, *p1, *p0);
        axpy(coeff[n], *p2, wf_out);
        p0 = p1;
        p1 = p2;
    }
}

/**
 * @brief Evaluate |psi(t)> for a series of time steps from initial state in wf_out[0]
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 * @param[in] h_sparse rescaled sparse Hamiltonian
 * @param[in] coeff Bessel table bes(n, k) for k = 0, 1, 2, ..., pool_size
 * @param[in] t0 buffer at step i-1, must be pre-allocated
 * @param[in] t1 buffer at step i, must be pre-allocated
 * @param[inout] wf_out |psi(t)> for t = 0, 1, 2, ..., pool_size, must be pre-allocated
 *
 * @note
 *   1. DO NOT use std::vector<Eigen::MatrixBase<coeff_t>> for declaring coeff. It will
 *      not compile.
 *   2. When called by 'tbpm_wft_fast', k in coeff can be discrete, e.g., 0, 1, 5, 17, 128.
 */
template <typename sparse_t>
void cheb_eval_wf_multi(
    const sparse_t* h_sparse,
    const std::vector<Eigen::VectorXcd>& coeff,
    Eigen::VectorXcd& t0,
    Eigen::VectorXcd& t1,
    std::vector<Eigen::VectorXcd>& wf_out)
{
    size_t dim = h_sparse->get_dim();
    assert(dim == t0.size());
    assert(dim == t1.size());
    assert(dim == wf_out.back().size());
    assert(coeff.size() == wf_out.size());

    // Initialize pointers
    Eigen::VectorXcd* p0 = &t0;
    Eigen::VectorXcd* p1 = &t1;
    Eigen::VectorXcd* p2 = nullptr;

    // Initial condition
    const Eigen::VectorXcd& wf_in = wf_out[0];
    h_sparse->mv(wf_in, t0);
    h_sparse->amxsy_out(2.0, t0, wf_in, t1);
    size_t pool_size = coeff.size() - 1;
    for (int k = 1; k <= pool_size; ++k) {
        axpbypcz_out(coeff[k][0], wf_in, coeff[k][1], t0, coeff[k][2], t1, wf_out[k]);
    }

    // Continue the series expansion
    size_t num_series = coeff.back().size();
    for (size_t n = 3; n < num_series; ++n) {
        p2 = p0;
        h_sparse->amxsy(2.0, *p1, *p0);
        for (int k = 1; k <= pool_size; ++k) {
            if (n < coeff[k].size()) {
                axpy(coeff[k][n], *p2, wf_out[k]);
            }
        }
        p0 = p1;
        p1 = p2;
    }
}

/**
 * @brief Evaluate Chebyshev moments <wf_in|Tn|wf_in>.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 * @param[in] h_sparse rescaled sparse Hamiltonian
 * @param[in] meta meta info
 * @param[in] wf_in initial wave function
 * @param[out] cheb_mu Chebyshev moments for n = 0, 1, 2, ..., num_mu, must be pre-allocated
 *
 * @note See chebyshev.py for the deduction formulae.
 */
template <typename sparse_t>
void cheb_eval_mu_dos(
    const sparse_t* h_sparse,
    const TBPMMeta* meta,
    const Eigen::VectorXcd& wf_in,
    Eigen::VectorXcd& cheb_mu)
{
    // Initialize results and buffers
    size_t dim = h_sparse->get_dim();
    assert(dim == wf_in.size());
    Eigen::VectorXcd t0(dim);
    Eigen::VectorXcd t1(dim);

    // Initialize pointers
    Eigen::VectorXcd* p0 = &t0;
    Eigen::VectorXcd* p1 = &t1;
    Eigen::VectorXcd* p2 = nullptr;

    // Initial condition
    // t0 = H * wf_in
    // t1 = 2 * H * t0 - wf_in
    h_sparse->mv(wf_in, t0);
    h_sparse->amxsy_out(2.0, t0, wf_in, t1);
    complex_t mu0 = vec_norm2(wf_in);
    complex_t mu1 = vec_dot(wf_in, t0);
    complex_t mu2 = vec_dot(wf_in, t1);
    cheb_mu[0] = mu0;
    cheb_mu[1] = mu1;
    cheb_mu[2] = mu2;

    // CAUTION: Avoid dividing odd number by 2 since the result is unpredictable!
    int num_mu = cheb_mu.size() - 1;
    int num_mu_half = 0;
    if (num_mu % 2 == 0) {
        num_mu_half = num_mu / 2;
    } else {
        num_mu_half = (num_mu + 1) / 2;
    }
    ProgressBar prog_bar(num_mu_half, 64, "Finished mu order");
    if (meta->is_master) {
        prog_bar.count(); // mu1
        prog_bar.count(); // mu2
    }

    // Continue the deduction
    // p2 = 2 * H * p1 - p0
    // mu[i] = <wf_in|p2>
    // mu[2*i-1] = 2 * <p2|p1> - mu[1]
    // mu[2*i] = 2 * <p2|p2> - mu[0]
    //
    // Deduction for i = 3 as example:
    // T(5) = 2T(3)T(2) - T(1)
    // T(6) = 2T(3)T(3) - T(0)
    for (int n = 3; n <= num_mu_half; ++n) {
        p2 = p0;
        h_sparse->amxsy(2.0, *p1, *p0);
        cheb_mu[n] = vec_dot(wf_in, *p2);

        int i_dec = 2 * n - 1;
        if ((i_dec > num_mu_half) && (i_dec <= num_mu)) {
            cheb_mu[i_dec] = 2.0 * vec_dot(*p2, *p1) - mu1;
        }
        i_dec++;
        if ((i_dec > num_mu_half) && (i_dec <= num_mu)) {
            cheb_mu[i_dec] = 2.0 * vec_norm2(*p2) - mu0;
        }

        // Update pointers
        p0 = p1;
        p1 = p2;
        if (meta->is_master) {
            prog_bar.count();
        }
    }
}


/**
 * @brief Evaluate Chebyshev moments of DC <psi1_{\alpha}|Tn|psi2_{\alpha}>.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 * @param[in] h_sparse rescaled sparse Hamiltonian
 * @param[in] meta meta info
 * @param[in] wf_in initial wave function
 * @param[out] cheb_mu Chebyshev moments of DC cheb_mu(e,n) for n = 0, 1, 2, ..., num_mu-1,
 *  for e = 0, 2, ..., num_energies-1, must be pre-allocated.
 *
 */
template <typename sparse_t>
void cheb_eval_mu_dc(
    const sparse_t* h_sparse,
    const TBPMMeta* meta,
    const Eigen::VectorXcd& psi1_a,
    const std::vector<Eigen::VectorXcd>& psi2_a,
    std::vector<Eigen::VectorXcd>& cheb_mu)
{
    // Initialize results and buffers
    size_t dim = h_sparse->get_dim();
    assert(dim == psi1_a.size());
    assert(dim == psi2_a.back().size());
    Eigen::VectorXcd t0(dim);
    Eigen::VectorXcd t1(dim);

    size_t num_energies = cheb_mu.size();   // number of energies (E)
    size_t num_bes_max  = cheb_mu.back().size();   // number of Chebyshev moments (N)

    // Initialize pointers
    Eigen::VectorXcd* p0 = &t0;
    Eigen::VectorXcd* p1 = &t1;
    Eigen::VectorXcd* p2 = nullptr;

    // Initial condition
    // t0 = H * psi1_a
    // t1 = 2 * H * t0 - psi1_a
    h_sparse->mv(psi1_a, t0);
    h_sparse->amxsy_out(2.0, t0, psi1_a, t1);

    for (size_t i_e = 0; i_e < num_energies; ++i_e) {
        cheb_mu[i_e][0] = vec_dot(psi1_a, psi2_a[i_e]);
        cheb_mu[i_e][1] = vec_dot(t0, psi2_a[i_e]);
        cheb_mu[i_e][2] = vec_dot(t1, psi2_a[i_e]);
    }

    ProgressBar prog_bar(num_bes_max, 64, "Finished mu order");
    if (meta->is_master) {
        prog_bar.count(); // mu1
        prog_bar.count(); // mu2
    }

    // Continue the series expansion
    // p2 = 2 * H * p1 - p0
    // cheb_mu[i, i_e] = <p2|psi2_a[i_e]>
    for (size_t n = 3; n < num_bes_max; ++n) {
        p2 = p0;
        h_sparse->amxsy(2.0, *p1, *p0);
        for (size_t i_e = 0; i_e < num_energies; ++i_e) {
            cheb_mu[i_e][n] = vec_dot(*p2, psi2_a[i_e]);
        }
        p0 = p1;
        p1 = p2;
        if (meta->is_master) {
            prog_bar.count();
        }
    }
}


/**
 * @brief Evaluate quasi-eigenstates.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 * @param[in] h_sparse rescaled sparse Hamiltonian
 * @param[in] meta meta info
 * @param[in] coeff integrated exp(iEt)*cn(t), dimension (num_bes, num_eng)
 * @param[in] wf_in initial state
 * @param[out] wf_qe quasi-eigenstates, must be pre-allocated, dimension (num_eng, n_wf)
 */
template <typename sparse_t>
void cheb_eval_qe(
    const sparse_t* h_sparse,
    const TBPMMeta* meta,
    const std::vector<Eigen::VectorXcd>& coeff,
    const Eigen::VectorXcd& wf_in,
    std::vector<Eigen::VectorXcd>& wf_qe)
{
    // Initialize results and buffers
    size_t dim = h_sparse->get_dim();
    assert(dim == wf_in.size());
    assert(dim == wf_qe.back().size());
    Eigen::VectorXcd t0(dim);
    Eigen::VectorXcd t1(dim);

    // Initialize pointers
    Eigen::VectorXcd* p0 = &t0;
    Eigen::VectorXcd* p1 = &t1;
    Eigen::VectorXcd* p2 = nullptr;

    // Initial condition
    h_sparse->mv(wf_in, t0);
    h_sparse->amxsy_out(2.0, t0, wf_in, t1);
    size_t num_eng = coeff.back().size();
    for (size_t i_e = 0; i_e < num_eng; ++i_e) {
        axpbypcz_out(coeff[0][i_e], wf_in, coeff[1][i_e], t0, coeff[2][i_e], t1, wf_qe[i_e]);
    }

    // Continue the series expansion
    size_t num_series = coeff.size();
    ProgressBar prog_bar(num_series, 64, "Finished Bessel order");
    if (meta->is_master) {
        prog_bar.count(); // coeff[0] * wf_in
        prog_bar.count(); // coeff[1] * t0
        prog_bar.count(); // coeff[2] * t1
    }

    for (size_t n = 3; n < num_series; ++n) {
        p2 = p0;
        h_sparse->amxsy(2.0, *p1, *p0);
        for (size_t i_e = 0; i_e < num_eng; ++i_e) {
            axpy(coeff[n][i_e], *p2, wf_qe[i_e]);
        }
        p0 = p1;
        p1 = p2;
        if (meta->is_master) {
            prog_bar.count();
        }
    }
}


/**
 * @brief Evaluate quasi-eigenstates.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 * @param[in] h_sparse rescaled sparse Hamiltonian
 * @param[in] meta meta info
 * @param[in] coeff integrated exp(iEt)*cn(t), dimension (num_eng, num_bes)
 * @param[in] wf_in initial state
 * @param[out] wf_qe quasi-eigenstates, must be pre-allocated, dimension (num_eng, n_wf)
 */
template <typename sparse_t>
void cheb_eval_qe_new(
    const sparse_t* h_sparse,
    const TBPMMeta* meta,
    const std::vector<Eigen::VectorXcd>& coeff,
    const Eigen::VectorXcd& wf_in,
    std::vector<Eigen::VectorXcd>& wf_qe)
{
    // Initialize results and buffers
    size_t dim = h_sparse->get_dim();
    assert(dim == wf_in.size());
    assert(dim == wf_qe.back().size());
    Eigen::VectorXcd t0(dim);
    Eigen::VectorXcd t1(dim);

    // Initialize pointers
    Eigen::VectorXcd* p0 = &t0;
    Eigen::VectorXcd* p1 = &t1;
    Eigen::VectorXcd* p2 = nullptr;

    // Initial condition
    h_sparse->mv(wf_in, t0);
    h_sparse->amxsy_out(2.0, t0, wf_in, t1);
    size_t num_eng = coeff.size();
    for (size_t i_e = 0; i_e < num_eng; ++i_e) {
        axpbypcz_out(coeff[i_e][0], wf_in, coeff[i_e][1], t0, coeff[i_e][2], t1, wf_qe[i_e]);
    }

    // Continue the series expansion
    size_t num_series = coeff.back().size();
    ProgressBar prog_bar(num_series, 64, "Finished Bessel order");
    if (meta->is_master) {
        prog_bar.count(); // coeff[0] * wf_in
        prog_bar.count(); // coeff[1] * t0
        prog_bar.count(); // coeff[2] * t1
    }

    for (size_t n = 3; n < num_series; ++n) {
        p2 = p0;
        h_sparse->amxsy(2.0, *p1, *p0);
        for (size_t i_e = 0; i_e < num_eng; ++i_e) {
            axpy(coeff[i_e][n], *p2, wf_qe[i_e]);
        }
        p0 = p1;
        p1 = p2;
        if (meta->is_master) {
            prog_bar.count();
        }
    }
}
} // namespace tbplas::tbpm
#endif