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

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

namespace {
    const int COMPLEX_STRIDE = 2;
    const int REAL_STRIDE = 1;
    const int SAME_STRIDE = 1;
    const double VALUE_ONE = 1.0;
}


Aurora::Matrix Aurora::complex(const Aurora::Matrix &matrix) {
    if (matrix.getValueType() == Complex) {
        std::cerr<<"complex not support complex value type"<<std::endl;
        return matrix;
    }
    auto output = (std::complex<double> *) mkl_malloc(matrix.getDataSize() * sizeof(std::complex<double>), 64);
    memset(output, 0, (matrix.getDataSize() * sizeof(std::complex<double>)));
    cblas_dcopy(matrix.getDataSize(), matrix.getData(), REAL_STRIDE, (double *) output, COMPLEX_STRIDE);
    return Aurora::Matrix::New((double *) output, matrix.getDimSize(0), matrix.getDimSize(1), matrix.getDimSize(2),
                               Complex);
}

Aurora::Matrix Aurora::real(const Aurora::Matrix &matrix) {
    if (matrix.getValueType() == Normal) {
        std::cerr<<"real only support complex value type"<<std::endl;
        return matrix;
    }
    auto output = (double *) mkl_malloc(matrix.getDataSize() * sizeof(double), 64);
    memset(output, 0, (matrix.getDataSize() * sizeof(double)));
    cblas_dcopy(matrix.getDataSize(), matrix.getData(),COMPLEX_STRIDE , (double *) output, REAL_STRIDE);
    return Aurora::Matrix::New((double *) output, matrix.getDimSize(0), matrix.getDimSize(1), matrix.getDimSize(2));
}

Aurora::Matrix Aurora::imag(const Aurora::Matrix &matrix) {
    if (matrix.getValueType() == Normal) {
        std::cerr<<"imag only support complex value type"<<std::endl;
        return matrix;
    }
    auto output = (double *) mkl_malloc(matrix.getDataSize() * sizeof(double), 64);
    memset(output, 0, (matrix.getDataSize() * sizeof(double)));
    cblas_dcopy(matrix.getDataSize(), matrix.getData()+1,COMPLEX_STRIDE , (double *) output, REAL_STRIDE);
    return Aurora::Matrix::New((double *) output, matrix.getDimSize(0), matrix.getDimSize(1), matrix.getDimSize(2));
}

Aurora::Matrix Aurora::ceil(const Aurora::Matrix &matrix) {
    auto output = (double *) mkl_malloc(matrix.getDataSize() * sizeof(double), 64);
    //for real part
    vdCeilI(matrix.getDataSize(), matrix.getData(), SAME_STRIDE, output, SAME_STRIDE);
    if (matrix.getValueType() == Complex) {
        //for imag part
        vdCeilI(matrix.getDataSize(), matrix.getData() + 1, SAME_STRIDE, output + 1, SAME_STRIDE);
    }
    return Aurora::Matrix::New(output, matrix);
}

Aurora::Matrix Aurora::ceil(const Aurora::Matrix &&matrix) {
    std::cout<<"RR ceil"<<std::endl;
    //for real part
    vdCeilI(matrix.getDataSize(), matrix.getData(), SAME_STRIDE, matrix.getData(), SAME_STRIDE);
    if (matrix.getValueType() == Complex) {
        //for imag part
        vdCeilI(matrix.getDataSize(), matrix.getData() + 1, SAME_STRIDE, matrix.getData() + 1, SAME_STRIDE);
    }
    return matrix;
}

Aurora::Matrix Aurora::round(const Aurora::Matrix &matrix) {
    auto output = (double *) mkl_malloc(matrix.getDataSize() * sizeof(double), 64);
    //for real part
    vdRoundI(matrix.getDataSize(), matrix.getData(), SAME_STRIDE, output, SAME_STRIDE);
    if (matrix.getValueType() == Complex) {
        //for imag part
        vdRoundI(matrix.getDataSize(), matrix.getData() + 1, SAME_STRIDE, output + 1, SAME_STRIDE);
    }
    return Aurora::Matrix::New(output, matrix);
}

Aurora::Matrix Aurora::round(const Aurora::Matrix &&matrix) {
    std::cout<<"RR round"<<std::endl;
    //for real part
    vdRoundI(matrix.getDataSize(), matrix.getData(), SAME_STRIDE, matrix.getData(), SAME_STRIDE);
    if (matrix.getValueType() == Complex) {
        //for imag part
        vdRoundI(matrix.getDataSize(), matrix.getData() + 1, SAME_STRIDE, matrix.getData() + 1, SAME_STRIDE);
    }
    return matrix;
}

