#include <gtest/gtest.h>
#include <vector>
#include "Matrix.h"
#include "Function.h"
#include "Function1D.h"
#include "Function2D.h"
#include "Function3D.h"
#include "TestUtility.h"


class Matrix_Test : public ::testing::Test {
protected:
    static void SetUpMatrixTester() {

    }

    static void TearDownTestCase() {
    }

    void SetUp() {
    }

    void TearDown() {
    }
};

TEST_F(Matrix_Test, MatrixCreate) {
    float *dataA = Aurora::malloc(9);
    float *dataB = new float[9];
    float *dataC = new float[9];
    float *dataD = new float[9];
//mkl matrix
    {
        Aurora::Matrix A = Aurora::Matrix::New(dataA, 3, 3);
        EXPECT_EQ(dataA, A.getData());
        EXPECT_EQ(9, A.getDataSize());
        EXPECT_EQ(2, A.getDims());
        EXPECT_EQ(3, A.getDimSize(0));
        EXPECT_EQ(3, A.getDimSize(1));
        EXPECT_EQ(Aurora::Normal, A.getValueType());
    }
// float [] matrix
    {
        Aurora::Matrix B = Aurora::Matrix::fromRawData(dataB, 3, 3);
        EXPECT_EQ(dataB, B.getData());
        EXPECT_EQ(9, B.getDataSize());
        EXPECT_EQ(2, B.getDims());
        EXPECT_EQ(3, B.getDimSize(0));
        EXPECT_EQ(3, B.getDimSize(1));
        EXPECT_EQ(Aurora::Normal, B.getValueType());
    }
// copy from float [] to mkl matrix
    {
        Aurora::Matrix C = Aurora::Matrix::copyFromRawData(dataC, 3, 3);
        EXPECT_NE(dataC, C.getData());
        EXPECT_EQ(9, C.getDataSize());
        EXPECT_EQ(2, C.getDims());
        EXPECT_EQ(3, C.getDimSize(0));
        EXPECT_EQ(3, C.getDimSize(1));
        EXPECT_EQ(Aurora::Normal, C.getValueType());
    }
// 2vector
    {
        Aurora::Matrix C = Aurora::Matrix::copyFromRawData(dataC, 9);
        EXPECT_NE(dataC, C.getData());
        EXPECT_EQ(9, C.getDataSize());
        EXPECT_EQ(2, C.getDims());
        EXPECT_EQ(9, C.getDimSize(0));
        EXPECT_EQ(1, C.getDimSize(1));
        EXPECT_EQ(Aurora::Normal, C.getValueType());
    }
// 2d vector
    {
        Aurora::Matrix C = Aurora::Matrix::fromRawData(dataC, 9, 1);
        EXPECT_EQ(dataC, C.getData());
        EXPECT_EQ(9, C.getDataSize());
        EXPECT_EQ(2, C.getDims());
        EXPECT_EQ(9, C.getDimSize(0));
        EXPECT_EQ(1, C.getDimSize(1));
        EXPECT_EQ(Aurora::Normal, C.getValueType());
    }
// 2d vector column major
    {
        Aurora::Matrix C = Aurora::Matrix::copyFromRawData(dataD, 1, 9);
        EXPECT_NE(dataD, C.getData());
        delete[] dataD;
        EXPECT_EQ(9, C.getDataSize());
        EXPECT_EQ(2, C.getDims());
        EXPECT_EQ(1, C.getDimSize(0));
        EXPECT_EQ(9, C.getDimSize(1));
        EXPECT_EQ(Aurora::Normal, C.getValueType());
    }
// 3d matrix
    {
        float *tempData = new float[9];
        Aurora::Matrix C = Aurora::Matrix::fromRawData(tempData, 3, 3, 1);
        EXPECT_EQ(tempData, C.getData());
        EXPECT_EQ(9, C.getDataSize());
        EXPECT_EQ(2, C.getDims());
        EXPECT_EQ(3, C.getDimSize(0));
        EXPECT_EQ(3, C.getDimSize(1));
        EXPECT_EQ(1, C.getDimSize(2));
        EXPECT_EQ(Aurora::Normal, C.getValueType());
    }
// 3d matrix 2
    {
        float *tempData = new float[9];
        Aurora::Matrix C = Aurora::Matrix::fromRawData(tempData, 3, 1, 3);
        EXPECT_EQ(tempData, C.getData());
        EXPECT_EQ(9, C.getDataSize());
        EXPECT_EQ(3, C.getDims());
        EXPECT_EQ(3, C.getDimSize(0));
        EXPECT_EQ(1, C.getDimSize(1));
        EXPECT_EQ(3, C.getDimSize(2));
        EXPECT_EQ(Aurora::Normal, C.getValueType());
    }
}

