/**
 * @brief Classes and functions for sparse matrix in CSR format based on MKL C API.
 *
 * @author Junyi Wang (jywang)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_TBPM_CSR_H
#define TBPLAS_TBPM_CSR_H
#include <cassert>
#include <iostream>
#include <type_traits>
#include <vector>

#include <eigen3/Eigen/Dense>

#include "base/datatypes.h"
#include "mkl_wrapper.h"

#define DESC_GENERAL_FULL_NONUNIT                                               \
    {                                                                           \
        SPARSE_MATRIX_TYPE_GENERAL, SPARSE_FILL_MODE_FULL, SPARSE_DIAG_NON_UNIT \
    }
#define DESC_HERMITIAN_UPPER_NONUNIT                                               \
    {                                                                              \
        SPARSE_MATRIX_TYPE_HERMITIAN, SPARSE_FILL_MODE_UPPER, SPARSE_DIAG_NON_UNIT \
    }
using namespace tbplas::math::mkl;
namespace tbplas::tbpm {
using tbplas::base::complex_t;

/**
 * @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 {
    using Scalar = typename value_t::Scalar;
    using math_scalar_t = tbplas::math::mkl::type<Scalar>;
    template <class scalar_t>
    using type = tbplas::math::mkl::type<scalar_t>;

private:
    MKL_INT* indices_ = nullptr;
    MKL_INT* indptr_ = nullptr;
    bool free_flag_ = false;
    // A存储了指向数据indices_、indptr_、values_的指针，在释放A前不可释放指针
    sparse_matrix_t A_;
    matrix_descr descr_;
    // rows=dim_
    MKL_INT rows_ = 0, cols_ = 0;

    // 销毁稀疏矩阵的函数模板
    void destroy()
    {
        if (A_ != nullptr) {
            mkl_sparse_destroy(A_);
            A_ = nullptr;
        }
    }

public:
    // WITH_FULL_CSR option is removed. But we keep this snippet for reference.
    // #if defined(WITH_FULL_CSR)
    //     explicit CSRMatrixSquare(const std::vector<idx_t>& indptr,
    //         const std::vector<idx_t>& indices,
    //         const Eigen::MatrixBase<value_t>& values, const matrix_descr descr = DESC_GENERAL_FULL_NONUNIT)
    // #else
    //     explicit CSRMatrixSquare(const std::vector<idx_t>& indptr,
    //         const std::vector<idx_t>& indices,
    //         const Eigen::MatrixBase<value_t>& values, const matrix_descr descr = DESC_HERMITIAN_UPPER_NONUNIT)
    // #endif
    explicit CSRMatrixSquare(const std::vector<idx_t>& indptr,
        const std::vector<idx_t>& indices,
        const Eigen::MatrixBase<value_t>& values, const matrix_descr descr = DESC_GENERAL_FULL_NONUNIT)
        : rows_(indptr.size() - 1)
        , cols_(*std::max_element(indices.begin(), indices.end()) + 1)
        , descr_(descr)

    {
        assert(indptr.size() >= 2);
        assert(indices.size() >= 1);
        assert(indices.size() == values.size());

        free_flag_ = trans_array_template_to_MKL_INT(indptr, indptr_);
        trans_array_template_to_MKL_INT(indices, indices_);

        // 匹配
        // creat mat
        auto* ptr = (math_scalar_t*)values.derived().data();
        sparse_status_t status = sparse_create_csr<Scalar>::call(
            &A_,
            SPARSE_INDEX_BASE_ZERO,
            rows_,
            cols_,
            indptr_,
            indptr_ + 1,
            indices_,
            ptr);
        if (status != SPARSE_STATUS_SUCCESS) {
            std::cerr << "Failed to create sparse matrix" << std::endl;
            std::exit(-1);
        }
    }

    CSRMatrixSquare(const CSRMatrixSquare& other)
        : indptr_(other.indptr_)
        , indices_(other.indices_)
        , free_flag_(other.free_flag_)
        , A_(other.A_)
        , descr_(other.descr_)
        , rows_(other.rows_)
        , cols_(other.cols_)
    {
    }

    CSRMatrixSquare(CSRMatrixSquare&& other) noexcept
        : indptr_(other.indptr_)
        , indices_(other.indices_)
        , free_flag_(other.free_flag_)
        , A_(other.A_)
        , descr_(other.descr_)
        , rows_(other.rows_)
        , cols_(other.cols_)
    {
        mkl_sparse_destroy(other.A_);
        if (other.free_flag_) {
            free(other.indptr_);
            free(other.indices_);
            other.indptr_ = nullptr;
            other.indices_ = nullptr;
            other.free_flag_ = false;
        }
        other.descr_ = DESC_GENERAL_FULL_NONUNIT;
        other.rows_ = 1;
        other.cols_ = 1;
    }

    CSRMatrixSquare& operator=(const CSRMatrixSquare& other)
    {
        if (this != &other) {
#if defined(__GNUC__)
            // Workaround for GCC as x_(other.x_) does not compile.
            indptr_ = other.indptr_;
            indices_ = other.indices_;
            free_flag_ = other.free_flag_;
            A_ = other.A_;
            descr_ = other.descr_;
            rows_ = other.rows_;
            cols_ = other.cols_;
#else
            indptr_(other.indptr_);
            indices_(other.indices_);
            free_flag_(other.free_flag_);
            A_(other.A_);
            descr_(other.descr_);
            rows_(other.rows_);
            cols_(other.cols_);
#endif
        }
        return *this;
    }

    CSRMatrixSquare& operator=(CSRMatrixSquare&& other) noexcept
    {
        if (this != &other) {
#if defined(__GNUC__)
            // Workaround for GCC as x_(other.x_) does not compile.
            indptr_ = other.indptr_;
            indices_ = other.indices_;
            free_flag_ = other.free_flag_;
            A_ = other.A_;
            descr_ = other.descr_;
            rows_ = other.rows_;
            cols_ = other.cols_;
#else
            indptr_(other.indptr_);
            indices_(other.indices_);
            free_flag_(other.free_flag_);
            A_(other.A_);
            descr_(other.descr_);
            rows_(other.rows_);
            cols_(other.cols_);
#endif
            mkl_sparse_destroy(other.A_);
            if (other.free_flag_) {
                free(other.indptr_);
                free(other.indices_);
                other.free_flag_ = false;
            }
            other.indptr_ = nullptr;
            other.indices_ = nullptr;
            other.descr_ = DESC_GENERAL_FULL_NONUNIT;
            other.rows_ = 1;
            other.cols_ = 1;
        }
        return *this;
    }

    ~CSRMatrixSquare()
    {
        destroy();
        rows_ = 1;
        cols_ = 1;
        if (free_flag_) {
            free(indices_);
            free(indptr_);
            free_flag_ = false;
        }
        indptr_ = nullptr;
        indices_ = nullptr;
    }

    std::tuple<MKL_INT, MKL_INT> get_matrix_size() const
    {
        return std::make_tuple(rows_, cols_);
    }

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

    void print() const
    {
        sparse_index_base_t indexing;
        MKL_INT rows, cols;
        MKL_INT *indptr, *rows_end, *indices;
        math_scalar_t* values;
        sparse_status_t status = sparse_export_csr<Scalar>::call(
            A_,
            &indexing,
            &rows,
            &cols,
            &indptr,
            &rows_end,
            &indices,
            &values);
        // 无法导出CSR格式的数据
        if (status != SPARSE_STATUS_SUCCESS) {
            std::cerr << "Unable to export data in CSR format" << std::endl;
            assert(status == SPARSE_STATUS_SUCCESS);
        }
        MKL_INT i = 0; // row index
        MKL_INT j = 0; // offset of indices and values
        MKL_INT k = 0; // column index
        for (i = 0; i < rows_; ++i) {
            for (j = indptr[i]; j < indptr[i + 1]; ++j) {
                k = indices[j];
                // 行索引 列索引 元素值
                std::cout << i << " " << k << " ";
                print_e(values[j]);
                std::cout << "\n";
            }
        }
    }

    /**
     * @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(rows_ == x.size());
        assert(rows_ == y.size());

        typedef typename vy_t::Scalar Scalar_y;
        // 也可以不检查x和y元素数据类型是否一致，在编译时sparse_mv报类型不匹配或者无法转换
        // 必须保证A_和x/y的元素数据类型一致，mkl不会检查和报错但不能计算出正确结果
        // typedef typename vx_t::Scalar Scalar_x;
        // static_assert(std::is_same_v<Scalar_x, Scalar_y>, "Scalar types must be the same!");
        math_scalar_t alpha, beta;
        assignment(alpha, 1.0);
        assignment(beta, 0.0);
        math_scalar_t* x_data = (math_scalar_t*)x.derived().data();
        math_scalar_t* y_data = (math_scalar_t*)y.derived().data();
        // Hermitian矩阵填入上三角,计算整体
        sparse_status_t status = sparse_mv<Scalar_y, Scalar>::call(
            SPARSE_OPERATION_NON_TRANSPOSE,
            alpha,
            A_,
            descr_,
            x_data,
            beta,
            y_data);
        if (status != SPARSE_STATUS_SUCCESS) {
            std::cerr << "Failed to calculate Y = mat * X" << std::endl;
            std::exit(-1);
        }
    }

    /**
     * @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(rows_ == x.size());
        assert(rows_ == y.size());
        typedef typename vy_t::Scalar Scalar_y;
        // typedef typename vx_t::Scalar Scalar_x;
        ////编译时检查
        // static_assert(std::is_same_v<Scalar_x, Scalar_y>, "Scalar types must be the same!");
        ////运行时检查
        // assert(typeid(Scalar_x) == typeid(Scalar_y));
        // 下一行写法有问题，assert运行时检查，std::is_same_v返回一个const bool，不符合assert的参数类型
        // assert(std::is_same_v<Scalar_x, Scalar_y>)
        math_scalar_t alpha, beta;
        assignment(alpha, a);
        assignment(beta, 0.0);
        math_scalar_t* x_data = (math_scalar_t*)x.derived().data();
        math_scalar_t* y_data = (math_scalar_t*)y.derived().data();
        sparse_status_t status = sparse_mv<Scalar_y, Scalar>::call(
            SPARSE_OPERATION_NON_TRANSPOSE,
            alpha,
            A_,
            descr_,
            x_data,
            beta,
            y_data);
        if (status != SPARSE_STATUS_SUCCESS) {
            std::cerr << "Failed to calculate Y = a * mat * X" << std::endl;
            std::exit(-1);
        }
    }

    /**
     * @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(rows_ == x.size());
        assert(rows_ == y.size());

        typedef typename vy_t::Scalar Scalar_y;
        math_scalar_t alpha, beta;
        assignment(alpha, a);
        assignment(beta, -1.0);
        math_scalar_t* x_data = (math_scalar_t*)x.derived().data();
        math_scalar_t* y_data = (math_scalar_t*)y.derived().data();
        sparse_status_t status = sparse_mv<Scalar_y, Scalar>::call(
            SPARSE_OPERATION_NON_TRANSPOSE,
            alpha,
            A_,
            descr_,
            x_data,
            beta,
            y_data);
        if (status != SPARSE_STATUS_SUCCESS) {
            std::cerr << "Failed to calculate Y = a * mat * X - Y" << std::endl;
            std::exit(-1);
        }
    }

    /**
     * @brief Calculate Y = a * mat * X + b * Y
     *
     * @tparam sa_t datatype of factor a
     * @tparam vx_t datatype of vector x
     * @tparam sb_t datatype of factor b
     * @tparam vy_t datatype of vector y
     * @param[in] a factor a
     * @param[in] x vector x
     * @param[in] b factor b
     * @param[inout] y result vector
     */
    template <typename sa_t, typename vx_t, typename sb_t, typename vy_t>
    void amxpby(const sa_t& a, const Eigen::MatrixBase<vx_t>& x, const sb_t& b,
        Eigen::MatrixBase<vy_t>& y) const
    {
        assert(rows_ == x.size());
        assert(rows_ == y.size());

        typedef typename vy_t::Scalar Scalar_y;
        math_scalar_t alpha, beta;
        assignment(alpha, a);
        assignment(beta, b);
        math_scalar_t* x_data = (math_scalar_t*)x.derived().data();
        math_scalar_t* y_data = (math_scalar_t*)y.derived().data();
        sparse_status_t status = sparse_mv<Scalar_y, Scalar>::call(
            SPARSE_OPERATION_NON_TRANSPOSE,
            alpha,
            A_,
            descr_,
            x_data,
            beta,
            y_data);
        if (status != SPARSE_STATUS_SUCCESS) {
            std::cerr << "Failed to calculate Y = a * mat * X + b * Y" << std::endl;
            std::exit(-1);
        }
    }

    /**
     * @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
     */
    // out=a*mat*x+0*out, out=-Y+out
    //  cblas_ ? axpby y=a*x+b*y(a,b:标量；x,y:列向量)
    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(rows_ == x.size());
        assert(rows_ == y.size());
        assert(rows_ == out.size());

        // a * mat * X - Y，所有参数的数据类型应相同
        typedef typename vy_t::Scalar Scalar_y;
        math_scalar_t alpha, beta;
        assignment(alpha, a);
        assignment(beta, 0.0);
        math_scalar_t* x_data = (math_scalar_t*)x.derived().data();
        math_scalar_t* y_data = (math_scalar_t*)y.derived().data();
        math_scalar_t* out_data = (math_scalar_t*)out.derived().data();
        //// out=a*mat*x+0*out
        sparse_status_t status = sparse_mv<Scalar_y, Scalar>::call(
            SPARSE_OPERATION_NON_TRANSPOSE,
            alpha,
            A_,
            descr_,
            x_data,
            beta,
            out_data);
        if (status != SPARSE_STATUS_SUCCESS) {
            std::cerr << "Failed to calculate OUT = a * mat * X - Y" << std::endl;
            std::exit(-1);
        }
        // out=-Y+out
        assignment(alpha, -1.0);
        MKL_INT inc = 1;
        axpy_<Scalar>::call(
            &rows_,
            &alpha,
            y_data,
            &inc,
            out_data,
            &inc);

        // 另一种方法
        // out=Y
        // out=ax-out
        // typedef typename vy_t::Scalar Scalar_y;
        // math_scalar_t alpha, beta;
        // assignment(alpha, a);
        // assignment(beta, -1.0);
        // math_scalar_t* x_data = (math_scalar_t*)x.derived().data();
        // math_scalar_t* y_data = (math_scalar_t*)y.derived().data();
        // math_scalar_t* out_data = (math_scalar_t*)out.derived().data();
        // copy<Scalar>::call(
        //     rows_,
        //     y_data,
        //     1,
        //     out_data,
        //     1);
        // sparse_status_t status = sparse_mv<Scalar_y, Scalar>::call(
        //     SPARSE_OPERATION_NON_TRANSPOSE,
        //     alpha,
        //     A_,
        //     descr_,
        //     x_data,
        //     beta,
        //     out_data);
        // if (status != SPARSE_STATUS_SUCCESS) {
        //     std::cerr << "Failed to calculate OUT = a * mat * X - Y" << std::endl;
        //     std::exit(-1);
        // }
    }

    /**
     * @brief Calculate prod = conj(x.trans) * 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
     */
    ////cblas_zdotc_sub:x共轭与y做点积，x\y为复数向量
    template <typename vx_t, typename vy_t>
    complex_t xmy(const Eigen::MatrixBase<vx_t>& x,
        const Eigen::MatrixBase<vy_t>& y) const
    {
        assert(rows_ == x.size());
        assert(rows_ == y.size());
        complex_t prod = 0.0;
        math_scalar_t* tmp;
        tmp = (math_scalar_t*)malloc(rows_ * sizeof(math_scalar_t));
        // tmp=mat*y
        typedef typename vy_t::Scalar Scalar_y;
        math_scalar_t* x_data = (math_scalar_t*)x.derived().data();
        math_scalar_t* y_data = (math_scalar_t*)y.derived().data();
        math_scalar_t alpha, beta;
        assignment(alpha, 1.0);
        assignment(beta, 0.0);
        sparse_status_t status = sparse_mv<Scalar, Scalar_y>::call(
            SPARSE_OPERATION_NON_TRANSPOSE,
            alpha,
            A_,
            descr_,
            y_data,
            beta,
            tmp);
        if (status != SPARSE_STATUS_SUCCESS) {
            std::cerr << "Failed to calculate prod = conj(x.trans) * mat * y" << std::endl;
            std::exit(-1);
        }
        // conj(x.T)*tmp
        math_scalar_t result;
        dotc<Scalar>::call(
            rows_,
            x_data,
            1,
            tmp,
            1,
            &result);
        assignment(prod, result);
        free(tmp);
        return prod;
    }
};

} // namespace tbplas::tbpm

#endif