/**
 * @brief Functions for dense matrix and vector operations.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 *
 * @note
 *   1. Abbreviation for vector-vector operations: add, sub(stract), mul(tiply), div(ide).
 *   2. Abbreviation for composite vector operations: p(lus), m(atrix). We do not use add
 *      for + since a indicates a scalar operand.
 */

#ifndef TBPLAS_TBPM_DENSE_H
#define TBPLAS_TBPM_DENSE_H

#include <cassert>
#include <vector>

#include <eigen3/Eigen/Dense>

#include "base/datatypes.h"

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

/**
 * @brief Calculate the inner product between double vectors
 *
 * @param[in] x left operand
 * @param[in] y right operand
 * @return the inner product
 */
inline double vec_dot(const Eigen::VectorXd& x, const Eigen::VectorXd& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
    double product = 0.0;
#pragma omp parallel for reduction(+ : product) schedule(static)
    for (size_t i = 0; i < dim; ++i) {
        product += x[i] * y[i];
    }
    return product;
}

/**
 * @brief Calculate the inner product between complex vectors
 *
 * @param[in] x left operand
 * @param[in] x right operand
 * @return the inner product
 *
 * @note The conjugate of left operand is taken automatically.
 */
inline complex_t vec_dot(const Eigen::VectorXcd& x, const Eigen::VectorXcd& y, bool check_dim = true)
{
    size_t dim = x.size();
    if (check_dim) {
        assert(dim == y.size());
    }
    double real_part = 0.0;
    double imag_part = 0.0;
    double a = 0.0;
    double b = 0.0;
    double c = 0.0;
    double d = 0.0;
#pragma omp parallel for reduction(+ : real_part, imag_part) schedule(static) private(a, b, c, d)
    for (size_t i = 0; i < dim; ++i) {
        a = x[i].real();
        b = x[i].imag();
        c = y[i].real();
        d = y[i].imag();
        real_part += a * c + b * d;
        imag_part += a * d - b * c;
    }
    return { real_part, imag_part };
}

/**
 * @brief Calculate the inner product between complex vectors without conjugation
 * of left operand.
 *
 * @param[in] x left operand
 * @param[in] x right operand
 * @param[in] check_dim whether to enable the dimension check assertion in debug mode
 * @return the inner product
 *
 * @note Currently, this function is used by fast TBPM algorithms, where we need
 * to calculate the product of two vectors with different dimensions. So we add
 * the check_dim parameter to disable the dimension check assertion in debug mode.
 * The default value of check_dim should be true, in order to be consistent with vec_dot.
 */
inline complex_t vec_dot_no_conj(const Eigen::VectorXcd& x, const Eigen::VectorXcd& y, bool check_dim = true)
{
    size_t dim = x.size();
    if (check_dim) {
        assert(dim == y.size());
    }
    double real_part = 0.0;
    double imag_part = 0.0;
    double a = 0.0;
    double b = 0.0;
    double c = 0.0;
    double d = 0.0;
#pragma omp parallel for reduction(+ : real_part, imag_part) schedule(static) private(a, b, c, d)
    for (size_t i = 0; i < dim; ++i) {
        a = x[i].real();
        b = x[i].imag();
        c = y[i].real();
        d = y[i].imag();
        real_part += a * c - b * d;
        imag_part += a * d + b * c;
    }
    return { real_part, imag_part };
}

/**
 * @brief Calculate the norm^2 of vector
 *
 * @tparam vx_t datatype of vector
 * @param[in] x the vector
 * @return norm of vector
 */
template <typename vx_t>
double vec_norm2(const Eigen::MatrixBase<vx_t>& x)
{
    double result = 0.0;
    size_t dim = x.size();
#pragma omp parallel for reduction(+ : result) schedule(static)
    for (size_t i = 0; i < dim; ++i) {
        result += (std::conj(x[i]) * x[i]).real();
    }
    return result;
}

/**
 * @brief Calculate the norm of vector
 *
 * @tparam vx_t datatype of vector
 * @param[in] x the vector
 * @return norm of vector
 */
template <typename vx_t>
double vec_norm(const Eigen::MatrixBase<vx_t>& x)
{
    return std::sqrt(vec_norm2(x));
}

/**
 * @brief Initialize vector X = a.
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @param[inout] x vector x
 * @param[in] a scalar a
 */
template <typename vx_t, typename sa_t>
void vec_init(Eigen::MatrixBase<vx_t>& x, const sa_t& a)
{
    size_t dim = x.size();
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        x[i] = a;
    }
}

