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


class CudaMatrix_Test : public ::testing::Test {
protected:
    static void SetUpCudaMatrixTester() {

    }

    static void TearDownTestCase() {
    }

    void SetUp() {
    }

    void TearDown() {
    }
};

TEST_F(CudaMatrix_Test, MatrixAdd) {
    auto A = Aurora::zeros(1000,1,1);
    auto B = Aurora::zeros(1000,1,1);
    for (size_t i = 0; i < 1000; i++)
    {
        A[i] = -1;
        B[i] =  i;
    }
    printf("Test CudaMatrix operator+(const CudaMatrix &aMatrix) const \r\n");
    //CudaMatrix operator+(const CudaMatrix &aMatrix) const
    auto C = A+B;
    auto dA = A.toDeviceMatrix();
    auto dB = B.toDeviceMatrix();
    auto dC = (dA+dB);
    auto dhC = dC.toHostMatrix();
    for (size_t i = 0; i < 1000; i++)
    {
        ASSERT_FLOAT_EQ(C[i],dhC[i]);
    }
    printf("Test CudaMatrix operator+(float aScalar) const \r\n");
    //CudaMatrix operator+(float aScalar) const
    auto D = C+0.5;
    auto dD = dC+0.5;
    auto dhD = dD.toHostMatrix();
    for (size_t i = 0; i < 1000; i++)
    {
        ASSERT_FLOAT_EQ(D[i],dhD[i]);
    }
    printf("Test CudaMatrix operator+(float aScalar, const CudaMatrix &aMatrix) \r\n");
    // CudaMatrix operator+(float aScalar, const CudaMatrix &aMatrix)
    dD = 0.5 + dC;
    dhD = dD.toHostMatrix();
    for (size_t i = 0; i < 1000; i++)
    {
        ASSERT_FLOAT_EQ(D[i],dhD[i]);
    }
    printf("Test CudaMatrix &operator+(float aScalar, CudaMatrix &&aMatrix) \r\n");
    // CudaMatrix &operator+(float aScalar, CudaMatrix &&aMatrix)
    {
        auto dD2 = 0.5 + (dA+dB);
        dhD = dD2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(D[i],dhD[i]);
        }
    }
    printf("Test CudaMatrix &operator+(CudaMatrix &&aMatrix, float aScalar) \r\n");
    // CudaMatrix &operator+(CudaMatrix &&aMatrix, float aScalar)
    {
        
        auto dD2 =  (dA+dB)+0.5;
        dhD = dD2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(D[i],dhD[i]);
        }
    }
    //CudaMatrix operator+(CudaMatrix &&aMatrix) const
    printf("Test CudaMatrix operator+(CudaMatrix &&aMatrix) const \r\n");
    {
        auto D = A+C;
        auto dD2 = dA+(dA+dB);
        dhD = dD2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(D[i],dhD[i]);
        }
    }
    //CudaMatrix operator+(CudaMatrix &&aMatrix,CudaMatrix &aOther)
    printf("Test CudaMatrix operator+(CudaMatrix &&aMatrix,CudaMatrix &aOther) \r\n");
    {
        auto D = A+C;
        auto dD2 = (dA+dB)+dA;
        dhD = dD2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(D[i],dhD[i]);
        }
    }
}

