/**
 * @brief Classes and functions for sparse matrix in CSR format.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 *
 * @note
 *   1. Abbreviation for composite vector operations: p(lus), s(ubstract), m(atrix).
 *      We do not use add for + since a indicates a scalar operand.
 */

#ifndef TBPLAS_TBPM_CSR_H
#define TBPLAS_TBPM_CSR_H

#include <cassert>
#include <iostream>
#include <vector>

#include <eigen3/Eigen/Dense>

#include "base/datatypes.h"

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

/**
 * @brief Collections of sizes and pointers of arrays in CSR square matrix,
 * for GPU programming.
 * 
 * @tparam idx_t datatype of indices
 * @tparam value_t template parameter for Eigen::MatrixBase
 */
template <typename idx_t, typename value_t>
class CSRMatrixSquareMeta {
public:
    idx_t dim = 0;
    size_t indptr_size = 0;
    size_t indices_size = 0;
    size_t values_size = 0;
    const idx_t* indptr_data = nullptr;
    const idx_t* indices_data = nullptr;
    const typename value_t::Scalar* values_data = nullptr;
};

/**
 * @brief Square sparse matrix in CSR format.
 *
 * @tparam idx_t datatype of indices
 * @tparam value_t template parameter for Eigen::MatrixBase
 *
 * @note For reducing memory usage, this class does not contain any copy of
 * matrix data, just pointers to them. So all copies will be shallow copy.
 */
template <typename idx_t, typename value_t>
class CSRMatrixSquare {
private:
    /// @brief Pointer to the vector of offsets, namely 'outerIndex'
    /// of eigen sparse matrix
    /// @note Eigen outIndex drops the last element of indptr. DO NOT use it
    /// to instantiate this class.
    const std::vector<idx_t>* indptr_ = nullptr;

    /// @brief Pointer to the vector of column indices, namely 'innerIndex'
    /// of eigen sparse matrix
    const std::vector<idx_t>* indices_ = nullptr;

    /// @brief Pointer to the eigen vector of the values
    const Eigen::MatrixBase<value_t>* values_ = nullptr;

    /// @brief Dimension of the square matrix
    idx_t dim_ = 0;

public:
    /**
     * @brief Construct a new CSRMatrixSquare object
     *
     * @param[in] indptr vector of the offsets
     * @param[in] indices vector of column indices
     * @param[in] values vector of the values
     */
    explicit CSRMatrixSquare(const std::vector<idx_t>& indptr,
        const std::vector<idx_t>& indices,
        const Eigen::MatrixBase<value_t>& values)
        : indptr_(&indptr)
        , indices_(&indices)
        , values_(&values)
        , dim_(indptr.size() - 1)
    {
        assert(indptr.size() >= 2);
        assert(indices.size() >= 1);
        assert(indices.size() == values.size());
    }

    CSRMatrixSquare(const CSRMatrixSquare& other)
        : indptr_(other.indptr_)
        , indices_(other.indices_)
        , values_(other.values_)
        , dim_(other.dim_)
    {
    }

    CSRMatrixSquare(CSRMatrixSquare&& other) noexcept
        : indptr_(other.indptr_)
        , indices_(other.indices_)
        , values_(other.values_)
        , dim_(other.dim_)
    {
        other.indptr_ = nullptr;
        other.indices_ = nullptr;
        other.values_ = nullptr;
        other.dim_ = 1;
    }

    CSRMatrixSquare& operator=(const CSRMatrixSquare& other)
    {
        if (this != &other) {
            indptr_(other.indptr_);
            indices_(other.indices_);
            values_(other.values_);
            dim_(other.dim_);
        }
        return *this;
    }

    CSRMatrixSquare& operator=(CSRMatrixSquare&& other) noexcept
    {
        if (this != &other) {
            indptr_(other.indptr_);
            indices_(other.indices_);
            values_(other.values_);
            dim_(other.dim_);
            other.indptr_ = nullptr;
            other.indices_ = nullptr;
            other.values_ = nullptr;
            other.dim_ = 1;
        }
        return *this;
    }