TEST_F(Matrix_Test, matrixSlice) {
    float *dataA = Aurora::malloc(8);
    float *dataB = Aurora::malloc(8);
    float *dataC = Aurora::malloc(8);
    for (int i = 0; i < 8; ++i) {
        dataA[i] = (float) (i);
        dataB[i] = (float) (1);
        dataC[i] = (float) (9 - i);
    }
    Aurora::Matrix A = Aurora::Matrix::New(dataA, 2, 2, 2);
    Aurora::Matrix B = Aurora::Matrix::New(dataB, 2, 2, 2);
    Aurora::Matrix C = Aurora::Matrix::New(dataC, 2, 2, 2);
//2D slice
    EXPECT_EQ(4.0, dataA[4]);
    A(Aurora::$, Aurora::$, 1) = B(0, Aurora::$, Aurora::$);
    EXPECT_EQ(1, dataA[4]);
    EXPECT_EQ(3, dataA[3]);
    A(Aurora::$, 1, Aurora::$) = B(Aurora::$, Aurora::$, 0);
    EXPECT_EQ(1.0, dataA[3]);
    EXPECT_EQ(0.0, dataA[0]);
    A(0, Aurora::$, Aurora::$) = B(Aurora::$, 0, Aurora::$);
    EXPECT_EQ(1.0, dataA[0]);
//vector slice
    A(0, Aurora::$, 0) = C(0, 0, Aurora::$);
    EXPECT_EQ(9.0, dataA[0]);
    A(Aurora::$, 0, 0) = C(0, Aurora::$, 1);
    EXPECT_EQ(5.0, dataA[0]);
//error slice
    EXPECT_EQ(1, A(Aurora::$, Aurora::$, Aurora::$).toMatrix().getDataSize());
    auto D = C(0, 0, 0).toMatrix();
    EXPECT_EQ(1, D.getDataSize());
    EXPECT_EQ(9, D.getData()[0]);


    float *dataD = Aurora::malloc(27);
    for (int i = 0; i < 27; ++i) {
        dataD[i] = (float) (i);
    }
    Aurora::Matrix D1 = Aurora::Matrix::New(dataD, 3, 3, 3);
    auto r1 = D1.block(0, 0, 1);
    EXPECT_EQ(2,r1.getDimSize(0));
    EXPECT_EQ(3,r1.getData()[2]);
    auto r2 = D1.block(1, 0, 0);
    EXPECT_EQ(1,r2.getDimSize(1));
    EXPECT_EQ(10,r2.getData()[4]);
    auto r3 = D1.block(2, 1, 2);
    EXPECT_EQ(2,r3.getDimSize(2));
    EXPECT_EQ(9,r3.getData()[0]);
}

