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

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

namespace tbplas::tbpm {
template <typename sparse_t>
class HaydockGPU final : public AbstractHaydock<sparse_t> {
private:
    CudaConfig cuda_;
public:
    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());
        n2.setZero();
        Eigen::VectorXcd n_buf(n1.size());
        n_buf.setZero();
        coefa.resize(n_depth);
        coefb.resize(n_depth);

        // 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;
        auto h_meta = h_sparse.get_meta();

        // Propagate
        cuda_.haydock_coef_GPU(
            meta.rescale, n_depth,
            h_meta.indptr_data, h_meta.indptr_size,
            h_meta.indices_data, h_meta.indices_size,
            h_meta.values_data, h_meta.values_size,
            p1->data(), p1->size(), p2->data(), p2->size(),
            p_buf->data(), p_buf->size(),
            coefa.data(), coefa.size(),
            coefb.data(), coefb.size());
    }
};

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