#include <iostream>
#include "Function3D.h"
#include "Function2D.h"
#include "Function.h"

//必须在Eigen之前
#include "AuroraDefs.h"

using namespace Aurora;

Matrix Aurora::interp3(const Matrix& aX, const Matrix& aY, const Matrix& aZ, const Matrix& aV, const Matrix& aX1, const Matrix& aY1, const Matrix& aZ1,InterpnMethod aMethod)
{
    if (aV.getDims() != 3)
    {
        return Matrix();
    }

    if ( aX1.getDimSize(0) != aY1.getDimSize(0) ||
         aX1.getDimSize(0) != aZ1.getDimSize(0) ||
         aZ1.getDimSize(0) != aY1.getDimSize(0))
    {
        return Matrix();
    }
    int nx1 = aX1.getDimSize(0);
    int zAxisNum = aV.getDimSize(2);
    int columnNum = aV.getDimSize(1);
    int rowNum = aV.getDimSize(0);

    if(aX.getDimSize(0) != columnNum || aY.getDimSize(0) != rowNum || aZ.getDimSize(0) != zAxisNum)
    {
      return Matrix();
    }
    std::shared_ptr<float> resultData = std::shared_ptr<float>(new float[nx1], std::default_delete<float[]>());
    for (int i = 0; i < nx1; ++i)
    {
        std::shared_ptr<float> zResultData = std::shared_ptr<float>(new float[zAxisNum], std::default_delete<float[]>());
        for(int j =0; j < zAxisNum; ++j)
        {
            zResultData.get()[j] = interp2(aX,aY,aV($,$,j).toMatrix(),aX1(i).toMatrix(),aY1(i).toMatrix(),aMethod).getData()[0];
        }
        Matrix xResult(zResultData,std::vector<int>{columnNum});
        resultData.get()[i] = interp1(aZ,xResult,aZ1(i).toMatrix(),aMethod).getData()[0];
    }
    return Matrix(resultData,std::vector<int>{nx1});
}

Matrix Aurora::interpn(const Matrix& aX, const Matrix& aY, const Matrix& aZ, const Matrix& aV, const Matrix& aX1, const Matrix& aY1, const Matrix& aZ1,InterpnMethod aMethod)
{
    return Aurora::interp3(aY,aX,aZ,aV,aY1,aX1,aZ1,aMethod);
}

Matrix Aurora::ones(int aRow, int aColumn, int aSlice) {
    if (aRow == 0 || aColumn == 0)
    {
        std::cerr<<"ones function can create matrix with dim unit cont =0";
        return Matrix();
    }
    int rowSize =  aRow;
    int colSize =  aColumn;
    int sliceSize = aSlice == 0 ? 1 : aSlice;
    size_t arraySize  = rowSize * colSize* sliceSize;
    float* data = malloc(arraySize);
    float one = 1.0;
    cblas_scopy(arraySize,&one,0,data,1);
    return Matrix::New(data,rowSize,colSize,aSlice);
}

Matrix Aurora::ones(int aSquareRow) {
    return Aurora::ones(aSquareRow, aSquareRow);
}

Matrix Aurora::zeros(int aRow, int aColumn, int aSlice) {
    if (aRow == 0 || aColumn == 0)
    {
        std::cerr<<"zeros function can create matrix with dim unit cont =0";
        return Matrix();
    }
    int rowSize =  aRow;
    int colSize =  aColumn;
    int sliceSize = aSlice == 0 ? 1 : aSlice;
    size_t arraySize  = rowSize * colSize* sliceSize;
    float* data = malloc(arraySize);
    float zero = 0.0;
    cblas_scopy(arraySize,&zero,0,data,1);
    return Matrix::New(data,rowSize,colSize,sliceSize);
}

Matrix Aurora::zeros(int aSquareRow) {
    return Aurora::zeros(aSquareRow, aSquareRow);
}