TEST_F(Matrix_Test, matrixOpertaor) {
//3D
    {
        float *dataA = new float[8];
        float *dataB = new float[8];
        float *dataD = new float[8];
        for (int i = 0; i < 8; ++i) {
            dataA[i] = (float) (i);
            dataB[i] = (float) (2);
        }
        Aurora::Matrix A = Aurora::Matrix::fromRawData(dataA, 2, 2, 2);
        Aurora::Matrix B = Aurora::Matrix::fromRawData(dataB, 2, 2, 2);
        Aurora::Matrix D = Aurora::Matrix::fromRawData(dataD, 4, 2);
        auto C = (A * B) - B;
        EXPECT_EQ(C.getData()[2], 2);
        C = A * B / 2.0;
        EXPECT_EQ(C.getData()[2], 2);
        C = (A * 8.0 / B) * (B + 1.0) - 2.0 + 8;
        EXPECT_EQ(C.getData()[2], 30);
//Error matrix
        C = A * D;
        EXPECT_EQ(C.getDataSize(), 0);
    }
//complex
    {
        float *dataA = Aurora::malloc(8, true);
        float *dataB = Aurora::malloc(8, true);
        float *dataD = new float[8];
        for (int i = 0; i < 16; ++i) {
            dataA[i] = (float) (i + 1);
            dataB[i] = (float) (2);
            if (i < 8) dataD[i] = (float) (2);
        }
        Aurora::Matrix A = Aurora::Matrix::New(dataA, 2, 2, 2, Aurora::Complex);
        Aurora::Matrix B = Aurora::Matrix::New(dataB, 2, 2, 2, Aurora::Complex);
        Aurora::Matrix D = Aurora::Matrix::fromRawData(dataD, 2, 2, 2);
        auto C = A + 5.0 - 2.0;
        EXPECT_EQ(C.getData()[2], 6);
        EXPECT_EQ(C.getData()[3], 4);
        EXPECT_TRUE(C.isComplex());
        C = A * B - B;
        EXPECT_EQ(C.getData()[2], -4);
        EXPECT_EQ(C.getData()[3], 12);
        EXPECT_TRUE(C.isComplex());
        C = A * B - 2.0 + 5.0;
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 14);
        C = A * B / 2.0 * 3.;
        EXPECT_EQ(C.getData()[2], -3);
        EXPECT_EQ(C.getData()[3], 21);
        C = A * B / B;
        EXPECT_EQ(C.getData()[2], 3);
        EXPECT_EQ(C.getData()[3], 4);
        C = A * B / (B * A / A);
        EXPECT_EQ(C.getData()[2], 3);
        EXPECT_EQ(C.getData()[3], 4);

        C = B * D;
        EXPECT_EQ(C.getData()[2], 4);
        EXPECT_EQ(C.getData()[3], 4);
        C = B / D;
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 1);
        C = B + D;
        EXPECT_EQ(C.getData()[2], 4);
        EXPECT_EQ(C.getData()[3], 2);
        C = B - D;
        EXPECT_EQ(C.getData()[2], 0);
        EXPECT_EQ(C.getData()[3], 2);

        C = B * (D * 1.0);
        EXPECT_EQ(C.getData()[2], 4);
        EXPECT_EQ(C.getData()[3], 4);
        C = B / (D * 1.0);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 1);
        C = B + (D * 1.0);
        EXPECT_EQ(C.getData()[2], 4);
        EXPECT_EQ(C.getData()[3], 2);
        C = B - (D * 1.0);
        EXPECT_EQ(C.getData()[2], 0);
        EXPECT_EQ(C.getData()[3], 2);

        C = (B * 1.0) * (D * 1.0);
        EXPECT_EQ(C.getData()[2], 4);
        EXPECT_EQ(C.getData()[3], 4);
        C = (B * 1.0) / (D * 1.0);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 1);
        C = (B * 1.0) + (D * 1.0);
        EXPECT_EQ(C.getData()[2], 4);
        EXPECT_EQ(C.getData()[3], 2);
        C = (B * 1.0) - (D * 1.0);
        EXPECT_EQ(C.getData()[2], 0);
        EXPECT_EQ(C.getData()[3], 2);

        C= -B;
        EXPECT_EQ(C.getData()[0], -2);
        EXPECT_EQ(C.getData()[1], -2);
    }
}

TEST_F(Matrix_Test, matrixCompare){
        float *dataA = new float[8];
        float *dataB = new float[8];
        for (int i = 0; i < 8; ++i) {
            dataA[i] = (float) (i);
            dataB[i] = (float) (2);
        }
        Aurora::Matrix A = Aurora::Matrix::fromRawData(dataA, 2, 2, 2);
        Aurora::Matrix B = Aurora::Matrix::fromRawData(dataB, 2, 2, 2);
        auto C = A < B;
        EXPECT_EQ(C.getData()[0], 1);
        EXPECT_EQ(C.getData()[1], 1);
        EXPECT_EQ(C.getData()[2], 0);
        EXPECT_EQ(C.getData()[3], 0);
        C = A <= B;
        EXPECT_EQ(C.getData()[0], 1);
        EXPECT_EQ(C.getData()[1], 1);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 0);
        C = A > B;
        EXPECT_EQ(C.getData()[0], 0);
        EXPECT_EQ(C.getData()[1], 0);
        EXPECT_EQ(C.getData()[2], 0);
        EXPECT_EQ(C.getData()[3], 1);
        C = A >= B;
        EXPECT_EQ(C.getData()[0], 0);
        EXPECT_EQ(C.getData()[1], 0);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 1);
        C = A == B;
        EXPECT_EQ(C.getData()[0], 0);
        EXPECT_EQ(C.getData()[1], 0);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 0);
        C = A == 1.0;
        EXPECT_EQ(C.getData()[0], 0);
        EXPECT_EQ(C.getData()[1], 1);
        EXPECT_EQ(C.getData()[2], 0);
        EXPECT_EQ(C.getData()[3], 0);

        C = A > 1.0;
        EXPECT_EQ(C.getData()[0], 0);
        EXPECT_EQ(C.getData()[1], 0);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 1);

        C = A >= 1.0;
        EXPECT_EQ(C.getData()[0], 0);
        EXPECT_EQ(C.getData()[1], 1);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 1);

        C = A < 3.0;
        EXPECT_EQ(C.getData()[0], 1);
        EXPECT_EQ(C.getData()[1], 1);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 0);

        C = A <= 2.0;
        EXPECT_EQ(C.getData()[0], 1);
        EXPECT_EQ(C.getData()[1], 1);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 0);

        C = A == 2.0;
        EXPECT_EQ(C.getData()[0], 0);
        EXPECT_EQ(C.getData()[1], 0);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 0);

}

