/**
 * @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_MKL_WRAPPER_H
#define TBPLAS_TBPM_MATH_MKL_WRAPPER_H
// #define FROM_BLAS
#include <complex>
#include <iostream>
#include <vector>
// #define FROM_BLAS
#ifdef FROM_BLAS
#include <mkl_blas.h>
#include <mkl_spblas.h>
#include <mkl_types.h>
#include <mkl_vml_functions.h>
#else
#include <mkl_cblas.h>
#include <mkl_spblas.h>
#include <mkl_types.h>
#include <mkl_vml_functions.h>
#endif
#include <type_traits>
namespace tbplas::math {

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

namespace mkl {
    /// Get the corresponding MKL C API type from the C++ type `scalar_t`
    template <class scalar_t>
    using type = typename tbplas::math::detail::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_MKL_INT(const std::vector<idx_t>& a, MKL_INT*& a_)
    {
        assert(sizeof(a[0]) <= sizeof(MKL_INT));
        bool flag = false;
        if constexpr (sizeof(a[0]) < sizeof(MKL_INT)) {
            flag = true;
            a_ = (MKL_INT*)malloc(a.size() * sizeof(MKL_INT));
            for (int i = 0; i < a.size(); ++i) {
                MKL_INT temp = a[i];
                a_[i] = temp;
            }
        } else {
            a_ = reinterpret_cast<MKL_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 = mkl_sparse_d_mv;
    };
    template <>
    struct sparse_mv<std::complex<double>, std::complex<double>> {
        static constexpr auto call = mkl_sparse_z_mv;
    };

    template <class scalar_t>
    struct sparse_export_csr;
    template <>
    struct sparse_export_csr<double> {
        static constexpr auto call = mkl_sparse_d_export_csr;
    };
    template <>
    struct sparse_export_csr<std::complex<double>> {
        static constexpr auto call = mkl_sparse_z_export_csr;
    };

    template <class scalar_t>
    struct sparse_create_csr;
    template <>
    struct sparse_create_csr<double> {
        static constexpr auto call = mkl_sparse_d_create_csr;
    };
    template <>
    struct sparse_create_csr<std::complex<double>> {
        static constexpr auto call = mkl_sparse_z_create_csr;
    };

    template <class scalar_t>
    struct axpy_;
#ifdef FROM_BLAS
    //(const MKL_INT *n, const double *alpha, const double *x, const MKL_INT *incx, double *y, const MKL_INT *incy)
    template <>
    struct axpy_<double> {
        static constexpr auto call = daxpy;
    };
    //(const MKL_INT* n, const MKL_Complex16* alpha, const MKL_Complex16* x, const MKL_INT* incx, MKL_Complex16* y, const MKL_INT* incy)
    template <>
    struct axpy_<std::complex<double>> {
        static constexpr auto call = zaxpy;
    };
#else
    template <>
    struct axpy_<double> {
        static void call(const MKL_INT* N, const double* alpha, const double* X,
            const MKL_INT* incX, double* Y, const MKL_INT* incY)
        {
            cblas_daxpy(*N, *alpha, X, *incX, Y, *incY);
        }
    };
    template <>
    struct axpy_<std::complex<double>> {
        static void call(const MKL_INT* N, const MKL_Complex16* alpha, const MKL_Complex16* X,
            const MKL_INT* incX, MKL_Complex16* Y, const MKL_INT* incY)
        {
            cblas_zaxpy(*N, alpha, X, *incX, Y, *incY);
        }
    };
#endif

    template <class scalar_t>
    struct axpby_;
#ifdef FROM_BLAS
    //(const MKL_INT *n, const double *alpha, const double *x, const MKL_INT *incx,const double* beta ,double *y, const MKL_INT *incy)
    template <>
    struct axpby_<double> {
        static constexpr auto call = daxpby;
    };
    //(const MKL_INT* n, const MKL_Complex16* alpha, const MKL_Complex16* x, const MKL_INT* incx,const MKL_Complex16* beta, MKL_Complex16* y, const MKL_INT* incy)
    template <>
    struct axpby_<std::complex<double>> {
        static constexpr auto call = zaxpby;
    };
#else
    template <>
    struct axpby_<double> {
        static void call(const MKL_INT* N, const double* alpha, const double* X,
            const MKL_INT* incX, const double* beta, double* Y, const MKL_INT* incY)
        {
            cblas_daxpby(*N, *alpha, X, *incX, *beta, Y, *incY);
        }
    };
    //(const MKL_INT N, const void *alpha, const void *X,const MKL_INT incX, const void* beta, void* Y, const MKL_INT incY)
    template <>
    struct axpby_<std::complex<double>> {
        static void call(const MKL_INT* N, const MKL_Complex16* alpha, const MKL_Complex16* X,
            const MKL_INT* incX, const MKL_Complex16* beta, MKL_Complex16* Y, const MKL_INT* incY)
        {
            cblas_zaxpby(*N, alpha, X, *incX, beta, Y, *incY);
        }
    };
#endif

    template <class scalar_t>
    struct copy;
#ifdef FROM_BLAS
    template <>
    struct copy<double> {
        static void call(const MKL_INT N, const double* X, const MKL_INT incX, double* Y, const MKL_INT incY)
        {
            dcopy(&N, X, &incX, Y, &incY);
        }
    };
    template <>
    struct copy<std::complex<double>> {
        static void call(const MKL_INT N, const MKL_Complex16* X, const MKL_INT incX, MKL_Complex16* Y, const MKL_INT incY)
        {
            zcopy(&N, X, &incX, Y, &incY);
        }
    };
#else
    //(const MKL_INT N, const double *X, const MKL_INT incX,double* Y, const MKL_INT incY)
    template <>
    struct copy<double> {
        static constexpr auto call = cblas_dcopy;
    };
    //(const MKL_INT N, const void *X, const MKL_INT incX,void* Y, const MKL_INT incY)
    template <>
    struct copy<std::complex<double>> {
        static constexpr auto call = cblas_zcopy;
    };
#endif

    template <class scalar_t>
    struct dotc;
#ifdef FROM_BLAS
    // double ddot(const MKL_INT *n, const double *x, const MKL_INT *incx, const double *y, const MKL_INT *incy)
    template <>
    struct dotc<double> {
        static void call(const MKL_INT N, const double* X, const MKL_INT incX, const double* Y, const MKL_INT incY, double* dot)
        {
            *dot = ddot(&N, X, &incX, Y, &incY);
        }
    };
    // void zdotc(MKL_Complex16* pres, const MKL_INT* n, const MKL_Complex16* x, const MKL_INT* incx, const MKL_Complex16* y, const MKL_INT* incy)
    template <>
    struct dotc<std::complex<double>> {
        static void call(const MKL_INT N, const MKL_Complex16* X, const MKL_INT incX, const MKL_Complex16* Y, const MKL_INT incY, MKL_Complex16* dotc)
        {
            zdotc(dotc, &N, X, &incX, Y, &incY);
        }
    };
#else
    // double cblas_ddot(const MKL_INT N, const double *X, const MKL_INT incX,const double* Y, const MKL_INT incY)
    template <>
    struct dotc<double> {
        static void call(const MKL_INT N, const double* X, const MKL_INT incX, const double* Y, const MKL_INT incY, double* dot)
        {
            *dot = cblas_ddot(N, X, incX, Y, incY);
        }
    };
    // void cblas_zdotc_sub(const MKL_INT N, const void* X, const MKL_INT incX,const void* Y, const MKL_INT incY, void* dotc)
    template <>
    struct dotc<std::complex<double>> {
        static constexpr auto call = cblas_zdotc_sub;
    };
#endif

    template <class scalar_t>
    struct dotu;
#ifdef FROM_BLAS
    // void zdotu(MKL_Complex16 *pres, const MKL_INT *n, const MKL_Complex16 *x, const MKL_INT *incx, const MKL_Complex16 *y, const MKL_INT *incy)
    template <>
    struct dotu<std::complex<double>> {
        static void call(const MKL_INT N, const MKL_Complex16* X, const MKL_INT incX, const MKL_Complex16* Y, const MKL_INT incY, MKL_Complex16* dotu)
        {
            zdotu(dotu, &N, X, &incX, Y, &incY);
        }
    };
#else
    // void cblas_zdotu_sub(const MKL_INT N, const void *X, const MKL_INT incX,const void* Y, const MKL_INT incY, void* dotu)
    template <>
    struct dotu<std::complex<double>> {
        static constexpr auto call = cblas_zdotu_sub;
    };
#endif

    template <class scalar_t>
    struct nrm2;
#ifdef FROM_BLAS
    // double dnrm2(const MKL_INT *n, const double *x, const MKL_INT *incx)
    template <>
    struct nrm2<double> {
        static double call(const MKL_INT N, const double* X, const MKL_INT incX)
        {
            return dnrm2(&N, X, &incX);
        }
    };
    // double dznrm2(const MKL_INT *n, const MKL_Complex16 *x, const MKL_INT *incx)
    template <>
    struct nrm2<std::complex<double>> {
        static double call(const MKL_INT N, const MKL_Complex16* X, const MKL_INT incX)
        {
            return dznrm2(&N, X, &incX);
        }
    };
#else
    template <>
    struct nrm2<double> {
        static constexpr auto call = cblas_dnrm2;
    };
    // double cblas_dznrm2(const MKL_INT N, const void* X, const MKL_INT incX)
    template <>
    struct nrm2<std::complex<double>> {
        static constexpr auto call = cblas_dznrm2;
    };
#endif

    template <class scalar_t>
    struct scal;
#ifdef FROM_BLAS
    // void dscal(const MKL_INT *n, const double *a, double *x, const MKL_INT *incx)
    template <>
    struct scal<double> {
        static void call(const MKL_INT N, const double* alpha, double* X, const MKL_INT incX)
        {
            dscal(&N, alpha, X, &incX);
        }
    };
    template <>
    struct scal<std::complex<double>> {
        static void call(const MKL_INT N, const MKL_Complex16* alpha, MKL_Complex16* X, const MKL_INT incX)
        {
            zscal(&N, alpha, X, &incX);
        }
    };
#else
    // void cblas_dscal(const MKL_INT N, const double alpha, double *X, const MKL_INT incX)
    template <>
    struct scal<double> {
        static void call(const MKL_INT N, const double* alpha, double* X, const MKL_INT incX)
        {
            cblas_dscal(N, *alpha, X, incX);
        }
    };
    // void cblas_zscal(const MKL_INT N, const void *alpha, void *X, const MKL_INT incX)
    template <>
    struct scal<std::complex<double>> {
        static constexpr auto call = cblas_zscal;
    };
#endif

    template <typename scalar_t>
    struct vmul;
    template <>
    struct vmul<double> {
        static constexpr auto call = vdMul;
    };
    template <>
    struct vmul<std::complex<double>> {
        static constexpr auto call = vzMul;
    };

    template <typename scalar_t>
    struct vdiv;
    template <>
    struct vdiv<double> {
        static constexpr auto call = vdDiv;
    };
    template <>
    struct vdiv<std::complex<double>> {
        static constexpr auto call = vzDiv;
    };

    template <typename scalar_t>
    struct vadd;
    template <>
    struct vadd<double> {
        static constexpr auto call = vdAdd;
    };
    template <>
    struct vadd<std::complex<double>> {
        static constexpr auto call = vzAdd;
    };

    template <typename scalar_t>
    struct vsub;
    template <>
    struct vsub<double> {
        static constexpr auto call = vdSub;
    };
    template <>
    struct vsub<std::complex<double>> {
        static constexpr auto call = vzSub;
    };

    template <typename scalar_t>
    inline void print_e(const scalar_t& i)
    {
        std::cout << i;
    };
    template <>
    inline void print_e(const MKL_Complex16& 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<MKL_Complex16, double>(MKL_Complex16& i, double a)
    {
        i.real = a;
        i.imag = 0;
    };
    template <>
    inline void assignment<MKL_Complex16, std::complex<double>>(MKL_Complex16& i, std::complex<double> a)
    {
        i.real = a.real();
        i.imag = a.imag();
    };
    template <>
    inline void assignment<std::complex<double>, MKL_Complex16>(std::complex<double>& i, MKL_Complex16 a)
    {
        i = std::complex<double>(a.real, a.imag);
    };
}
} // namespace tbplas::math
#endif