/**
 * @brief Deep copy X to Y.
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @param[in] x vector x
 * @param[out] y vector y
 */
template <typename vx_t, typename vy_t>
void vec_copy(const Eigen::MatrixBase<vx_t>& x, Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        y[i] = x[i];
    }
}

/**
 * @brief Truncate a decreasing vector according to threshold.
 *
 * @tparam vx_t datatype of vector x
 * @param[inout] x vector to truncate
 * @param[in] eps threshold below which elements will be dropped
 */
template <typename vx_t>
void vec_trunc(vx_t& x, const double& eps)
{
    size_t dim = x.size();
    assert(dim > 0);
    size_t i_max = 0;
    bool truncate = false;
    for (size_t i = 0; i < dim - 1; ++i) {
        if ((std::abs(x[i]) <= eps) && (std::abs(x[i + 1]) <= eps)) {
            i_max = i;
            truncate = true;
            break;
        }
    }
    if (truncate) {
        x.conservativeResize(i_max + 1); // +1 to avoid yielding empty vector
    }
}

/**
 * @brief Allocate matrix X as std::vector<Eigen::VectorX*>.
 *
 * @tparam vx_t datatype of matrix X
 * @param[out] x matrix X
 * @param[in] num_row number of rows
 * @param[in] num_col number of columns
 */
template <typename vx_t>
void mat_alloc(std::vector<vx_t>& x, size_t num_row, size_t num_col)
{
    x = std::vector<vx_t>(num_col);
    for (vx_t& v : x) {
        v.resize(num_row);
    }
}

/**
 * @brief Initialize matrix X = a.
 *
 * @tparam vx_t datatype of matrix X
 * @tparam sa_t datatype of scalar a
 * @param[out] x matrix X
 * @param[in] a scalar a
 */
template <typename vx_t, typename sa_t>
void mat_init(std::vector<vx_t>& x, const sa_t& a)
{
    for (vx_t& v : x) {
        vec_init(v, a);
    }
}

/**
 * @brief Deep copy vector X to each column of matrix Y.
 *
 * @tparam vx_t datatype of x
 * @tparam vy_t datatype of y
 * @param[in] x vector x
 * @param[out] y matrix y
 */
template <typename vx_t, typename vy_t>
void mat_col_copy(const Eigen::MatrixBase<vx_t>& x, std::vector<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y[0].size());
    for (vy_t& v : y) {
        vec_copy(x, v);
    }
}

/**
 * @brief Calculate Y = a * X + Y
 *
 * @tparam sa_t datatype of scalar a
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @param[in] a scalar a
 * @param[in] x vector x
 * @param[inout] y vector y containing the result
 */
template <typename sa_t, typename vx_t, typename vy_t>
void axpy(const sa_t& a, const Eigen::MatrixBase<vx_t>& x,
    Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        y[i] += a * x[i];
    }
}

/**
 * @brief Calculate Z = a * X + b * Y + Z
 *
 * @tparam sa_t datatype of scalar a
 * @tparam vx_t datatype of vector x
 * @tparam sb_t datatype of scalar b
 * @tparam vy_y datatype of vector y
 * @tparam vz_t datatype of vector z
 * @param[in] a scalar a
 * @param[in] x vector x
 * @param[in] b scalar b
 * @param[in] y vector y
 * @param[inout] z vector z
 */
template <typename sa_t, typename vx_t, typename sb_t, typename vy_y, typename vz_t>
void axpbypz(const sa_t& a, const Eigen::MatrixBase<vx_t>& x,
    const sb_t& b, const Eigen::MatrixBase<vy_y>& y,
    Eigen::MatrixBase<vz_t>& z)
{
    size_t dim = x.size();
    assert(dim == y.size());
    assert(dim == z.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        z[i] += a * x[i] + b * y[i];
    }
}

/**
 * @brief Calculate OUT = a * X + b * Y + c * Z
 *
 * @tparam sa_t datatype of scalar a
 * @tparam vx_t datatype of vector x
 * @tparam sb_t datatype of scalar b
 * @tparam vy_y datatype of vector y
 * @tparam sc_t datatype of scalar c
 * @tparam vz_t datatype of vector z
 * @tparam vout_t datatype of vector out
 * @param[in] a scalar a
 * @param[in] x vector x
 * @param[in] b scalar b
 * @param[in] y vector y
 * @param[in] c scalar c
 * @param[in] z vector z
 * @param[out] out vector out
 */
template <typename sa_t, typename vx_t,
    typename sb_t, typename vy_y,
    typename sc_t, typename vz_t,
    typename vout_t>
