#include "lindhard_core.h"
#include "base/consts.h"
#include "base/funcs.h"

namespace tbplas::diag {
using base::conj_t;
using base::exp_it;
using base::UIMAG;

AbstractDPCore::AbstractDPCore(
    const size_t& num_orb,
    const size_t& num_states,
    const size_t& num_qpt,
    const int& conv)
    : convention(conv)
{
    eng_k = Eigen::VectorXd::Zero(num_states);
    eng_kq = Eigen::VectorXd::Zero(num_states);
    wfn_k = Eigen::MatrixXcd::Zero(num_orb, num_states);
    wfn_kq = Eigen::MatrixXcd::Zero(num_orb, num_states);
    f_k = Eigen::VectorXd::Zero(num_states);
    f_kq = Eigen::VectorXd::Zero(num_states);
    phase = Eigen::MatrixXcd::Zero(num_orb, num_qpt);
    delta_eng = Eigen::MatrixXd::Zero(num_states, num_states);
    product = Eigen::MatrixXcd::Zero(num_states, num_states);
}

void DefaultDPCore::calc_product(const int& iq)
{
    // Reset arrays.
    delta_eng.setZero();
    product.setZero();

    size_t num_states = eng_k.size();
    // OpenMP requires variables to be defined in advance.
    int nn = 0;
    complex_t prod = 0.0;
    if (convention == 1) {
#pragma omp parallel for private(nn, prod)
        for (ptrdiff_t jj = 0; jj < num_states; ++jj) {
            for (nn = 0; nn < num_states; ++nn) {
                delta_eng(nn, jj) = eng_k(jj) - eng_kq(nn);
                prod = wfn_kq.col(nn).dot(wfn_k.col(jj));
                product(nn, jj) = prod * conj_t(prod) * (f_k(jj) - f_kq(nn));
            }
        }
    } else {
#pragma omp parallel for private(nn, prod)
        for (ptrdiff_t jj = 0; jj < num_states; ++jj) {
            for (nn = 0; nn < num_states; ++nn) {
                delta_eng(nn, jj) = eng_k(jj) - eng_kq(nn);
                prod = (wfn_kq.col(nn).conjugate().array() * wfn_k.col(jj).array() * phase.col(iq).array()).sum();
                product(nn, jj) = prod * conj_t(prod) * (f_k(jj) - f_kq(nn));
            }
        }
    }
}

void DefaultDPCore::calc_dyn_pol(
    const Eigen::VectorXd& omegas,
    const double& delta,
    const int& iq,
    Eigen::MatrixXcd& dyn_pol) const
{
    complex_t cdelta = complex_t(0.0, delta);
#pragma omp parallel for
    for (ptrdiff_t iw = 0; iw < omegas.size(); ++iw) {
        dyn_pol(iw, iq) += (product.array() / (delta_eng.array() + omegas(iw) + cdelta)).sum();
    }
}

AbstractACCore::AbstractACCore(
    const size_t& num_orb,
    const size_t& num_states)
{
    eng_k = Eigen::VectorXd::Zero(num_states);
    wfn_k = Eigen::MatrixXcd::Zero(num_orb, num_states);
    f_k = Eigen::VectorXd::Zero(num_states);
    delta_eng = Eigen::MatrixXd::Zero(num_states, num_states);
    product = Eigen::MatrixXcd::Zero(num_states, num_states);
    v_alpha = Eigen::MatrixXcd::Zero(num_orb, num_orb);
    v_beta = Eigen::MatrixXcd::Zero(num_orb, num_orb);
}

void DefaultACCore::calc_product(
    const Eigen::Matrix3Xd& k_points_cart,
    const int& ik,
    const int& comp_alpha,
    const int& comp_beta)
{
    // Reset arrays.
    v_alpha.setZero();
    v_beta.setZero();
    delta_eng.setZero();
    product.setZero();

    size_t num_states = eng_k.size();
    size_t num_hop = hop_i.size();

    // Construct velocity operators
    // NOTE: this must be done in serial mode because different threads
    // may have the same ib1 and ib2.
    for (ptrdiff_t ih = 0; ih < num_hop; ++ih) {
        complex_t ph = exp_it(k_points_cart.col(ik).dot(dr_nm.col(ih))) * hop_v[ih];
        complex_t f1 = UIMAG * ph;
        complex_t f2 = UIMAG * conj_t(ph);
        double x_alpha = dr_nm(comp_alpha, ih);
        double x_beta = dr_nm(comp_beta, ih);
        // It's hard to distinguish ii and jj, so we use nn instead of ii.
        size_t nn = hop_i[ih];
        size_t jj = hop_j[ih];
        v_alpha(nn, jj) += f1 * x_alpha;
        v_alpha(jj, nn) -= f2 * x_alpha;
        v_beta(nn, jj) += f1 * x_beta;
        v_beta(jj, nn) -= f2 * x_beta;
    }

    // Construct delta_eng and prod_df
    int nn = 0;
    complex_t prod1 = 0.0;
    complex_t prod2 = 0.0;
#pragma omp parallel for private(nn, prod1, prod2)
    for (ptrdiff_t jj = 0; jj < num_states; ++jj) {
        for (nn = 0; nn < num_states; ++nn) {
            delta_eng(nn, jj) = eng_k(jj) - eng_k(nn);
            if (abs(delta_eng(nn, jj)) >= 1.0e-7) {
                prod1 = wfn_k.col(nn).dot(v_alpha * wfn_k.col(jj));
                prod2 = wfn_k.col(jj).dot(v_beta * wfn_k.col(nn));
                product(nn, jj) = prod1 * prod2 * (f_k(jj) - f_k(nn)) / delta_eng(nn, jj);
            }
        }
    }
}

void DefaultACCore::calc_ac_cond(
    const Eigen::VectorXd& omegas,
    const double& delta,
    Eigen::VectorXcd& ac_cond) const
{
    complex_t cdelta = complex_t(0.0, delta);
#pragma omp parallel for
    for (ptrdiff_t iw = 0; iw < omegas.size(); ++iw) {
        ac_cond(iw) += (product.array() / (delta_eng.array() - omegas(iw) - cdelta)).sum();
    }
}

#ifdef WITH_CUDA
/// @todo Implement the methods
void DPCoreGPU::calc_product(const int& iq)
{
    std::cerr << "ERROR: DPCoreGPU::calc_product not implemented yet!" << std::endl;
    std::exit(-1);
}

void DPCoreGPU::calc_dyn_pol(
    const Eigen::VectorXd& omegas,
    const double& delta,
    const int& iq,
    Eigen::MatrixXcd& dyn_pol) const
{
    std::cerr << "ERROR: DPCoreGPU::calc_dyn_pol not implemented yet!" << std::endl;
    std::exit(-1);
}

void ACCoreGPU::calc_product(
    const Eigen::Matrix3Xd& k_points_cart,
    const int& ik,
    const int& comp_alpha,
    const int& comp_beta)
{
    std::cerr << "ERROR: ACCoreGPU::calc_product not implemented yet!" << std::endl;
    std::exit(-1);
}

void ACCoreGPU::calc_ac_cond(
    const Eigen::VectorXd& omegas,
    const double& delta,
    Eigen::VectorXcd& ac_cond) const
{
    std::cerr << "ERROR: ACCoreGPU::calc_ac_cond not implemented yet!" << std::endl;
    std::exit(-1);
}
#endif // WITH_CUDA
} // namespace tbplas::diag