#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]);
        }
    }
}