#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, MatrixSub) {
    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)
    D = 0.5-C;
    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 E = C-0.5;
        auto dE2 =  (dA-dB)-0.5;
        auto dhE = dE2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(E[i],dhE[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 = C-A;
        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, MatrixDiv) {
    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)
    D = 0.5/C;
    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 E = C/0.5;
        auto dE2 =  (dA/dB)/0.5;
        auto dhE = dE2.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(E[i],dhE[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 = C/A;
        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, MatrixPow){
    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+0.2*i;
    }
    auto dA= A.toDeviceMatrix();
    {
        auto R = A^0;
        auto dR = dA^0;
        auto dhR = dR.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R = A^1;
        auto dR = dA^1;
        auto dhR = dR.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
        {
        auto R = A^2;
        auto dR = dA^2;
        auto dhR = dR.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R = A^3;
        auto dR = dA^3;
        auto dhR = dR.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R = A^5;
        auto dR = dA^5;
        auto dhR = dR.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }

}

TEST_F(CudaMatrix_Test, MatrixNeg){
    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+0.2*i;
    }
    auto dA= A.toDeviceMatrix();
    {
        auto R = -A;
        auto dR = -dA;
        auto dhR = dR.toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
}

TEST_F(CudaMatrix_Test, MatrixCompare){
    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+0.2*i;
    }
    auto dA= A.toDeviceMatrix();
    auto dB= B.toDeviceMatrix();
    {
        auto R= (A<B);
        auto dhR = (dA<dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (A>B);
        auto dhR = (dA>dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }    
    {
        auto R= (A<=B);
        auto dhR = (dA<=dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (A>=B);
        auto dhR = (dA>=dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (A==B);
        auto dhR = (dA==dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (A!=B);
        auto dhR = (dA!=dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (9<B);
        auto dhR = (9<dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (9>B);
        auto dhR = (9>dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (9<=B);
        auto dhR = (9<=dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (9>=B);
        auto dhR = (9>=dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (9==B);
        auto dhR = (9==dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[i]);
        }
    }
    {
        auto R= (9!=B);
        auto dhR = (9!=dB).toHostMatrix();
        for (size_t i = 0; i < 1000; i++)
        {
            ASSERT_FLOAT_EQ(R[i],dhR[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]);
        }       
    }

}