    ~CSRMatrixSquare()
    {
        indptr_ = nullptr;
        indices_ = nullptr;
        values_ = nullptr;
        dim_ = 1;
    }

    /**
     * @brief Print non-zero matrix elements.
     *
     */
    void print() const
    {
        const std::vector<idx_t>& indptr = *indptr_;
        const std::vector<idx_t>& indices = *indices_;
        const Eigen::MatrixBase<value_t>& values = *values_;
        idx_t i = 0; // row index
        idx_t j = 0; // offset of indices and values
        idx_t k = 0; // column index
        for (i = 0; i < dim_; ++i) {
            for (j = indptr[i]; j < indptr[i + 1]; ++j) {
                k = indices[j];
                std::cout << i << " " << k << " " << values(j) << "\n";
            }
        }
    }

    /**
     * @brief Get the dim attribute
     *
     * @return dimension of the matrix
     */
    [[nodiscard]] idx_t get_dim() const { return dim_; }

    /**
     * @brief Get sizes and pointer of attribute arrays.
     * 
     * @return Meta info containing the sizes and pointers
     */
    [[nodiscard]]  CSRMatrixSquareMeta<idx_t, value_t> get_meta() const {
        CSRMatrixSquareMeta<idx_t, value_t> meta;
        meta.dim = dim_;
        meta.indptr_size = indptr_->size();
        meta.indices_size = indices_->size();
        meta.values_size = values_->derived().size();
        meta.indptr_data = indptr_->data();
        meta.indices_data = indices_->data();
        meta.values_data = values_->derived().data();
        return meta;
    }

    /**
     * @brief Calculate Y = mat * X
     *
     * @tparam vx_t datatype for vector x
     * @tparam vy_t datatype for vector y
     * @param[in] x incoming vector
     * @param[out] y result vector
     */
    template <typename vx_t, typename vy_t>
    void mv(const Eigen::MatrixBase<vx_t>& x,
        Eigen::MatrixBase<vy_t>& y) const
    {
        assert(dim_ == x.size());
        assert(dim_ == y.size());
        const std::vector<idx_t>& indptr = *indptr_;
        const std::vector<idx_t>& indices = *indices_;
        const Eigen::MatrixBase<value_t>& values = *values_;
        idx_t i = 0; // row index
        idx_t j = 0; // offset of indices and values
        idx_t k = 0; // column index
#pragma omp parallel for private(j, k)
        for (i = 0; i < dim_; ++i) {
            y[i] = 0.0;
            for (j = indptr[i]; j < indptr[i + 1]; ++j) {
                k = indices[j];
                y[i] += values[j] * x[k]; // Different j and k may access the same i so use +=.
            }
        }
    }

    /**
     * @brief Calculate Y = a * mat * X
     *
     * @tparam sa_t datatype of factor a
     * @tparam vx_t datatype for vector x
     * @tparam vy_t datatype for vector y
     * @param[in] a the factor
     * @param[in] x incoming vector
     * @param[out] y result vector
     */
    template <typename sa_t, typename vx_t, typename vy_t>
    void amv(const sa_t& a, const Eigen::MatrixBase<vx_t>& x,
        Eigen::MatrixBase<vy_t>& y) const
    {
        assert(dim_ == x.size());
        assert(dim_ == y.size());
        const std::vector<idx_t>& indptr = *indptr_;
        const std::vector<idx_t>& indices = *indices_;
        const Eigen::MatrixBase<value_t>& values = *values_;
        idx_t i = 0; // row index
        idx_t j = 0; // offset of indices and values
        idx_t k = 0; // column index
#pragma omp parallel for private(j, k)
        for (i = 0; i < dim_; ++i) {
            y[i] = 0.0;
            for (j = indptr[i]; j < indptr[i + 1]; ++j) {
                k = indices[j];
                y[i] += a * values[j] * x[k];
            }
        }
    }