Aurora::Matrix Aurora::sqrt(const Aurora::Matrix& matrix) {

    if (matrix.getValueType() != Complex) {
        auto output = (double *) mkl_malloc(matrix.getDataSize() * sizeof(double), 64);
        vdSqrtI(matrix.getDataSize(), matrix.getData(), SAME_STRIDE, output, SAME_STRIDE);
        return Aurora::Matrix::New(output, matrix);
    }
    std::cerr<<"sqrt not support complex"<<std::endl;
    return Aurora::Matrix();
}

Aurora::Matrix Aurora::sqrt(const Aurora::Matrix&& matrix) {
    std::cout<<"RR sqrt"<<std::endl;
    if (matrix.getValueType() != Complex) {
        vdSqrtI(matrix.getDataSize(), matrix.getData(), SAME_STRIDE, matrix.getData(), SAME_STRIDE);
        return matrix;
    }
    std::cerr<<"sqrt not support complex"<<std::endl;
    return Aurora::Matrix();
}

Aurora::Matrix Aurora::abs(const Aurora::Matrix &matrix) {
    auto output = (double *) mkl_malloc(matrix.getDataSize() * sizeof(double), 64);
    if (matrix.getValueType()==Normal){
        vdAbsI(matrix.getDataSize(), matrix.getData(), SAME_STRIDE, output, SAME_STRIDE);
    }
    else{
        vzAbsI(matrix.getDataSize(), (std::complex<double> *)matrix.getData(), SAME_STRIDE,output, SAME_STRIDE);
    }
    return Aurora::Matrix::New(output, matrix);
}

Aurora::Matrix Aurora::abs(const Aurora::Matrix&& matrix) {
    std::cout<<"RR abs"<<std::endl;
    if (matrix.getValueType()==Normal){
        vdAbsI(matrix.getDataSize(), matrix.getData(), SAME_STRIDE, matrix.getData(), SAME_STRIDE);
        return matrix;
    }
    //TODO：考虑尝试是不是使用realloc缩短已分配的内存的方式重用matrix
    else{
        auto output = (double *) mkl_malloc(matrix.getDataSize() * sizeof(double), 64);
        vzAbsI(matrix.getDataSize(), (std::complex<double> *)matrix.getData(), SAME_STRIDE,output, SAME_STRIDE);
        return Aurora::Matrix::New(output, matrix);
    }

}

Aurora::Matrix Aurora::sign(const Aurora::Matrix &matrix) {
    if (matrix.getValueType()==Normal){
        auto ret = matrix.deepCopy();
        Eigen::Map<Eigen::VectorXd> retV(ret.getData(),ret.getDataSize());
        retV = retV.array().sign();
        return ret;
    }
    else{
        //sign(x) = x./abs(x)，前提是 x 为复数。
        auto output = (double *) mkl_malloc(matrix.getDataSize() * sizeof(std::complex<double>), 64);
        Matrix absMatrix = abs(matrix);
        vdDivI(matrix.getDataSize(), matrix.getData(),COMPLEX_STRIDE,
                absMatrix.getData(), REAL_STRIDE,output,COMPLEX_STRIDE);
        vdDivI(matrix.getDataSize(), matrix.getData()+1,COMPLEX_STRIDE,
               absMatrix.getData(), REAL_STRIDE,output+1,COMPLEX_STRIDE);
        return Aurora::Matrix::New(output, matrix);
    }
}

Aurora::Matrix Aurora::sign(const Aurora::Matrix&& matrix) {
    std::cout<<"RR sign"<<std::endl;
    if (matrix.getValueType()==Normal){
        Eigen::Map<Eigen::VectorXd> retV(matrix.getData(),matrix.getDataSize());
        retV = retV.array().sign();
        return matrix;
    }
    else{
        //sign(x) = x./abs(x)，前提是 x 为复数。
        Matrix absMatrix = abs(matrix);
        vdDivI(matrix.getDataSize(), matrix.getData(),COMPLEX_STRIDE,
               absMatrix.getData(), REAL_STRIDE,matrix.getData(),COMPLEX_STRIDE);
        vdDivI(matrix.getDataSize(), matrix.getData()+1,COMPLEX_STRIDE,
               absMatrix.getData(), REAL_STRIDE,matrix.getData()+1,COMPLEX_STRIDE);
        return matrix;
    }
}