/**
 * @brief Functions for Recursion Haydock method.
 *
 * @author Zewen Wu, Miao Zhang and Yunhai Li
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_TBPM_HAYDOCK_H
#define TBPLAS_TBPM_HAYDOCK_H

#include <iostream>

#include <eigen3/Eigen/Dense>

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

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

/**
 * @brief Abstract Haydock class.
 *
 * @tparam sparse_t datatype of sparse Hamiltonian matrix
 */
template <typename sparse_t>
class AbstractHaydock {
public:
    virtual void haydock_coef(
        const sparse_t& h_sparse,
        const TBPMMeta& meta,
        int n_depth,
        Eigen::VectorXcd& n1,
        Eigen::VectorXcd& coefa,
        Eigen::VectorXd& coefb)
        = 0;

    /**
     * @brief Calculate local DOS using Haydock recursive method.
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] config user-defined calculation configurations
     * @param[in] meta additional meta info.
     * @param[out] energy energies of LDOS in eV
     * @param[out] ldos normalized LDOS in 1/eV
     */
    void ldos_haydock(
        const sparse_t& h_sparse,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::VectorXd& energy,
        Eigen::VectorXd& ldos)
    {
        assert(!config.ldos_orbital_indices.empty());
        PARA_EXEC(std::cout << "Getting Haydock coefficients." << std::endl);

        // Initialize coefficient vectors for Haydock recursion
        int n_depth = config.ldos_recursion_depth;
        size_t n_wf = h_sparse.get_dim();
        Eigen::VectorXcd coefa = Eigen::VectorXcd::Zero(n_depth);
        Eigen::VectorXd coefb = Eigen::VectorXd::Zero(n_depth);
        Eigen::VectorXcd wf0 = Eigen::VectorXcd::Zero(n_wf);

        // Initialize random generator with the given seed
        RandomGenerator rand_gen(config.seed);

        // Loop over random 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);

            // Initialize wave function vectors
            Eigen::VectorXcd wf_temp(config.ldos_orbital_indices.size());

            // Generate random initial wave function for the current sample
            rand_gen.random_state(wf_temp, seed_i);

            // Assign the random wave function values to the specified site indices
            vec_init(wf0, 0.0);
            size_t counter = 0;
            for (const auto& i : config.ldos_orbital_indices) {
                wf0[i] = wf_temp[counter];
                ++counter;
            }

            // Calculate Haydock coefficients for the current sample
            Eigen::VectorXcd a_work = Eigen::VectorXcd::Zero(n_depth);
            Eigen::VectorXd b_work = Eigen::VectorXd::Zero(n_depth);
            haydock_coef(h_sparse, meta, n_depth, wf0, a_work, b_work);
            coefa += a_work;
            coefb += b_work;
        }
        coefa /= num_samples;
        coefb /= num_samples;

        PARA_EXEC(std::cout << "Calculating LDOS with Green's function." << std::endl);

        // Initialize energy and LDOS vectors
        int num_time_steps = config.num_time_steps;
        int num_eng = 2 * num_time_steps + 1;
        energy = Eigen::VectorXd(num_eng);
        ldos = Eigen::VectorXd(num_eng);
        for (int i = 0; i < num_eng; ++i) {
            energy[i] = 0.5 * (i - num_time_steps) * meta.energy_range / num_time_steps;
        }

        // Calculate LDOS using Green's function, with parallelization
        complex_t g00 = 0.0;
#pragma omp parallel for private(g00)
        for (int i = 0; i < num_eng; ++i) {
            g00 = green_function(energy[i], config.ldos_delta, coefa, coefb, n_depth);
            ldos[i] = -1.0 / tbplas::base::PI * g00.imag();
        }
    }

private:
    /**
     * @brief Calculate G0(energy+i*delta).
     *
     * @param[in] energy energy in eV
     * @param[in] delta infinitesimal in eV
     * @param[in] coefa coefficients {a_i} from haydock_coef
     * @param[in] coefb coefficients {b_i} from haydock_coef
     * @param[in] n_depth depth of recursion
     * @return complex_t
     *
     * @note We make it inline to be header-only.
     */
    inline static complex_t green_function(
        double energy,
        double delta,
        const Eigen::VectorXcd& coefa,
        const Eigen::VectorXd& coefb,
        int n_depth)
    {
        complex_t e_cmplx = { energy, delta };
        complex_t g00 = { 0.0, 0.0 };
        for (int i = n_depth - 1; i >= 0; --i) {
            g00 = 1.0 / (e_cmplx - coefa[i] - coefb[i] * coefb[i] * g00);
        }
        return g00;
    }
};