TEST_F(CudaMatrix_Test, MatrixMul) {
    auto A = Aurora::zeros(1000,1,1);
    auto B = Aurora::zeros(1000,1,1);
    for (size_t i = 0; i < 1000; i++)
    {
        A[i] = -1;
        B[i] =  i;
    }
    printf("Test CudaMatrix operator*(const CudaMatrix &aMatrix) const \r\n");
    //CudaMatrix operator*(const CudaMatrix &aMatrix) const
    auto C = A*B;
    auto dA = A.toDeviceMatrix();
    auto dB = B.toDeviceMatrix();
    auto dC = (dA*dB);
    auto dhC = dC.toHostMatrix();
    for (size_t i = 0; i < 1000; i++)
    {
        ASSERT_FLOAT_EQ(C[i],dhC[i]);
    }
    printf("Test CudaMatrix operator*(float aScalar) const \r\n");
    //CudaMatrix operator*(float aScalar) const
    auto D = C*0.5;
    auto dD = dC*0.5;
    auto dhD = dD.toHostMatrix();
    for (size_t i = 0; i < 1000; i++)
    {
        ASSERT_FLOAT_EQ(D[i],dhD[i]);
    }

    printf("Test CudaMatrix operator*(float aScalar, const CudaMatrix &aMatrix) \r\n");
    // CudaMatrix operator*(float aScalar, const CudaMatrix &aMatrix)
    dD = 0.5 * dC;
    dhD = dD.toHostMatrix();
    for (size_t i = 0; i < 1000; i++)
    {
        ASSERT_FLOAT_EQ(D[i],dhD[i]);
    }
    printf("Test CudaMatrix &operator*(float aScalar, CudaMatrix &&aMatrix) \r\n");
    // CudaMatrix &operator*(float aScalar, CudaMatrix &&aMatrix)
    {
        auto dD2 = 0.5 * (dA*dB);
        dhD = dD2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(D[i],dhD[i]);
        }
    }
    printf("Test CudaMatrix &operator*(CudaMatrix &&aMatrix, float aScalar) \r\n");
    // CudaMatrix &operator*(CudaMatrix &&aMatrix, float aScalar)
    {
        
        auto dD2 =  (dA*dB)*0.5;
        dhD = dD2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(D[i],dhD[i]);
        }
    }
    //CudaMatrix operator*(CudaMatrix &&aMatrix) const
    printf("Test CudaMatrix operator*(CudaMatrix &&aMatrix) const \r\n");
    {
        auto D1 = A*C;
        auto dD2 = dA*(dA*dB);
        dhD = dD2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(D1[i],dhD[i]);
        }
    }
    //CudaMatrix operator*(CudaMatrix &&aMatrix,CudaMatrix &aOther)
    printf("Test CudaMatrix operator*(CudaMatrix &&aMatrix,CudaMatrix &aOther) \r\n");
    {
        auto D1 = A*C;
        auto dD2 = (dA*dB)*dA;
        dhD = dD2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(D1[i],dhD[i]);
        }
    }
}

TEST_F(CudaMatrix_Test, matrixfunction)
{
    printf("Test CudaMatrix block(int aDim,int aBeginIndx, int aEndIndex) const\r\n");
    {
        float *dataA = new float[9]{1,6,9,4,1,0,5,8,1};
        float *dataB = new float[3]{1,2,3};
        Aurora::CudaMatrix A = Aurora::Matrix::fromRawData(dataA, 3, 3).toDeviceMatrix();
        Aurora::CudaMatrix B = Aurora::Matrix::fromRawData(dataB, 1, 3).toDeviceMatrix();

        Aurora::Matrix block1 = A.block(0, 1, 2).toHostMatrix();
        Aurora::Matrix block2 = A.toHostMatrix().block(0, 1, 2);
        for (size_t i = 0; i < block1.getDataSize(); i++)
        {   
            ASSERT_FLOAT_EQ(block1[i], block2[i]);
        }

        block2 = A.toHostMatrix();
        A.setBlockValue(0, 1, 2,-1);
        block1 = A.toHostMatrix();
        block2.setBlockValue(0, 1, 2,-1);
        for (size_t i = 0; i < block1.getDataSize(); i++)
        {   
            ASSERT_FLOAT_EQ(block1[i], block2[i]);
        }

        Aurora::CudaMatrix C = Aurora::zeros(2,3).toDeviceMatrix();
        EXPECT_TRUE(A.setBlock(0, 0, 1, C));
        block1 = A.block(0, 0, 1).toHostMatrix();
        block2 = C.toHostMatrix();
        for(size_t i = 0; i < C.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(block1[i], block2[i]);
        }       
    }

}