/**
 * @brief Functions for calculating DC conductivity using Kubo-Bastin method.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */
#ifdef WITH_CUDA
#ifndef TBPLAS_TBPM_KUBO_BASTIN_GPU_H
#define TBPLAS_TBPM_KUBO_BASTIN_GPU_H

#include "cuda/cudaConfig.h"
#include "kubo_bastin.h"

namespace tbplas::tbpm {
template <typename sparse_t>
class KuboBastinGPU final : public AbstractKuboBastin<sparse_t> {
private:
    CudaConfig cuda_;

public:
    /**
     * @brief Calculate \mu_{mn} in eqn.70 of CPC paper.
     *
     * @param[in] h_sparse scaled sparse Hamiltonian
     * @param[in] curr_x sparse current operator along x direction
     * @param[in] curr_y sparse current operator along y direction
     * @param[in] config user-defined calculation configurations
     * @param[in] meta additional meta info.
     * @param[out] mu_avg averaged \mu_{mn}, dimension (num_kernel, num_kernel)
     */
    void kbdc_mu(
        const sparse_t& h_sparse,
        const sparse_t& curr_x,
        const sparse_t& curr_y,
        const TBPMConfig& config,
        const TBPMMeta& meta,
        Eigen::MatrixXcd& mu_avg) final
    {
        PARA_EXEC(std::cout << "Calculating mu of DC conductivity." << std::endl);

        // Initialize wave functions
        size_t n_wf = h_sparse.get_dim();
        int num_kernel = config.dckb_num_kernel;
        Eigen::VectorXcd wf0(n_wf); // Initial random state
        Eigen::VectorXcd t0(n_wf); // Tn(H)|wf0> and Tm(H)*v_alpha|wf0>
        Eigen::VectorXcd t1(n_wf); // Tn(H)|wf0> and Tm(H)*v_alpha|wf0>
        std::vector<Eigen::VectorXcd> wf_vb_tn; // v_beta*Tn(H)|wf0>
        mat_alloc(wf_vb_tn, n_wf, num_kernel);

        // Initialize correlation functions
        mu_avg = Eigen::MatrixXcd::Zero(num_kernel, num_kernel);
        Eigen::MatrixXcd mu_sample(num_kernel, num_kernel);
        mu_sample.setZero();

        // Initialize utilities
        RandomGenerator rand_gen(config.seed);
        ProgressBar prog_bar(num_kernel, 256, "Finished i");

        // Get Jackson kernel
        // First element has a factor of 0.5 due to the 1/(1+\delta(0,n)) factor in eqn.70.
        Eigen::VectorXd kernel(num_kernel);
        jackson_kernel(kernel);
        kernel[0] *= 0.5;
        bool is_row_major = false;
        if (mu_sample.IsRowMajor) {
            is_row_major = true;
        }

        // Main loop
        int num_samples = config.num_random_samples;
        int rank = meta.mpi_rank;
        auto h_meta = h_sparse.get_meta();
        auto curr_x_meta = curr_x.get_meta();
        auto curr_y_meta = curr_y.get_meta();
        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);
            PARA_EXEC(prog_bar.reset());

            // Evaluate v_beta*Tn(H)|wf0> and save to wf_tb_vn
            PARA_EXEC(std::cout << "Evaluating v_beta*Tn(H)|wf0>." << std::endl);
            rand_gen.random_state(wf0, seed_i);
            vec_copy(wf0, t0);

            // Propagate
            cuda_.kbdc_mu_GPU2(
                wf0.data(), wf0.size(),
                config.dckb_component, num_kernel,
                h_meta.indptr_data, h_meta.indptr_size,
                h_meta.indices_data, h_meta.indices_size,
                h_meta.values_data, h_meta.values_size,
                curr_x_meta.indptr_data, curr_x_meta.indptr_size,
                curr_x_meta.indices_data, curr_x_meta.indices_size,
                curr_x_meta.values_data, curr_x_meta.values_size,
                curr_y_meta.indptr_data, curr_y_meta.indptr_size,
                curr_y_meta.indices_data, curr_y_meta.indices_size,
                curr_y_meta.values_data, curr_y_meta.values_size,
                t0.data(), t0.size(),
                t1.data(), t1.size(),
                mu_sample.data(),
                is_row_major, prog_bar, wf_vb_tn);

            // Apply Jackson kernel and average
            // Note: eqn.70 of CPC paper is \mu_{nm} rather than \mu_{mn}. But it makes no sense
            // swapping m and n. Also, PRL 114, 116602 (2015) uses \mu_{mn}. So we do not take the
            // transpose here.
            double factor = 1.0 / static_cast<double>(num_samples);
            int n = 0;
#pragma omp parallel for private(n)
            for (int m = 0; m < num_kernel; ++m) {
                for (n = 0; n < num_kernel; ++n) {
                    mu_avg(m, n) += factor * kernel(m) * kernel(n) * mu_sample(m, n);
                }
            }
        }
    }
};

} // namespace tbplas::tbpm
#endif // TBPLAS_TBPM_KUBO_BASTIN_GPU_H
#endif // WITH_CUDA