template <typename sparse_t>
class HaydockCPU final : public AbstractHaydock<sparse_t> {
public:
    /**
     * @brief Calculate the a_i and b_i in eqn. 46-50 of CPC paper.
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] meta additional meta info.
     * @param[in] n_depth recursion depth
     * @param[inout] n1 Initial wave function for recursion, namely l0 in the paper.
     * @param[out] coefa a_i for i = 0, 1, 2, ..., n_depth-1
     * @param[out] coefb b_i for i = 1, 2, 3, ..., n_depth
     *
     * @note Although a_i begins from 0 while b_i begins from 1, they actually have
     * the same length, since it is a_i and b_i+1 that appear in the same depth of
     * recursion. In some papers, e.g. 10.1016/j.physleta.2012.09.009, the notation
     * is that both a_i and b_i begin from 1.
     */
    void haydock_coef(
        const sparse_t& h_sparse,
        const TBPMMeta& meta,
        int n_depth,
        Eigen::VectorXcd& n1,
        Eigen::VectorXcd& coefa,
        Eigen::VectorXd& coefb) final
    {
        // Intermediate arrays
        Eigen::VectorXcd n2(n1.size());
        Eigen::VectorXcd n_buf(n1.size());
        coefa.resize(n_depth);
        coefb.resize(n_depth);

        /* Formulation of the method

        NOTE: We begin from i=1 and use ni instead of li, which are slightly different
        from the CPC paper.

        Initial condition:
          a1 = <n1|H|n1>
          m2 = (H - a1)|n1>
          b2 = |m2|
          n2 = m2 / b2

        Iteration:
          ai = <ni|H|ni>
          mi+1 = (H - ai)|ni> - bi|ni-1>
          bi+1 = |mi+1|
          ni+1 = mi+1 / bi+1

        where mi+1 is the un-normalized wave function at iteration i+1.

        In our implementation:
          p1, n1: ni-1
          p2, n2: ni
          p_buf, n_buf: H|ni> and mi+1 (buffer for intermediate results)
        */

        // Get initial values for a and b
        h_sparse.amv(meta.rescale, n1, n_buf); // n_buf = H|n1>
        coefa[0] = vec_dot(n1, n_buf); // a1 = <n1|H|n1> = <n1|n_buf>
        axpy(-coefa[0], n1, n_buf); // m2 = (H - a1)|n1> = |n_buf> - a1|n1>
        coefb[0] = vec_norm(n_buf); // b2 = |m2|
        vec_div_sca(n_buf, coefb[0], n2); // n2 = m2 / b2

        // Initialize pointers
        Eigen::VectorXcd* p1 = &n1;
        Eigen::VectorXcd* p2 = &n2;
        Eigen::VectorXcd* p_buf = &n_buf;

        // Recursion
        base::ProgressBar prog_bar(n_depth, 100, "Finished depth");
        PARA_EXEC(prog_bar.count()); // We have done the 0th step in initialization.
        for (int i = 1; i < n_depth; ++i) {
            h_sparse.amv(meta.rescale, *p2, *p_buf); // n_buf = H|ni>
            coefa[i] = vec_dot(*p2, *p_buf); // ai = <ni|H|ni> = <n1|n_buf>
            axpbypz(-coefa[i], *p2, -coefb[i - 1], *p1, *p_buf); // mi+1 = (H - ai)|ni> - bi|ni-1> = n_buf - ai|ni> - bi|ni-1>
            coefb[i] = vec_norm(*p_buf); // bi+1 = |mi+1|
            vec_self_div_sca(*p_buf, coefb[i]); // ni+1 = mi+1 / bi+1

            // Now p1->ni-1, p2->ni, p_buf->ni+1. We want p1->ni, p2->ni+1 and p_buf->ni-1.
            // This can be done with a cyclic permutation, which can be decomposed into pairwise
            // permutations.
            std::swap(p_buf, p2);
            std::swap(p_buf, p1);
            PARA_EXEC(prog_bar.count());
        }
    }
};
} // namespace tbplas::tbpm
#endif