TEST_F(Matrix_Test, matrixfunction){
    float *dataA = new float[9]{1,6,9,4,1,0,5,8,1};
    float *dataB = new float[3]{1,2,3};
    Aurora::Matrix A = Aurora::Matrix::fromRawData(dataA, 3, 3);
    Aurora::Matrix B = Aurora::Matrix::fromRawData(dataB, 1, 3);
    auto C = B.mul(A);
    EXPECT_FLOAT_EQ(C.getData()[0], 40);
    EXPECT_FLOAT_EQ(C.getData()[1], 6);
    EXPECT_FLOAT_EQ(C.getData()[2], 24);
    B.forceReshape(3, 1, 1);
    C = A.mul(B);
    EXPECT_FLOAT_EQ(C.getData()[0], 24);
    EXPECT_FLOAT_EQ(C.getData()[1], 32);
    EXPECT_FLOAT_EQ(C.getData()[2], 12);
    float *dataD = new float[9]{2.1,3,9,-3,1,0,51,-8,1};
    Aurora::Matrix D = Aurora::Matrix::fromRawData(dataD, 3, 3);
    C = A.mul(D);
    EXPECT_FLOAT_EQ(C.getData()[0], 59.1);
    EXPECT_FLOAT_EQ(C.getData()[1], 87.6);
    EXPECT_FLOAT_EQ(C.getData()[2], 27.9);
    EXPECT_FLOAT_EQ(C.getData()[3], 1);
    EXPECT_FLOAT_EQ(C.getData()[4], -17);
    EXPECT_FLOAT_EQ(C.getData()[5], -27);    
    EXPECT_FLOAT_EQ(C.getData()[6], 24);
    EXPECT_FLOAT_EQ(C.getData()[7], 306);
    EXPECT_FLOAT_EQ(C.getData()[8], 460);

    float *dataE = Aurora::random(60);
    Aurora::Matrix E = Aurora::Matrix::fromRawData(dataE, 3, 4, 5);
    Aurora::Matrix block1 = E.block(0, 1, 2);
    for (size_t i = 0; i < block1.getDimSize(2); i++)
    {   
        for (size_t j = 0; j < block1.getDimSize(1); j++)
        {
            for (size_t k = 1; k < E.getDimSize(0); k++)
            {
                auto index1 = k+j*E.getDimSize(0)+i*E.getDimSize(1)*E.getDimSize(0);
                auto index2 = k-1+j*block1.getDimSize(0)+i*block1.getDimSize(1)*block1.getDimSize(0);
                EXPECT_FLOAT_AE(E[index1],block1[index2]);
            }
        }
    }
    E.setBlockValue(0, 1, 2,-1);
    for (size_t i = 0; i < block1.getDimSize(2); i++)
    {   
        for (size_t j = 0; j < block1.getDimSize(1); j++)
        {
            for (size_t k = 1; k < E.getDimSize(0); k++)
            {
                EXPECT_FLOAT_AE(E(k,i,j).toMatrix().getScalar(),-1);
            }
        }
    }
    auto z = Aurora::zeros(1,4,5);
    EXPECT_TRUE(E.setBlock(0, 0, 0,z));
    for (size_t i = 0; i < 4; i++)
    {
        for (size_t j = 0; j < 5; j++)
        {
            EXPECT_FLOAT_AE(E(0,i,j).toMatrix().getScalar(), 0.0);
        }
    }
    auto z2 = Aurora::zeros(3,2,5);
    EXPECT_FALSE(E.setBlock(1, 1, 2,z));
    EXPECT_TRUE(E.setBlock(1, 1, 2,z2));
    for (size_t j = 0; j < 5; j++)
    {
        for (size_t i = 1; i < 3; i++)
        {
            for (size_t k = 0; k < 3; k++)
            {
                EXPECT_FLOAT_AE(E(k,i,j).toMatrix().getScalar(), 0.0); 
            }
        }
    }
    auto o2 = Aurora::ones(3,4,2);
    EXPECT_FALSE(E.setBlock(2, 1, 2,z));
    EXPECT_FALSE(E.setBlock(2, 1, 2,z2));
    EXPECT_TRUE(E.setBlock(2, 1, 2,o2));
    for (size_t j = 1; j < 3; j++)
    {
        for (size_t i = 1; i < 3; i++)
        {
            for (size_t k = 0; k < 3; k++)
            {
                EXPECT_FLOAT_AE(E(k,i,j).toMatrix().getScalar(), 1.0); 
            }
        }
    }
    auto ZZ =  Aurora::zeros(1,20);
    auto OO = Aurora::ones(1,10);
    EXPECT_TRUE(ZZ.setBlock(1, 10, 19, OO));
    for (size_t i = 0; i < ZZ.getDataSize(); i++)
    {
        EXPECT_FLOAT_AE(ZZ[i], i>9?1.0:0.0);
    }


}