void axpbypcz_out(const sa_t& a, const Eigen::MatrixBase<vx_t>& x,
    const sb_t& b, const Eigen::MatrixBase<vy_y>& y,
    const sc_t& c, const Eigen::MatrixBase<vz_t>& z,
    Eigen::MatrixBase<vout_t>& out)
{
    size_t dim = x.size();
    assert(dim == y.size());
    assert(dim == z.size());
    assert(dim == out.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        out[i] = a * x[i] + b * y[i] + c * z[i];
    }
}

/**
 * @brief Calculate prod = x.conj() * m * y for diagonal matrix m
 *
 * @tparam vx_t datatype of vector x
 * @tparam vm_t datatype of diagonal matrix m
 * @tparam vy_t datatype of vector y
 * @param[in] x vector x
 * @param[in] m diagonal matrix m
 * @param[in] y vector y
 * @return the product
 */
template <typename vx_t, typename vm_t, typename vy_t>
complex_t xmy_diag(const Eigen::MatrixBase<vx_t>& x,
    const Eigen::MatrixBase<vm_t>& m,
    const Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == m.size());
    assert(dim == y.size());
    double real_part = 0.0;
    double imag_part = 0.0;
    complex_t tmp = 0.0;
#pragma omp parallel for reduction(+ : real_part, imag_part) private(tmp) schedule(static)
    for (size_t i = 0; i < dim; ++i) {
        tmp = std::conj(x[i]) * m[i] * y[i];
        real_part += tmp.real();
        imag_part += tmp.imag();
    }
    return { real_part, imag_part };
}

/* -------- Generated with gen_cpp.py -------- */

/**
 * @brief Calculate Z = X + Y
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @tparam vz_t datatype of vector z
 * @param[in] x left operand
 * @param[in] y right operand
 * @param[out] z result
 */
template <typename vx_t, typename vy_t, typename vz_t>
void vec_add_vec(const Eigen::MatrixBase<vx_t>& x, const Eigen::MatrixBase<vy_t>& y,
    Eigen::MatrixBase<vz_t>& z)
{
    size_t dim = x.size();
    assert(dim == y.size());
    assert(dim == z.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        z[i] = x[i] + y[i];
    }
}

/**
 * @brief Calculate Z = X - Y
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @tparam vz_t datatype of vector z
 * @param[in] x left operand
 * @param[in] y right operand
 * @param[out] z result
 */
template <typename vx_t, typename vy_t, typename vz_t>
void vec_sub_vec(const Eigen::MatrixBase<vx_t>& x, const Eigen::MatrixBase<vy_t>& y,
    Eigen::MatrixBase<vz_t>& z)
{
    size_t dim = x.size();
    assert(dim == y.size());
    assert(dim == z.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        z[i] = x[i] - y[i];
    }
}

/**
 * @brief Calculate Z = X * Y
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @tparam vz_t datatype of vector z
 * @param[in] x left operand
 * @param[in] y right operand
 * @param[out] z result
 */
template <typename vx_t, typename vy_t, typename vz_t>
void vec_mul_vec(const Eigen::MatrixBase<vx_t>& x, const Eigen::MatrixBase<vy_t>& y,
    Eigen::MatrixBase<vz_t>& z)
{
    size_t dim = x.size();
    assert(dim == y.size());
    assert(dim == z.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        z[i] = x[i] * y[i];
    }
}

/**
 * @brief Calculate Z = X / Y
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @tparam vz_t datatype of vector z
 * @param[in] x left operand
 * @param[in] y right operand
 * @param[out] z result
 */
template <typename vx_t, typename vy_t, typename vz_t>
void vec_div_vec(const Eigen::MatrixBase<vx_t>& x, const Eigen::MatrixBase<vy_t>& y,
    Eigen::MatrixBase<vz_t>& z)
{
    size_t dim = x.size();
    assert(dim == y.size());
    assert(dim == z.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        z[i] = x[i] / y[i];
    }
}

/**
 * @brief Calculate Y = X + a
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @tparam vy_t datatype of vector y
 * @param[in] x left operand
 * @param[in] a right operand
 * @param[out] y result
 */
template <typename vx_t, typename sa_t, typename vy_t>
void vec_add_sca(const Eigen::MatrixBase<vx_t>& x, const sa_t& a,
    Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        y[i] = x[i] + a;
    }
}

/**
 * @brief Calculate Y = X - a
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @tparam vy_t datatype of vector y
 * @param[in] x left operand
 * @param[in] a right operand
 * @param[out] y result
 */
