/**
 * @brief Header shared by dense and csr based on MKL.
 *
 * @author Junyi Wang (jywang)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_TBPM_MATH_AOCL_WRAPPER_H
#define TBPLAS_TBPM_MATH_AOCL_WRAPPER_H
#include <aoclsparse.h>
#include <blis.hh>
#include <complex>
#include <iostream>
#include <type_traits>
#include <vector>
namespace tbplas::math {

namespace detail_aocl {
    template <class scalar_t>
    struct typemap;
    template <>
    struct typemap<double> {
        using type = double;
    };
    template <>
    struct typemap<std::complex<double>> {
        using type = aoclsparse_double_complex;
    };
}

namespace aocl {
    /// Get the corresponding MKL C API type from the C++ type `scalar_t`
    template <class scalar_t>
    using type = typename tbplas::math::detail_aocl::typemap<scalar_t>::type;

    /// #include <type_traits>
    template <typename T>
    struct always_false : std::false_type { };

    template <typename idx_t>
    bool trans_array_template_to_aoclsparse_int(const std::vector<idx_t>& a, aoclsparse_int*& a_)
    {
        // std::cout << sizeof(aoclsparse_int) << std::endl;
        assert(sizeof(a[0]) <= sizeof(aoclsparse_int));
        bool flag = false;
        if constexpr (sizeof(a[0]) < sizeof(aoclsparse_int)) {
            flag = true;
            a_ = (aoclsparse_int*)malloc(a.size() * sizeof(aoclsparse_int));
            for (int i = 0; i < a.size(); ++i) {
                aoclsparse_int temp = a[i];
                a_[i] = temp;
            }
        } else {
            a_ = reinterpret_cast<aoclsparse_int*>((idx_t*)a.data());
        }
        return flag;
    }

    // static constexpr auto call要求`call` 是常量表达式，如zaxpy，在编译时就可以被计算
    // static void call 不要求在编译时进行求值，可以带参数

    template <class vx_t, class vy_t>
    struct sparse_mv;
    template <>
    struct sparse_mv<double, double> {
        static constexpr auto call = aoclsparse_dmv;
    };
    template <>
    struct sparse_mv<std::complex<double>, std::complex<double>> {
        static constexpr auto call = aoclsparse_zmv;
    };
    // aoclsparse_status aoclsparse_dmv(aoclsparse_operation op, const double *alpha, aoclsparse_matrix A, const aoclsparse_mat_descr descr, const double *x, const double *beta, double *y)
    // aoclsparse_status aoclsparse_zmv(aoclsparse_operation op, const aoclsparse_double_complex *alpha, aoclsparse_matrix A, const aoclsparse_mat_descr descr, const aoclsparse_double_complex *x, const aoclsparse_double_complex *beta, aoclsparse_double_complex *y)

    template <class scalar_t>
    struct sparse_export_csr;
    template <>
    struct sparse_export_csr<double> {
        static constexpr auto call = aoclsparse_export_dcsr;
    };
    template <>
    struct sparse_export_csr<std::complex<double>> {
        static constexpr auto call = aoclsparse_export_zcsr;
    };
    // aoclsparse_status aoclsparse_export_dcsr(const aoclsparse_matrix mat, aoclsparse_index_base *base, aoclsparse_int *m, aoclsparse_int *n, aoclsparse_int *nnz, aoclsparse_int **row_ptr, aoclsparse_int **col_ind, double **val)
    // aoclsparse_status aoclsparse_export_zcsr(const aoclsparse_matrix mat, aoclsparse_index_base *base, aoclsparse_int *m, aoclsparse_int *n, aoclsparse_int *nnz, aoclsparse_int **row_ptr, aoclsparse_int **col_ind, aoclsparse_double_complex **val)
    template <class scalar_t>
    struct sparse_create_csr;
    template <>
    struct sparse_create_csr<double> {
        static constexpr auto call = aoclsparse_create_dcsr;
    };
    template <>
    struct sparse_create_csr<std::complex<double>> {
        static constexpr auto call = aoclsparse_create_zcsr;
    };
    // aoclsparse_status aoclsparse_create_dcsr(aoclsparse_matrix *mat, aoclsparse_index_base base, aoclsparse_int M, aoclsparse_int N, aoclsparse_int nnz, aoclsparse_int *row_ptr, aoclsparse_int *col_idx, double *val)
    // aoclsparse_status aoclsparse_create_zcsr(aoclsparse_matrix *mat, aoclsparse_index_base base, aoclsparse_int M, aoclsparse_int N, aoclsparse_int nnz, aoclsparse_int *row_ptr, aoclsparse_int *col_idx, aoclsparse_double_complex *val)

    template <class scalar_t>
    struct dotc;
    // template< typename T, typename TR >TR dot(int64_t n,T const *x, int64_t incx,T const *y, int64_t incy )
    template <>
    struct dotc<double> {
        static constexpr double (*call)(int64_t, const double*, int64_t, const double*, int64_t) = blis::dot<double, double>;
    };
    // template< typename T >T dotc(int64_t n,T const *x, int64_t incx,T const *y, int64_t incy )
    template <>
    struct dotc<std::complex<double>> {
        static constexpr std::complex<double> (*call)(int64_t, const std::complex<double>*, int64_t, const std::complex<double>*, int64_t) = blis::dotc<std::complex<double>>;
    };
    // template< typename T>void axpy(int64_t n,T alpha,T const *x, int64_t incx,T *y, int64_t incy)

    // template< typename T >void copy(int64_t n,T const *x, int64_t incx,T *y, int64_t incy )

    // template< typename T >T dotu(int64_t n,T const *x, int64_t incx,T const *y, int64_t incy )

    // template< typename T >real_type<T>nrm2(int64_t n,T const * x, int64_t incx )

    // template< typename TA, typename TB >void scal(int64_t n,TA alpha,TB* x, int64_t incx )

    template <typename scalar_t>
    inline void print_e(const scalar_t& i)
    {
        std::cout << i;
    };
    template <>
    inline void print_e(const aoclsparse_double_complex& i)
    {
        std::cout << i.real << "+" << i.imag << "i";
    };

    template <typename scalar_t, typename value_t>
    inline void assignment(scalar_t& i, value_t a)
    {
        i = a;
    };
    template <>
    inline void assignment<aoclsparse_double_complex, double>(aoclsparse_double_complex& i, double a)
    {
        i.real = a;
        i.imag = 0;
    };
    template <>
    inline void assignment<aoclsparse_double_complex, std::complex<double>>(aoclsparse_double_complex& i, std::complex<double> a)
    {
        i.real = a.real();
        i.imag = a.imag();
    };
    template <>
    inline void assignment<std::complex<double>, aoclsparse_double_complex>(std::complex<double>& i, aoclsparse_double_complex a)
    {
        i = std::complex<double>(a.real, a.imag);
    };
}
} // namespace tbplas::math
#endif