//
// Created by Krad on 2023/4/6.
//

#include "Function.h"

//必须在mkl.h和Eigen的头之前，<complex>之后
#define MKL_Complex16 std::complex<double>
#include "mkl.h"
#include <string.h>
#include <algorithm>

#include <Eigen/Core>
#include <Eigen/Eigen>
#include <Eigen/Dense>

namespace Aurora {
    double immse(double *dataA, double *dataB, int size) {
        auto temp = new double[size];
        vdSub(size, dataA, dataB, temp);
        vdSqr(size, temp, temp);
        double result = cblas_dasum(size, temp, 1) / (double) size;
        delete[] temp;
        return result;
    }

    double *std(int rows, int cols, double *input) {
        auto std = new double[cols];
//    #pragma omp parallel for num_threads(4)
        for (int i = 0; i < cols; ++i) {
            double *p = input + i * rows;
            double mean = cblas_dasum(rows, p, 1) / rows;
            vdSubI(rows, p, 1, &mean, 0, p, 1);
            vdSqr(rows, p, p);
            std[i] = cblas_dasum(rows, p, 1) / (rows - 1);
        }
        vdSqrt(cols, std, std);
        return std;
    }

    double *inv(int cols, double *pMatrix) {
        int size = cols * cols;
        int *ipiv = new int[cols];
        auto result = new double[size];
        memcpy(result, pMatrix, sizeof(double) * size);
        LAPACKE_dgetrf(LAPACK_ROW_MAJOR, cols, cols, result, cols, ipiv);
        LAPACKE_dgetri(LAPACK_ROW_MAJOR, cols, result, cols, ipiv);
        delete[] ipiv;
        return result;
    }

    std::complex<double> *fft(long int size, std::complex<double> *input) {
        DFTI_DESCRIPTOR_HANDLE my_desc_handle = NULL;

        auto output = new std::complex<double>[size];
        //使用ce数据fft
        MKL_LONG status;

        //创建 Descriptor, 精度 double , 输入类型实数, 维度1
        status = DftiCreateDescriptor(&my_desc_handle, DFTI_DOUBLE, DFTI_COMPLEX, 1, size);
        //通过 setValue 配置Descriptor
        //使用单独的输出数据缓存
        status = DftiSetValue(my_desc_handle, DFTI_PLACEMENT, DFTI_NOT_INPLACE);

        //提交 修改配置后的Descriptor(实际上会进行FFT的计算初始化)
        status = DftiCommitDescriptor(my_desc_handle);
        //执行计算
        DftiComputeForward(my_desc_handle, input, output);
        //释放资源
        status = DftiFreeDescriptor(&my_desc_handle);
        return output;
    }

    std::complex<double> *ifft(long int size, std::complex<double> *input) {
        DFTI_DESCRIPTOR_HANDLE my_desc_handle = NULL;
        auto output = new std::complex<double>[size];
        MKL_LONG status;

        //创建 Descriptor, 精度 double , 输入类型实数, 维度1
        status = DftiCreateDescriptor(&my_desc_handle, DFTI_DOUBLE, DFTI_COMPLEX, 1, size);
        //通过 setValue 配置Descriptor
        //使用单独的输出数据缓存
        status = DftiSetValue(my_desc_handle, DFTI_PLACEMENT, DFTI_NOT_INPLACE);
        //设置DFTI_BACKWARD_SCALE ！！！很关键，不然值不对
        status = DftiSetValue(my_desc_handle, DFTI_BACKWARD_SCALE, 1.0f / size);
        status = DftiCommitDescriptor(my_desc_handle);
        //提交 修改配置后的Descriptor(实际上会进行FFT的计算初始化)
        status = DftiCommitDescriptor(my_desc_handle);
        //执行计算
        DftiComputeBackward(my_desc_handle, input, output);
        //释放资源
        status = DftiFreeDescriptor(&my_desc_handle);
        return output;
    }

    double *real(int size, std::complex<double> *input) {
        auto input_d = (double *) input;
        auto output = new double[size];
        const int complex_stride = 2;
        const int real_stride = 1;
        cblas_dcopy(size, input_d, complex_stride, output, real_stride);
        return output;
    }

    std::complex<double> *complex(int size, double *input) {
        auto output = (std::complex<double> *) mkl_malloc(size * sizeof(std::complex<double>), 64);
        memset(output, 0, size * sizeof(std::complex<double>));
        const int complex_stride = 2;
        const int real_stride = 1;
        cblas_dcopy(size, input, real_stride, (double *) output, complex_stride);
        return output;
    }

    std::complex<double> *hilbert(int size, double *input) {
        auto complexInput = complex(size, input);
        auto x = fft(size, complexInput);
        auto h = new double[size];
        auto two = 2.0;
        auto zero = 0.0;
        cblas_dcopy(size, &zero, 0, h, 1);
        cblas_dcopy(size / 2, &two, 0, h, 1);
        h[size / 2] = ((size << 31) >> 31) ? 2.0 : 1.0;
        h[0] = 1.0;
        auto p = (double *) x;
        vdMulI(size, p, 2, h, 1, p, 2);
        vdMulI(size, p + 1, 2, h, 1, p + 1, 2);
        auto result = ifft(size, x);
        delete[] h;
        delete[] x;
        return result;
    }

    double *malloc(size_t size, bool complex) {
        if (!complex) return (double *) mkl_malloc(size * sizeof(double), 64);
        size_t complex_size = size * sizeof(std::complex<double>);
        return (double *) mkl_malloc(complex_size, 64);
    }

    void free(void* ptr){
        mkl_free(ptr);
    }

    double * mul(double scalar, double *input, int size) {
        double* output = malloc(size);
        vdMulI(size,input,1,&scalar,0,output,1);
        return output;
    }

    double *mul(double *inputA, double *inputB, int size) {
        double* output = malloc(size);
        vdMulI(size,inputA,1,inputB,1,output,1);
        return output;
    }

    double *mulz(std::complex<double> *inputA, std::complex<double> *inputB, int size) {
        double* output = malloc(size,true);
        vzMulI(size,inputA,1,inputB,1,(std::complex<double>*)output,1);
        return output;
    }

    double *random( int size) {
        double * data = malloc(size);
        Eigen::Map<Eigen::VectorXd> srcV(data,size);
        srcV.setRandom();
        return data;
    }


}