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



#define DISPLAY_MATRIX

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

    }

    static void TearDownTestCase() {
    }

    void SetUp() {
    }

    void TearDown() {
    }
};

TEST_F(Matrix_Test, MatrixCreate) {
    double *dataA = Aurora::malloc(9);
    double *dataB = new double[9];
    double *dataC = new double[9];
    double *dataD = new double[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());
        DISPLAY_MATRIX(A);
    }
// double [] 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());
        DISPLAY_MATRIX(B);
    }
// copy from double [] 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());
        DISPLAY_MATRIX(C);
    }
// 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());
        DISPLAY_MATRIX(C);
    }
// 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());
        DISPLAY_MATRIX(C);
    }
// 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());
        DISPLAY_MATRIX(C);
    }
// 3d matrix
    {
        double *tempData = new double[9];
        Aurora::Matrix C = Aurora::Matrix::fromRawData(tempData, 3, 3, 1);
        EXPECT_EQ(dataD, 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());
        DISPLAY_MATRIX(C);
    }
// 3d matrix 2
    {
        double *tempData = new double[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());
        DISPLAY_MATRIX(C);
    }
}

TEST_F(Matrix_Test, matrixSlice) {
    double *dataA = Aurora::malloc(8);
    double *dataB = Aurora::malloc(8);
    double *dataC = Aurora::malloc(8);
    for (int i = 0; i < 8; ++i) {
        dataA[i] = (double) (i);
        dataB[i] = (double) (1);
        dataC[i] = (double) (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]);
}

TEST_F(Matrix_Test, matrixOpertaor) {
//3D
    {
        double *dataA = new double[8];
        double *dataB = new double[8];
        double *dataD = new double[8];
        for (int i = 0; i < 8; ++i) {
            dataA[i] = (double) (i);
            dataB[i] = (double) (2);
        }
        Aurora::Matrix A = Aurora::Matrix::fromRawData(dataA, 2, 2, 2);
        DISPLAY_MATRIX(A);
        Aurora::Matrix B = Aurora::Matrix::fromRawData(dataB, 2, 2, 2);
        DISPLAY_MATRIX(B);
        Aurora::Matrix D = Aurora::Matrix::fromRawData(dataD, 4, 2);
        DISPLAY_MATRIX(D);
        auto C = (A * B) - B;
        DISPLAY_MATRIX(C);
        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
    {
        double *dataA = Aurora::malloc(8, true);
        double *dataB = Aurora::malloc(8, true);
        double *dataD = new double[8];
        for (int i = 0; i < 16; ++i) {
            dataA[i] = (double) (i + 1);
            dataB[i] = (double) (2);
            if (i < 8) dataD[i] = (double) (2);
        }
        Aurora::Matrix A = Aurora::Matrix::New(dataA, 2, 2, 2, Aurora::Complex);
        DISPLAY_MATRIX(A);
        Aurora::Matrix B = Aurora::Matrix::New(dataB, 2, 2, 2, Aurora::Complex);
        DISPLAY_MATRIX(B);
        Aurora::Matrix D = Aurora::Matrix::fromRawData(dataD, 2, 2, 2);
        DISPLAY_MATRIX(D);
        auto C = A + 5.0 - 2.0;
        DISPLAY_MATRIX(C);
        EXPECT_EQ(C.getData()[2], 6);
        EXPECT_EQ(C.getData()[3], 4);
        EXPECT_TRUE(C.isComplex());
        C = A * B - B;
        DISPLAY_MATRIX(C);
        EXPECT_EQ(C.getData()[2], -4);
        EXPECT_EQ(C.getData()[3], 12);
        EXPECT_TRUE(C.isComplex());
        C = A * B - 2.0 + 5.0;
        DISPLAY_MATRIX(C);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 14);
        C = A * B / 2.0 * 3.;
        DISPLAY_MATRIX(C);
        EXPECT_EQ(C.getData()[2], -3);
        EXPECT_EQ(C.getData()[3], 21);
        C = A * B / B;
        DISPLAY_MATRIX(C);
        EXPECT_EQ(C.getData()[2], 3);
        EXPECT_EQ(C.getData()[3], 4);
        C = A * B / (B * A / A);
        DISPLAY_MATRIX(C);
        EXPECT_EQ(C.getData()[2], 3);
        EXPECT_EQ(C.getData()[3], 4);

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

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

        C = (B * 1.0) * (D * 1.0);
        DISPLAY_MATRIX(C);
        EXPECT_EQ(C.getData()[2], 4);
        EXPECT_EQ(C.getData()[3], 4);
        C = (B * 1.0) / (D * 1.0);
        DISPLAY_MATRIX(C);
        EXPECT_EQ(C.getData()[2], 1);
        EXPECT_EQ(C.getData()[3], 1);
        C = (B * 1.0) + (D * 1.0);
        DISPLAY_MATRIX(C);
        EXPECT_EQ(C.getData()[2], 4);
        EXPECT_EQ(C.getData()[3], 2);
        C = (B * 1.0) - (D * 1.0);
        DISPLAY_MATRIX(C);
        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){
        double *dataA = new double[8];
        double *dataB = new double[8];
        for (int i = 0; i < 8; ++i) {
            dataA[i] = (double) (i);
            dataB[i] = (double) (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){
    double *dataA = new double[9]{1,6,9,4,1,0,5,8,1};
    double *dataB = new double[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_DOUBLE_EQ(C.getData()[0], 40);
    EXPECT_DOUBLE_EQ(C.getData()[1], 6);
    EXPECT_DOUBLE_EQ(C.getData()[2], 24);
    B.forceReshape(3, 1, 1);
    C = A.mul(B);
    EXPECT_DOUBLE_EQ(C.getData()[0], 24);
    EXPECT_DOUBLE_EQ(C.getData()[1], 32);
    EXPECT_DOUBLE_EQ(C.getData()[2], 12);
    double *dataD = new double[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_DOUBLE_EQ(C.getData()[0], 59.1);
    EXPECT_DOUBLE_EQ(C.getData()[1], 87.6);
    EXPECT_DOUBLE_EQ(C.getData()[2], 27.9);
    EXPECT_DOUBLE_EQ(C.getData()[3], 1);
    EXPECT_DOUBLE_EQ(C.getData()[4], -17);
    EXPECT_DOUBLE_EQ(C.getData()[5], -27);    
    EXPECT_DOUBLE_EQ(C.getData()[6], 24);
    EXPECT_DOUBLE_EQ(C.getData()[7], 306);
    EXPECT_DOUBLE_EQ(C.getData()[8], 460);
}