/**
 * @brief Functions and classes for the Fermi-Dirac operator.
 *
 * @author Zewen Wu, Miao Zhang and Yunhai Li
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_TBPM_FERMI_DIRAC_H
#define TBPLAS_TBPM_FERMI_DIRAC_H

#include <eigen3/Eigen/Dense>
#include <eigen3/unsupported/Eigen/FFT>

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

namespace tbplas::tbpm {
/**
 * @brief Get occupation number according to Fermi-Dirac distribution, i.e.,
 * f(x) between eqn.36 and eqn.36 of CPC paper.
 *
 * @param[in] beta rescaled or unrescaled 1/kBT
 * @param[in] mu rescaled or unrescaled chemical potential
 * @param[in] x rescaled or unrescaled energy
 * @return double
 *
 * @note The formula follows:
 *   f(x) = z*exp(-beta*x) / (1+z*exp(-beta*x))
 * with z = exp(beta*mu), which yields:
 *   f(x) = exp(beta*(mu-x)) / (1+exp(beta*(mu-x)
 *
 * @note We decided to make it a general function taking both rescaled and unrescaled
 * arguments, and do not assume x to be within [-1, 1].
 */
double fermi_dirac(const double& beta, const double& mu, const double& x);

/**
 * @brief Fermi-Dirac operator based on Chebyshev expansion
 *
 * @tparam sparse_t datatype of sparse Hamiltonian
 */
template <typename sparse_t>
class FermiDiracOperator {
public:
    /**
     * @brief Construct a new Fermi Dirac Operator object
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     */
    explicit FermiDiracOperator(
        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_cheb_coeff(config, meta); // initialize cheb_coeff_ and cheb_coeff_omf_
    }

    /**
     * @brief Act on wave function
     *
     * @param[inout] wf_in incoming wave function to act on
     * @param[in] one_minus_fermi whether the operator is 1-f(H)
     */
    void act(Eigen::VectorXcd& wf_in, bool one_minus_fermi = false)
    {
        const Eigen::VectorXd* cheb = one_minus_fermi ? &cheb_coeff_omf_ : &cheb_coeff_;
        cheb_eval_wf(h_sparse_, *cheb, wf_in, t0_, t1_, wf_out_);
        vec_copy(wf_out_, wf_in);
    }

    /**
     * @brief Act on wave function
     *
     * @param[in] wf_in incoming wave function to act on
     * @param[out] wf_out wave function after acting the operator
     * @param[in] one_minus_fermi whether the operator is 1-f(H)
     */
    void act_out(const Eigen::VectorXcd& wf_in, Eigen::VectorXcd& wf_out, bool one_minus_fermi = false)
    {
        const Eigen::VectorXd* cheb = one_minus_fermi ? &cheb_coeff_omf_ : &cheb_coeff_;
        cheb_eval_wf(h_sparse_, *cheb, wf_in, t0_, t1_, wf_out);
    }

    /**
     * @brief Act on wave function with the f^{1/2}(H) operator.
     *
     * @param[in] wf_in incoming wave function to act on
     * @param[out] wf_out wave function after acting the f^{1/2}(H) operator
     */
    void act_sqrt_fd(const Eigen::VectorXcd& wf_in, Eigen::VectorXcd& wf_out) {
        // Use the newly computed cheb_coeff_sqrt_fd_
        cheb_eval_wf(h_sparse_, cheb_coeff_sqrt_fd_, wf_in, t0_, t1_, wf_out);
    }

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

    /// @brief Chebyshev coefficients for f(H)
    Eigen::VectorXd cheb_coeff_;

    /// @brief Chebyshev coefficients for 1 - f(H)
    Eigen::VectorXd cheb_coeff_omf_;

    /// @brief Chebyshev coefficients for f^{1/2}(H)
    Eigen::VectorXd cheb_coeff_sqrt_fd_;

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

    /**
     * @brief Set up the 'cheb_coeff_' attribute
     *
     * @param[in] config user-defined calculation parameters
     * @param[in] meta additional meta info.
     */
    void build_cheb_coeff(const TBPMConfig& config, const TBPMMeta& meta)
    {
        int num_fft = config.num_fermi_fft_steps;
        double theta = base::TWOPI / num_fft;
        Eigen::VectorXd fd(num_fft);
        Eigen::VectorXd fd_omf(num_fft);
        Eigen::VectorXd fd_sqrt(num_fft);

        // Build fd and fd_omf
        // We reuse the results of Fermi-Dirac distribution to save time, since f(H) and 1-f(H)
        // always appear in the same time.
#pragma omp parallel for simd
        for (int n = 0; n < num_fft; ++n) {
            fd[n] = fermi_dirac(meta.beta_re, meta.mu_re, std::cos(theta * n));
            fd_omf[n] = 1.0 - fd[n];
            fd_sqrt[n] = std::sqrt(std::max(0.0, fd[n])); // Use std::max to prevent sqrt of negative due to precision
        }

        // FFT on fd and fd_omf to yield the coefficients, according to eqn.35 of CPC paper.
        // Note there is a typo in the equation. It should be N rather than pi on the denominator.
        // The factor of 2.0 comes from our definition of ck in eqn.34. As we are doing FFT, there
        // is no 1/num_fft factor in the results. So we have to add it back. That's why we multiply
        // the results by 2.0 / num_fft. For inverse FFT, the result has already the factor of
        // 1/num_fft, and we only need the factor of 2.0.
        Eigen::FFT<double> fft_plan;
        Eigen::VectorXcd work(num_fft);
        fft_plan.fwd(work, fd);
        fd = work.real() * 2.0 / num_fft;
        fft_plan.fwd(work, fd_omf);
        fd_omf = work.real() * 2.0 / num_fft;
        fft_plan.fwd(work, fd_sqrt);
        fd_sqrt = work.real() * 2.0 / num_fft;

        // Truncate fd and fd_omf since we have extended [0, pi] to [0, 2pi] in utilizing FFT
        int num_cheb = num_fft / 2;
        fd.conservativeResize(num_cheb);
        fd_omf.conservativeResize(num_cheb);
        fd_sqrt.conservativeResize(num_cheb);

        // Build and apply Jackson kernel
        Eigen::VectorXd kernel(num_cheb);
        jackson_kernel(kernel);
        vec_self_mul_vec(fd, kernel);
        vec_self_mul_vec(fd_omf, kernel);
        vec_self_mul_vec(fd_sqrt, kernel); 

        // First coefficient of Chebyshev expansion has an additional factor of 0.5.
        fd[0] *= 0.5;
        fd_omf[0] *= 0.5;
        fd_sqrt[0] *= 0.5;

        // Truncate fd an fd_omf according to significance
        vec_trunc(fd, config.fermi_cheb_precision);
        vec_trunc(fd_omf, config.fermi_cheb_precision);
        vec_trunc(fd_sqrt, config.fermi_cheb_precision);

        // Finally, update chebyshev coefficients
        cheb_coeff_ = fd;
        cheb_coeff_omf_ = fd_omf;
        cheb_coeff_sqrt_fd_ = fd_sqrt;
    }
};
} // namespace tbplas::tbpm
#endif