Matrix Aurora::size(const Matrix &aMatrix)
{
    if (aMatrix.isScalar()){
        float * output = Aurora::malloc(1);
        output[0]=1;
        return Matrix::New(output,1,1,1);
    }
    else if (aMatrix.isVector()){
        float * output = Aurora::malloc(2);
        output[0]=aMatrix.getDimSize(0);
        output[1]=aMatrix.getDimSize(1);
        return Matrix::New(output,2,1,1);
    }
    //3D
    else if (aMatrix.getDimSize(2)>1){
        float * output = Aurora::malloc(3);
        output[0]=aMatrix.getDimSize(0);
        output[1]=aMatrix.getDimSize(1);
        output[2]=aMatrix.getDimSize(2);
        return Matrix::New(output,3,1,1);
    }
    //2D matrix
    else{
        float * output = Aurora::malloc(2);
        output[0]=aMatrix.getDimSize(0);
        output[1]=aMatrix.getDimSize(1);
        return Matrix::New(output,2,1,1);
    }
}

int Aurora::size(const Matrix &aMatrix,int dims)
{
    return aMatrix.getDimSize(dims-1);
}

Matrix Aurora::meshgridInterp3(const Matrix& aX, const Matrix& aY, const Matrix& aZ, const Matrix& aV, const Matrix& aX1, const Matrix& aY1, const Matrix& aZ1,InterpnMethod aMethod, float aExtrapval)
{
    std::vector<Matrix> zTemps;
    for(int rIndex=0; rIndex<aV.getDimSize(2); ++rIndex)
    {
        std::vector<Matrix> yTemps;
        for(int xIndex=0; xIndex<aV.getDimSize(0); ++xIndex)
        {
            Matrix test = aV($, xIndex, rIndex).toMatrix();
            Matrix a = interp1(aY, test, aY1, aMethod);
            yTemps.push_back(a);
        }

        std::vector<Matrix> xTemps;
        for(int yIndex=0; yIndex<yTemps[0].getDataSize(); ++yIndex)
        {
            float* xTempData = new float[yTemps.size()];
            for(int i=0;i<yTemps.size();++i)
            {
                xTempData[i] =  yTemps[i][yIndex];
            }
            Matrix xTemp = Matrix::fromRawData(xTempData, yTemps.size());
            Matrix b = interp1(aX, xTemp, aX1, aMethod);
            xTemps.push_back(b);
        }

        Matrix result2D = xTemps[0];
        for(int i=1; i<xTemps.size(); ++i)
        {
            result2D = horzcat(result2D, xTemps[i]);
        }
        result2D = transpose(result2D);
        zTemps.push_back(result2D);
    }

    std::vector<Matrix>  resultVector;
    for(int i=0; i<zTemps[0].getDataSize(); ++i)
    {
        float* resultTemp = new float[zTemps.size()]; 
        for(int j=0; j<zTemps.size(); ++j)
        {
            resultTemp[j] = zTemps[j][i];
        }
        Matrix zTemp = Matrix::fromRawData(resultTemp, zTemps.size());
        Matrix c = interp1(aZ, zTemp, aZ1, aMethod);
        resultVector.push_back(c);
    }

    Matrix result = resultVector[0];
    for(int i=1; i<resultVector.size(); ++i)
    {
        result = horzcat(result, resultVector[i]);
    }
    result = transpose(result);
    result.forceReshape(aX1.getDataSize(), aY1.getDataSize(), aZ1.getDataSize());

    float minXSOS = aX[0];
    float maxXSOS = aX[aX.getDataSize() - 1];
    float minYSOS = aY[0];
    float maxYSOS = aY[aY.getDataSize() - 1];
    float minZSOS = aZ[0];
    float maxZSOS = aZ[aZ.getDataSize() - 1];
    for(size_t zIndex=0; zIndex<aZ1.getDataSize(); ++zIndex)
    {
        for(size_t yIndex=0; yIndex<aY1.getDataSize(); ++yIndex)
        {
            for(size_t xIndex=0; xIndex<aX1.getDataSize(); ++xIndex)
            {
                size_t index = zIndex * aY1.getDataSize() * aX1.getDataSize() +
                               yIndex * aX1.getDataSize() + xIndex;
                if(aX1[xIndex] < minXSOS || aX1[xIndex] > maxXSOS ||
                   aY1[yIndex] < minYSOS || aY1[yIndex] > maxYSOS ||
                   aZ1[zIndex] < minZSOS || aZ1[zIndex] > maxZSOS)
                   {
                        result[index] = aExtrapval;
                   }
            }
        }
    }

    return result;
}