    /**
     * @brief  Calculate Y = a * mat * X - Y
     *
     * @tparam sa_t datatype of factor a
     * @tparam vx_t datatype for vector x
     * @tparam vy_t datatype for vector y
     * @param[in] a the factor
     * @param[in] x incoming vector
     * @param[inout] y result vector
     */
    template <typename sa_t, typename vx_t, typename vy_t>
    void amxsy(const sa_t& a, const Eigen::MatrixBase<vx_t>& x,
        Eigen::MatrixBase<vy_t>& y) const
    {
        assert(dim_ == x.size());
        assert(dim_ == y.size());
        const std::vector<idx_t>& indptr = *indptr_;
        const std::vector<idx_t>& indices = *indices_;
        const Eigen::MatrixBase<value_t>& values = *values_;
        idx_t i = 0; // row index
        idx_t j = 0; // offset of indices and values
        idx_t k = 0; // column index
#pragma omp parallel for private(j, k)
        for (i = 0; i < dim_; ++i) {
            y[i] = -y[i];
            for (j = indptr[i]; j < indptr[i + 1]; ++j) {
                k = indices[j];
                y[i] += a * values[j] * x[k];
            }
        }
    }

    /**
     * @brief Calculate OUT = a * mat * X - Y
     *
     * @tparam sa_t datatype of factor a
     * @tparam vx_t datatype for vector x
     * @tparam vy_t datatype for vector y
     * @tparam vout_t datatype for vector out
     * @param[in] a the factor
     * @param[in] x incoming vector y
     * @param[in] y incoming vector y
     * @param[out] out result vector
     */
    template <typename sa_t, typename vx_t, typename vy_t, typename vout_t>
    void amxsy_out(const sa_t& a, const Eigen::MatrixBase<vx_t>& x, const Eigen::MatrixBase<vy_t>& y,
        Eigen::MatrixBase<vout_t>& out) const
    {
        assert(dim_ == x.size());
        assert(dim_ == y.size());
        assert(dim_ == out.size());
        const std::vector<idx_t>& indptr = *indptr_;
        const std::vector<idx_t>& indices = *indices_;
        const Eigen::MatrixBase<value_t>& values = *values_;
        idx_t i = 0; // row index
        idx_t j = 0; // offset of indices and values
        idx_t k = 0; // column index
#pragma omp parallel for private(j, k)
        for (i = 0; i < dim_; ++i) {
            out[i] = -y[i];
            for (j = indptr[i]; j < indptr[i + 1]; ++j) {
                k = indices[j];
                out[i] += a * values[j] * x[k];
            }
        }
    }

    /**
     * @brief Calculate prod = x.conj() * mat * y
     *
     * @tparam vx_t datatype of vector x
     * @tparam vy_t datatype of vector y
     * @param[in] x vector x
     * @param[in] y vector y
     * @return the product
     */
    template <typename vx_t, typename vy_t>
    complex_t xmy(const Eigen::MatrixBase<vx_t>& x,
        const Eigen::MatrixBase<vy_t>& y) const
    {
        assert(dim_ == x.size());
        assert(dim_ == y.size());
        const std::vector<idx_t>& indptr = *indptr_;
        const std::vector<idx_t>& indices = *indices_;
        const Eigen::MatrixBase<value_t>& values = *values_;
        idx_t i = 0; // row index
        idx_t j = 0; // offset of indices and values
        idx_t k = 0; // column index
        double real_part = 0.0;
        double imag_part = 0.0;
        complex_t tmp = 0.0;
#pragma omp parallel for private(j, k, tmp) reduction(+ : real_part, imag_part)
        for (i = 0; i < dim_; ++i) {
            for (j = indptr[i]; j < indptr[i + 1]; ++j) {
                k = indices[j];
                tmp = std::conj(x[i]) * values[j] * y[k];
                real_part += tmp.real();
                imag_part += tmp.imag();
            }
        }
        return { real_part, imag_part };
    }
};
} // namespace tbplas::tbpm

#endif