#include "Function1D.h"
#include "Function.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>

using namespace Aurora;

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 = malloc(matrix.getDataSize() ,true);
    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 *) malloc(matrix.getDataSize());
    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 = malloc(matrix.getDataSize());
    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 = malloc(matrix.getDataSize());
    //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) {
    //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 = malloc(matrix.getDataSize());
    //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) {
    //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 = malloc(matrix.getDataSize());
        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(Aurora::Matrix&& matrix) {
    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 = malloc(matrix.getDataSize());
    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(Aurora::Matrix&& matrix) {
    if (matrix.getValueType()==Normal){
        vdAbsI(matrix.getDataSize(), matrix.getData(), SAME_STRIDE, matrix.getData(), SAME_STRIDE);
        return matrix;
    }
    //TODO：考虑尝试是不是使用realloc缩短已分配的内存的方式重用matrix
    else{
        auto output = malloc(matrix.getDataSize());
        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 = malloc(matrix.getDataSize(),true);
        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(Aurora::Matrix&& matrix) {
    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;
    }
}

Matrix Aurora::interp1(const Matrix& aX, const Matrix& aV, const Matrix& aX1, InterpnMethod aMethod)
{
    const int nx = aX.getDimSize(0);
    const int ny = 1;
    int nx1 = aX1.getDimSize(0);
    std::shared_ptr<double> resultData = std::shared_ptr<double>(Aurora::malloc(nx1), Aurora::free);
    std::vector<int> resultInfo = {nx1};
    Aurora::Matrix result(resultData, resultInfo);
    DFTaskPtr task ;
    int status = dfdNewTask1D(&task, nx, aX.getData(), DF_NO_HINT, ny, aV.getData(), DF_NO_HINT);
    if (status != DF_STATUS_OK)
    {
        return Matrix();
    }
    MKL_INT sorder;
    MKL_INT stype;
    if (aMethod == InterpnMethod::Spline)
    {
        sorder = DF_PP_CUBIC;
        stype = DF_PP_BESSEL;
    }
    else
    {
        sorder = DF_PP_LINEAR;
        stype = DF_PP_BESSEL;
    }
    double* scoeffs = Aurora::malloc(ny * (nx-1) * sorder);
     status = dfdEditPPSpline1D(task, sorder,DF_PP_NATURAL , DF_BC_NOT_A_KNOT, 0, DF_NO_IC, 0, scoeffs, DF_NO_HINT);
    if (status != DF_STATUS_OK)
    {
        return Matrix();
    }
    status = dfdConstruct1D( task, DF_PP_SPLINE, DF_METHOD_STD );
    if (status != DF_STATUS_OK)
    {
        return Matrix();
    }

    int dorder = 1;
    status = dfdInterpolate1D(task, DF_INTERP, DF_METHOD_PP, nx1, aX1.getData(), DF_NO_HINT, 1, &dorder,
                              DF_NO_APRIORI_INFO, resultData.get(), DF_MATRIX_STORAGE_ROWS, nullptr);

    status = dfDeleteTask(&task);

    Aurora::free(scoeffs);
    return result;
}

Matrix Aurora::repmat(const Matrix& aMatrix,int aRowTimes, int aColumnTimes)
{
    if(aRowTimes < 1 || aColumnTimes < 1 || aMatrix.getDims() > 2 || aMatrix.isNull())
    {
        return Matrix();
    }
    int complexStep = aMatrix.getValueType();
    int originalDataSize = aMatrix.getDataSize() * complexStep;
    double* resultData = Aurora::malloc(originalDataSize * aRowTimes * aColumnTimes);
    int row = aMatrix.getDimSize(0);
    int column = aMatrix.getDimSize(1);
    double* originalData = aMatrix.getData();
    double* resultDataTemp = resultData;
    for(int i=0; i<column; ++i)
    {
        for(int j=1; j<=aRowTimes; ++j)
        {
            cblas_dcopy(row*complexStep,originalData, 1, resultDataTemp, 1);
            resultDataTemp += row*complexStep;
        }
        originalData += row*complexStep;
    }

    resultDataTemp = resultData;
    int step = originalDataSize * aRowTimes;
    for(int i=1; i<aColumnTimes; ++i)
    {
        cblas_dcopy(step, resultData, 1, resultData + i*step, 1);
    }

    std::vector<int> resultInfo;
    resultInfo.push_back(aRowTimes * row);
    column = aColumnTimes*column;
    if (column > 1)
    {
        resultInfo.push_back(column);
    }

    return Matrix(std::shared_ptr<double>(resultData, Aurora::free),resultInfo, aMatrix.getValueType());
}

Matrix Aurora::repmat(const Matrix& aMatrix,int aRowTimes, int aColumnTimes, int aSliceTimes)
{
    if(aRowTimes < 1 || aColumnTimes < 1 || aSliceTimes < 1 || aMatrix.getDims() > 2 || aMatrix.isNull())
    {
        return Matrix();
    }

    int complexStep = aMatrix.getValueType();
    Matrix resultTemp = Aurora::repmat(aMatrix, aRowTimes, aColumnTimes);
    int resultTempDataSize = resultTemp.getDataSize() * complexStep;
    double* resultData = Aurora::malloc(resultTempDataSize * aSliceTimes);
    std::copy(resultTemp.getData(), resultTemp.getData() + resultTempDataSize, resultData);
    for(int i=1; i<aSliceTimes; ++i)
    {
        cblas_dcopy(resultTempDataSize, resultData, 1, resultData + i*resultTempDataSize, 1);
    }
    std::vector<int> resultInfo;
    int row = resultTemp.getDimSize(0);
    int column = resultTemp.getDimSize(1);
    resultInfo.push_back(row);
    if (column > 1 || aSliceTimes > 1)
    {
        resultInfo.push_back(column);
    }
    if (aSliceTimes > 1)
    {
        resultInfo.push_back(aSliceTimes);
    }

    return Matrix(std::shared_ptr<double>(resultData, Aurora::free), resultInfo, aMatrix.getValueType());
}


Matrix Aurora::polyval(const Matrix &aP, const Matrix &aX) {
    auto result = malloc(aX.getDataSize());
    auto powArg = new double[aP.getDataSize()];
    for (int j = aP.getDataSize(), i = 0; j > 0; --j, ++i) {
        powArg[i] = (double) (j - 1);
    }
    auto temp = new double[aP.getDataSize()];
    for (int i = 0; i < aX.getDataSize(); ++i) {
        vdPowI(aP.getDataSize(), aX.getData() + i, 0, powArg, 1, temp, 1);
        vdMul(aP.getDataSize(), aP.getData(), temp, temp);
        result[i] = cblas_dasum(3, temp, 1);
    }
    delete[] powArg;
    delete[] temp;

    return Matrix::New(result,aX);
}

Matrix Aurora::log(const Matrix& aMatrix, int aBaseNum)
{
    size_t size = aMatrix.getDataSize();
    double* data = Aurora::malloc(size);
    vdLn(size, aMatrix.getData(), data);
    if(aBaseNum != -1)
    {
        double baseNum = aBaseNum;
        double temp;
        for (size_t i = 0; i < size; i++)
        {
            vdLn(1, &baseNum, &temp);
            data[i] /= temp;
        }
    }

    return Aurora::Matrix(std::shared_ptr<double>(data,Aurora::free),std::vector<int>{aMatrix.getDimSize(0),aMatrix.getDimSize(1),aMatrix.getDimSize(2)});
}