template <typename vx_t, typename sa_t, typename vy_t>
void vec_sub_sca(const Eigen::MatrixBase<vx_t>& x, const sa_t& a,
    Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        y[i] = x[i] - a;
    }
}

/**
 * @brief Calculate Y = X * a
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @tparam vy_t datatype of vector y
 * @param[in] x left operand
 * @param[in] a right operand
 * @param[out] y result
 */
template <typename vx_t, typename sa_t, typename vy_t>
void vec_mul_sca(const Eigen::MatrixBase<vx_t>& x, const sa_t& a,
    Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        y[i] = x[i] * a;
    }
}

/**
 * @brief Calculate Y = X / a
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @tparam vy_t datatype of vector y
 * @param[in] x left operand
 * @param[in] a right operand
 * @param[out] y result
 */
template <typename vx_t, typename sa_t, typename vy_t>
void vec_div_sca(const Eigen::MatrixBase<vx_t>& x, const sa_t& a,
    Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        y[i] = x[i] / a;
    }
}

/**
 * @brief Calculate X += Y
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @param[inout] x left operand
 * @param[in] y right operand
 */
template <typename vx_t, typename vy_t>
void vec_self_add_vec(Eigen::MatrixBase<vx_t>& x, const Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        x[i] += y[i];
    }
}

/**
 * @brief Calculate X -= Y
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @param[inout] x left operand
 * @param[in] y right operand
 */
template <typename vx_t, typename vy_t>
void vec_self_sub_vec(Eigen::MatrixBase<vx_t>& x, const Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        x[i] -= y[i];
    }
}

/**
 * @brief Calculate X *= Y
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @param[inout] x left operand
 * @param[in] y right operand
 */
template <typename vx_t, typename vy_t>
void vec_self_mul_vec(Eigen::MatrixBase<vx_t>& x, const Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        x[i] *= y[i];
    }
}

/**
 * @brief Calculate X /= Y
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @param[inout] x left operand
 * @param[in] y right operand
 */
template <typename vx_t, typename vy_t>
void vec_self_div_vec(Eigen::MatrixBase<vx_t>& x, const Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size());
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        x[i] /= y[i];
    }
}

/**
 * @brief Calculate X += a
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @param[inout] x left operand
 * @param[in] a right operand
 */
template <typename vx_t, typename sa_t>
void vec_self_add_sca(Eigen::MatrixBase<vx_t>& x, const sa_t& a)
{
    size_t dim = x.size();
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        x[i] += a;
    }
}

/**
 * @brief Calculate X -= a
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @param[inout] x left operand
 * @param[in] a right operand
 */
template <typename vx_t, typename sa_t>
void vec_self_sub_sca(Eigen::MatrixBase<vx_t>& x, const sa_t& a)
{
    size_t dim = x.size();
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        x[i] -= a;
    }
}

/**
 * @brief Calculate X *= a
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @param[inout] x left operand
 * @param[in] a right operand
 */
template <typename vx_t, typename sa_t>
void vec_self_mul_sca(Eigen::MatrixBase<vx_t>& x, const sa_t& a)
{
    size_t dim = x.size();
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        x[i] *= a;
    }
}

/**
 * @brief Calculate X /= a
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @param[inout] x left operand
 * @param[in] a right operand
 */
template <typename vx_t, typename sa_t>
void vec_self_div_sca(Eigen::MatrixBase<vx_t>& x, const sa_t& a)
{
    size_t dim = x.size();
#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        x[i] /= a;
    }
}


/**
 * @brief Calculate Y += X.cwiseAbs2().
 * Specifically, for each element i: Y[i] += |X[i]|^2
 *
 * @tparam vx_t datatype of vector x (e.g., Eigen::VectorXcd)
 * @tparam vy_t datatype of vector y (e.g., Eigen::VectorXcd or Eigen::VectorXd)
 * @param[in] x The input vector whose elements' absolute squares are to be added.
 * @param[inout] y The output vector to which the absolute squares are added.
 */
template <typename vx_t, typename vy_t>
void vec_abs2_add(const Eigen::MatrixBase<vx_t>& x, Eigen::MatrixBase<vy_t>& y)
{
    size_t dim = x.size();
    assert(dim == y.size()); // Ensure dimensions match

#pragma omp parallel for simd
    for (size_t i = 0; i < dim; ++i) {
        // For complex types, std::norm(x) is just |x|^2.
        // For real types, std::norm(x) is just x*x.
        y[i] += std::norm(x[i]);
        // y[i] += (std::conj(x[i]) * x[i]).real();
    }
}
} // namespace tbplas::tbpm

#endif