#include <gtest/gtest.h>
#include <chrono>
#include "AuroraDefs.h"
#include "CudaMatrix.h"
#include "Function.h"
#include "Matrix.h"
#include "TestUtility.h"

#include "Function2D.h"
#include "Function2D.cuh"

class Function2D_Cuda_Test:public ::testing::Test
{
protected:
    static void SetUpFunction2DCudaTester(){

    }
    static void TearDownTestCase(){
    }
    public:
    Aurora::Matrix B;
    Aurora::CudaMatrix dB;

    void SetUp(){

    }
    void TearDown(){
    }

};

TEST_F(Function2D_Cuda_Test, min)
{
    // big data for test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix,
    //     Aurora::FunctionDirection direction, long &rowIdx, long &colIdx)
    {
        float *dataB = Aurora::random(4096*41472);
        B = Aurora::Matrix::fromRawData(dataB, 4096, 41472);
        dB = B.toDeviceMatrix();
        long r,c;
        // column
        auto ret1 = Aurora::min(B, Aurora::Column,r,c);

        auto ret2 = Aurora::min(dB, Aurora::Column,r,c);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }

        //row
        ret1 = Aurora::min(B, Aurora::FunctionDirection::Row,r,c);

        ret2 = Aurora::min(dB, Aurora::FunctionDirection::Row,r,c);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
        
        //ALL
        long r1,c1;
        ret1 = Aurora::min(B, Aurora::All,r,c);

        ret2 = Aurora::min(dB, Aurora::All,r1,c1);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));
        ASSERT_EQ(r, r1);
        ASSERT_EQ(c, c1);

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // different size speed 
    // Aurora::Matrix Aurora::min(const Aurora::Matrix &aMatrix,
    //     Aurora::FunctionDirection direction, long &rowIdx, long &colIdx)
    // in col wise
    {
        float *dataB = Aurora::random(3157*111);
        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        dB = B.toDeviceMatrix();
        long r,c;

        auto ret1 = Aurora::min(B, Aurora::FunctionDirection::Column,r,c);

        auto ret2 = Aurora::min(dB, Aurora::FunctionDirection::Column,r,c);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
        B.forceReshape( 111,3157, 1);
        dB = B.toDeviceMatrix();

        ret1 = Aurora::min(B, Aurora::FunctionDirection::Column,r,c);

        ret2 = Aurora::min(dB, Aurora::FunctionDirection::Column,r,c);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
        //ALL
        long r1,c1;
        ret1 = Aurora::min(B, Aurora::All,r,c);

        ret2 = Aurora::min(dB, Aurora::All,r1,c1);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));
        ASSERT_EQ(r, r1);
        ASSERT_EQ(c, c1);

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix, float aValue)
    {
        float *dataB = Aurora::random(3157*111);
        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        dB = B.toDeviceMatrix();
        long r,c;
        auto  start_time_ = std::chrono::high_resolution_clock::now();

        auto ret1 = Aurora::min(B, 500.5f);


        auto ret2 = Aurora::min(dB, 500.5f);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix, 
    //                          const Aurora::Matrix &aOther)
    // with same size matrix
    {
        float *dataB = Aurora::random(3157*111);
        float *dataA = Aurora::random(3157*111);
        auto A = Aurora::Matrix::fromRawData(dataA, 3157, 111);

        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        auto dA = A.toDeviceMatrix();
        dB = B.toDeviceMatrix();
        long r,c;

        auto ret1 = Aurora::min(B, A);

        auto ret2 = Aurora::min(dB, dA);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix, 
    //                          const Aurora::Matrix &aOther)
    // with col-vec and matrix 
    {
        float *dataB = Aurora::random(3157*111);
        float *dataA = Aurora::random(3157);
        auto A = Aurora::Matrix::fromRawData(dataA, 3157, 1);

        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        auto dA = A.toDeviceMatrix();
        dB = B.toDeviceMatrix();
        long r,c;

        auto ret1 = Aurora::min(B, A);

        auto ret2 = Aurora::min(dB, dA);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }

        ret2 = Aurora::min(dA, dB);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix, 
    //                          const Aurora::Matrix &aOther)
    // with row-vec and matrix 
    {
        float *dataB = Aurora::random(3157*111);
        float *dataA = Aurora::random(111);
        auto A = Aurora::Matrix::fromRawData(dataA, 1, 111);

        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        auto dA = A.toDeviceMatrix();
        dB = B.toDeviceMatrix();
        long r,c;
        auto  start_time_ = std::chrono::high_resolution_clock::now();

        auto ret1 = Aurora::min(B, A);

        auto ret2 = Aurora::min(dB, dA);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
}

TEST_F(Function2D_Cuda_Test, max)
{
    // big data for test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix,
    //     Aurora::FunctionDirection direction, long &rowIdx, long &colIdx)
    {
        float *dataB = Aurora::random(4096*41472);
        B = Aurora::Matrix::fromRawData(dataB, 4096, 41472);
        dB = B.toDeviceMatrix();
        long r,c;
        auto ret1 = Aurora::max(B, Aurora::FunctionDirection::Column,r,c);

        auto ret2 = Aurora::max(dB, Aurora::FunctionDirection::Column,r,c);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }

        ret1 = Aurora::max(B, Aurora::FunctionDirection::Row,r,c);

        ret2 = Aurora::max(dB, Aurora::FunctionDirection::Row,r,c);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
        //ALL
        long r1,c1;
        ret1 = Aurora::max(B, Aurora::All,r,c);

        ret2 = Aurora::max(dB, Aurora::All,r1,c1);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));
        ASSERT_EQ(r, r1);
        ASSERT_EQ(c, c1);

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // different size speed 
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix,
    //     Aurora::FunctionDirection direction, long &rowIdx, long &colIdx)
    // in col wise
    {
        float *dataB = Aurora::random(3157*111);
        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        dB = B.toDeviceMatrix();
        long r,c;

        auto ret1 = Aurora::max(B, Aurora::FunctionDirection::Column,r,c);

        auto ret2 = Aurora::max(dB, Aurora::FunctionDirection::Column,r,c);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
        B.forceReshape( 111,3157, 1);
        dB = B.toDeviceMatrix();

        ret1 = Aurora::max(B, Aurora::FunctionDirection::Column,r,c);

        ret2 = Aurora::max(dB, Aurora::FunctionDirection::Column,r,c);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }

        //ALL
        long r1,c1;
        ret1 = Aurora::max(B, Aurora::All,r,c);

        ret2 = Aurora::max(dB, Aurora::All,r1,c1);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));
        ASSERT_EQ(r, r1);
        ASSERT_EQ(c, c1);

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix, float aValue)
    {
        float *dataB = Aurora::random(3157*111);
        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        dB = B.toDeviceMatrix();
        long r,c;

        auto ret1 = Aurora::max(B, 500.5f);

        auto ret2 = Aurora::max(dB, 500.5f);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix, 
    //                          const Aurora::Matrix &aOther)
    // with same size matrix
    {
        float *dataB = Aurora::random(3157*111);
        float *dataA = Aurora::random(3157*111);
        auto A = Aurora::Matrix::fromRawData(dataA, 3157, 111);

        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        auto dA = A.toDeviceMatrix();
        dB = B.toDeviceMatrix();
        long r,c;
        auto  start_time_ = std::chrono::high_resolution_clock::now();

        auto ret1 = Aurora::max(B, A);

        auto ret2 = Aurora::max(dB, dA);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix, 
    //                          const Aurora::Matrix &aOther)
    // with col-vec and matrix 
    {
        float *dataB = Aurora::random(3157*111);
        float *dataA = Aurora::random(3157);
        auto A = Aurora::Matrix::fromRawData(dataA, 3157, 1);
        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        auto dA = A.toDeviceMatrix();
        dB = B.toDeviceMatrix();
        long r,c;

        auto ret1 = Aurora::max(B, A);

        // mat x vec
        auto ret2 = Aurora::max(dB, dA);
        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
        // vec x mat
        ret2 = Aurora::max(dA, dB);
        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));
        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
    // test
    // Aurora::Matrix Aurora::max(const Aurora::Matrix &aMatrix, 
    //                          const Aurora::Matrix &aOther)
    // with row-vec and matrix 
    {
        float *dataB = Aurora::random(3157*111);
        float *dataA = Aurora::random(111);
        auto A = Aurora::Matrix::fromRawData(dataA, 1, 111);

        B = Aurora::Matrix::fromRawData(dataB, 3157, 111);
        auto dA = A.toDeviceMatrix();
        dB = B.toDeviceMatrix();
        long r,c;
        auto  start_time_ = std::chrono::high_resolution_clock::now();

        auto ret1 = Aurora::max(B, A);
        // mat x vec
        auto ret2 = Aurora::max(dB, dA);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
        //vec x mat
        ret2 = Aurora::max(dA, dB);

        ASSERT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        ASSERT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        ASSERT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            ASSERT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
    }
}

TEST_F(Function2D_Cuda_Test, sum)
{
    //
    {
        // float *dataB = Aurora::random(4096*50000);
        float* dataB = new float[4096*5000];
        for (size_t i = 0; i < 4096*5000; i++)
        {
            dataB[i] = (i%2==0?1.0f:0.0f);
        }
        
        B = Aurora::Matrix::fromRawData(dataB, 4096, 5000);
        // B = Aurora::Matrix::fromRawData(dataB, 200, 200);

        auto dD = B.toDeviceMatrix();

        auto ret1 = Aurora::sum(B, Aurora::FunctionDirection::All);

        auto ret2 = Aurora::sum(dD, Aurora::FunctionDirection::All);

        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }


        ret1 = Aurora::sum(B, Aurora::FunctionDirection::Column);

        ret2 = Aurora::sum(dD, Aurora::FunctionDirection::Column);

        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i))
        }


        ret1 = Aurora::sum(B, Aurora::FunctionDirection::Row);

        ret2 = Aurora::sum(dD, Aurora::FunctionDirection::Row);

        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i))
        }
    }
    //complex type
    {
        float* dataB = new float[3000*2000*2];
        for (size_t i = 0; i < 3000*4000; i++)
        {
            dataB[i] = i%2==0?2.0f:1.0f;
        }
        
        B = Aurora::Matrix::fromRawData(dataB,3000, 2000,1,Aurora::Complex);

        auto dD = B.toDeviceMatrix();

        auto ret1 = Aurora::sum(B, Aurora::FunctionDirection::All);

        auto ret2 = Aurora::sum(dD, Aurora::FunctionDirection::All);

        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize()*2; i++)
        {
            EXPECT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }

        ret1 = Aurora::sum(B, Aurora::FunctionDirection::Column);

        ret2 = Aurora::sum(dD, Aurora::FunctionDirection::Column);
        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i))
        }


        ret1 = Aurora::sum(B, Aurora::FunctionDirection::Row);

        ret2 = Aurora::sum(dD, Aurora::FunctionDirection::Row);

        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i))
        }
    }
}

TEST_F(Function2D_Cuda_Test, mean)
{
    //
    {
        float* dataB = new float[4096*500];
        for (size_t i = 0; i < 4096*500; i++)
        {
            dataB[i] = (float)(i%2==0?1:0);
        }
        
        B = Aurora::Matrix::fromRawData(dataB, 4096, 500);
        // B = Aurora::Matrix::fromRawData(dataB, 200, 200);

        auto dD = B.toDeviceMatrix();

        auto ret1 = Aurora::mean(B, Aurora::FunctionDirection::All);

        auto ret2 = Aurora::mean(dD, Aurora::FunctionDirection::All);

        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }


        ret1 = Aurora::mean(B, Aurora::FunctionDirection::Column);

        ret2 = Aurora::mean(dD, Aurora::FunctionDirection::Column);

        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i))
        }


        ret1 = Aurora::mean(B, Aurora::FunctionDirection::Row);

        ret2 = Aurora::mean(dD, Aurora::FunctionDirection::Row);

        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));

        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i))
        }
    }
}

TEST_F(Function2D_Cuda_Test, sort)
{
    //
    {
        float* dataB = Aurora::random(25000000);
        B = Aurora::Matrix::fromRawData(dataB, 500, 500);
        // B = Aurora::Matrix::fromRawData(dataB, 200, 200);

        auto dD = B.toDeviceMatrix();

        auto ret1 = Aurora::sort(B, Aurora::Column);

        auto ret2 = Aurora::sort(dD,Aurora::Column);

        EXPECT_EQ(ret1.getDimSize(0),ret2.getDimSize(0));
        EXPECT_EQ(ret1.getDimSize(1),ret2.getDimSize(1));
        EXPECT_EQ(ret1.getDimSize(2),ret2.getDimSize(2));
        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_EQ(ret1[i], ret2.getValue(i))<<", index at :"<<i;
        }
        
    }
}

TEST_F(Function2D_Cuda_Test, fft)
{
    {
        float *input =  new float[20]{1,1,0,2,2,0,1,1,0,2,1,1,0,2,2,0,1,1,0,2};
        auto ma = Aurora::Matrix::fromRawData(input,10,2).toDeviceMatrix();
        auto fftrett = Aurora::fft(ma);
        auto fftrettH = fftrett.toHostMatrix();
        std::complex<float>* result = (std::complex<float>*)fftrettH.getData();
        //检验fft结果与matlab是否对应
        EXPECT_FLOAT_EQ(0.0729, fourDecimalRound(result[1].real()));
        EXPECT_FLOAT_EQ(2.4899, fourDecimalRound(result[2].imag()));
        EXPECT_FLOAT_EQ(0.0729, fourDecimalRound(result[11].real()));
        EXPECT_FLOAT_EQ(2.4899, fourDecimalRound(result[12].imag()));
        //检验fft的结果是否共轭
        EXPECT_FLOAT_EQ(0, result[4].imag()+result[6].imag());
        EXPECT_FLOAT_EQ(0, result[4].real()-result[6].real());


        auto ret= Aurora::ifft(fftrett);
        auto retH = ret.toHostMatrix();
        std::complex<float>* ifftResult = (std::complex<float>*)retH.getData();
        EXPECT_FLOAT_EQ(fourDecimalRound(ifftResult[1].real()),1.0);
        EXPECT_FLOAT_EQ(fourDecimalRound(ifftResult[3].real()),2.0);
        EXPECT_FLOAT_EQ(fourDecimalRound(ifftResult[11].real()),1.0);
        EXPECT_FLOAT_EQ(fourDecimalRound(ifftResult[13].real()),2.0);
        auto fftrett2 = Aurora::fft(ma,5);
        auto fftrett2H =fftrett2.toHostMatrix();

        auto result2 = (std::complex<float>*)fftrett2H.getData();
        //检验fft结果与matlab是否对应
        EXPECT_FLOAT_EQ(0.3090, fourDecimalRound(result2[1].real()));
        EXPECT_FLOAT_EQ(-1.3143, fourDecimalRound(result2[2].imag()));
        EXPECT_FLOAT_EQ(-0.8090, fourDecimalRound(result2[7].real()));
        EXPECT_FLOAT_EQ(1.3143, fourDecimalRound(result2[8].imag()));

        auto fftrett3 = Aurora::fft(ma,12);
        auto fftrett3H = fftrett3.toHostMatrix();

        auto result3 = (std::complex<float>*)fftrett3H.getData();
        //检验fft结果与matlab是否对应
        EXPECT_FLOAT_EQ(-1.0, fourDecimalRound(result3[1].real()));
        EXPECT_FLOAT_EQ(-3.4641, fourDecimalRound(result3[4].imag()));
        EXPECT_FLOAT_EQ(-1.0, fourDecimalRound(result3[13].real()));
        EXPECT_FLOAT_EQ(-3.4641, fourDecimalRound(result3[16].imag()));
    }
    {
        float *input =  new float[20]{1,1,0,2,2,0,1,1,0,2,1,1,0,2,2,0,1,1,0,2};
        auto ma = Aurora::Matrix::fromRawData(input,10,1, 1,Aurora::Complex);
        auto maD = ma.toDeviceMatrix();
        auto ret1 = Aurora::fft(ma);
        auto ret2 = Aurora::fft(maD);
        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i));
        }
        ret1 = Aurora::fft(ma,12);
        ret2 = Aurora::fft(maD,12);
        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i));
        }
        ret1 = Aurora::fft(ma,8);
        ret2 = Aurora::fft(maD,8);
        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i));
        }

    }
    {
        float *input =  new float[20]{1,1,0,2,2,0,1,1,0,2,1,1,0,2,2,0,1,1,0,2};
        auto ma = Aurora::Matrix::fromRawData(input,10,1,1, Aurora::Complex);
        auto maD = ma.toDeviceMatrix();
        auto ret1 = Aurora::ifft(Aurora::fft(ma),12);
        auto ret2 = Aurora::ifft(Aurora::fft(maD),12);
        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i));
        }
        ret1 = Aurora::ifft(Aurora::fft(ma),8);
        ret2 = Aurora::ifft(Aurora::fft(maD),8);
        for (size_t i = 0; i < ret1.getDataSize(); i++)
        {
            EXPECT_FLOAT_AE(ret1[i], ret2.getValue(i));
        }

    }
    {
        float * data = new float[36];
        for (size_t i = 0; i < 18; i++)
        {
            data[i*2] = i%9+1;
            data[i*2+1] = 1;
        }
        auto fm = Aurora::Matrix::fromRawData(data,9,2,1,Aurora::Complex);
        auto fmD = fm.toDeviceMatrix();
        Aurora::fftshift(fm);
        Aurora::fftshift(fmD);
        for (size_t i = 0; i < fm.getDataSize()*2; i++)
        {
            EXPECT_FLOAT_AE(fm[i], fmD.getValue(i));
        }
        Aurora::ifftshift(fm);
        Aurora::ifftshift(fmD);
        for (size_t i = 0; i < fm.getDataSize()*2; i++)
        {
            EXPECT_FLOAT_AE(fm[i], fmD.getValue(i));
        }
    }
    {
        float * data = new float[250000];
        for (size_t i = 0; i < 250000; i++)
        {
            data[i] = i%500;
        }
        auto fm = Aurora::Matrix::fromRawData(data,500,250,1,Aurora::Complex);
        auto fmD = fm.toDeviceMatrix();
        Aurora::fftshift(fm);
        Aurora::fftshift(fmD);
        for (size_t i = 0; i < fm.getDataSize()*2; i++)
        {
            EXPECT_FLOAT_AE(fm[i], fmD.getValue(i));
        }
        Aurora::ifftshift(fm);
        Aurora::ifftshift(fmD);
        for (size_t i = 0; i < fm.getDataSize()*2; i++)
        {
            EXPECT_FLOAT_AE(fm[i], fmD.getValue(i));
        }
    }
    {
        float * data = new float[333*500];
        for (size_t i = 0; i < 333*500; i++)
        {
            data[i] = i%333;
        }
        auto fm = Aurora::Matrix::fromRawData(data,333,250,1,Aurora::Complex);
        auto fmD = fm.toDeviceMatrix();
        Aurora::fftshift(fm);
        Aurora::fftshift(fmD);
        for (size_t i = 0; i < fm.getDataSize()*2; i++)
        {
            ASSERT_FLOAT_EQ(fm[i], fmD.getValue(i))<<"index:"<<i;
        }
        Aurora::ifftshift(fm);
        Aurora::ifftshift(fmD);
        for (size_t i = 0; i < fm.getDataSize()*2; i++)
        {
            EXPECT_FLOAT_AE(fm[i], fmD.getValue(i));
        }

    }
}

TEST_F(Function2D_Cuda_Test, immse) {
    auto matrixHost1 = Aurora::Matrix::fromRawData(Aurora::random(10000), 50,200);
    auto matrixHost2 = Aurora::Matrix::fromRawData(Aurora::random(10000), 50,200);
    auto matrixDevice1 = matrixHost1.toDeviceMatrix();
    auto matrixDevice2 = matrixHost1.toDeviceMatrix();
    auto result1 = Aurora::immse(matrixHost1, matrixHost2);
    auto result2 = Aurora::immse(matrixDevice1, matrixDevice2);
    EXPECT_FLOAT_AE(result1, result2);
}

TEST_F(Function2D_Cuda_Test, sortRows) {
    auto matrixHost1 = Aurora::Matrix::fromRawData(Aurora::random(10000), 50,200);
    Aurora::Matrix matrixHost2;
    auto matrixDevice1 = matrixHost1.toDeviceMatrix();
    Aurora::CudaMatrix matrixDevice2;
    auto result1 = Aurora::sortrows(matrixHost1, &matrixHost2);
    auto result2 = Aurora::sortrows(matrixDevice1, matrixDevice2).toHostMatrix();
    auto result3 = matrixHost2;
    auto result4 = matrixDevice2.toHostMatrix();
    ASSERT_FLOAT_EQ(result1.getDataSize(), result2.getDataSize());
    for (size_t i = 0; i < result1.getDataSize(); i++)
    {
        ASSERT_FLOAT_EQ(result1[i], result2[i]);
    }

    ASSERT_FLOAT_EQ(result3.getDataSize(), result4.getDataSize());
    for (size_t i = 0; i < result3.getDataSize(); i++)
    {
        ASSERT_FLOAT_EQ(result3[i], result4[i]);
    }
}

TEST_F(Function2D_Cuda_Test, inv) {
    auto matrixHost = Aurora::Matrix::fromRawData(new float[16]{4,6,7,8,9,3,7,5,4,3,2,1,2,3,4,5}, 4,4);
    auto matrixDevice = matrixHost.toDeviceMatrix();
    auto result1 = Aurora::inv(matrixHost);
    auto result2 = Aurora::inv(matrixDevice).toHostMatrix();
    ASSERT_FLOAT_EQ(result1.getDataSize(), result2.getDataSize());
    for (size_t i = 0; i < result1.getDataSize(); i++)
    {
        EXPECT_FLOAT_AE(result1[i], result2[i]);
    }
}

TEST_F(Function2D_Cuda_Test, dot) {
    auto matrixHost1 = Aurora::Matrix::fromRawData(new float[200000], 1000,200);
    auto matrixHost2 = Aurora::Matrix::fromRawData(new float[200000], 1000,200);
    for(unsigned int i=0; i<200000;++i)
    {
        matrixHost1[i] = i;
        matrixHost2[i] = i + 1;
    }
    auto matrixDevice1 = matrixHost1.toDeviceMatrix();
    Aurora::CudaMatrix matrixDevice2 = matrixHost2.toDeviceMatrix();
    auto result1 = Aurora::dot(matrixHost1, matrixHost2);
    auto result2 = Aurora::dot(matrixDevice1, matrixDevice2).toHostMatrix();
    ASSERT_FLOAT_EQ(result1.getDataSize(), result2.getDataSize());
    for (size_t i = 0; i < result1.getDataSize(); i++)
    {
        EXPECT_FLOAT_AE(result1[i], result2[i]);
    }
}

TEST_F(Function2D_Cuda_Test, prod) {
    auto matrixHost = Aurora::Matrix::fromRawData(new float[20], 4,5);
    for(unsigned int i=0; i<20;++i)
    {
        matrixHost[i] = i + 1;
    }
    auto matrixDevice = matrixHost.toDeviceMatrix();
    auto result1 = Aurora::prod(matrixHost);
    auto result2 = Aurora::prod(matrixDevice).toHostMatrix();
    ASSERT_FLOAT_EQ(result1.getDataSize(), result2.getDataSize());
    for (size_t i = 0; i < result1.getDataSize(); i++)
    {
        EXPECT_FLOAT_AE(result1[i], result2[i]);
    }

    matrixHost = Aurora::Matrix::fromRawData(new float[20], 20,1);
    for(unsigned int i=0; i<20;++i)
    {
        matrixHost[i] = i + 1;
    }
    matrixDevice = matrixHost.toDeviceMatrix();
    result1 = Aurora::prod(matrixHost);
    result2 = Aurora::prod(matrixDevice).toHostMatrix();
    ASSERT_FLOAT_EQ(result1.getDataSize(), result2.getDataSize());
    for (size_t i = 0; i < result1.getDataSize(); i++)
    {
        EXPECT_FLOAT_AE(result1[i], result2[i]);
    }

    auto matrixHostComplex = Aurora::Matrix::fromRawData(new float[20], 2,5, 1,Aurora::Complex);
    for(unsigned int i=0; i<20;++i)
    {
        matrixHost[i] = i + 1;
    }

    matrixDevice = matrixHostComplex.toDeviceMatrix();
    result1 = Aurora::prod(matrixHostComplex);
    result2 = Aurora::prod(matrixDevice).toHostMatrix();
    ASSERT_FLOAT_EQ(result1.getDataSize(), result2.getDataSize());
    ASSERT_FLOAT_EQ(result1.getValueType(), result2.getValueType());
    for (size_t i = 0; i < result1.getDataSize() * result1.getValueType(); i++)
    {
        EXPECT_FLOAT_AE(result1[i], result2[i]);
    }
}

TEST_F(Function2D_Cuda_Test, sub2ind) {
    float* dI1= new float[4]{1,2,1,2};
    Aurora::Matrix I1(std::shared_ptr<float>(dI1,std::default_delete<float[]>()),std::vector<int>{4});
    float* dI2= new float[4]{2,2,1,1};
    Aurora::Matrix I2(std::shared_ptr<float>(dI2,std::default_delete<float[]>()),std::vector<int>{4});
    float* dI4= new float[4]{1,1,2,2};
    Aurora::Matrix I3(std::shared_ptr<float>(dI4,std::default_delete<float[]>()),std::vector<int>{4});
    float* dsz= new float[3]{2,2,2};
    Aurora::Matrix sz(std::shared_ptr<float>(dsz,std::default_delete<float[]>()),std::vector<int>{3});
    auto result1 = Aurora::sub2ind(sz, {I1, I2, I3});
    auto result2 = Aurora::sub2ind(sz.toDeviceMatrix(), {I1.toDeviceMatrix(), I2.toDeviceMatrix(), I3.toDeviceMatrix()}).toHostMatrix();
    EXPECT_FLOAT_EQ(result1.getDataSize(), result2.getDataSize());
    for(unsigned int i=0; i<result1.getDataSize(); ++i)
    {
        EXPECT_FLOAT_EQ(result1[i], result2[i]);
    }
}

TEST_F(Function2D_Cuda_Test, ifft_symmetric) {
    float *input =  new float[18]{10,2,1,3,4,4,16,3,1,2,15,-2,1,-3,4,-4,1,-3};
    auto matrixHost = Aurora::Matrix::fromRawData(input,9,1,1,Aurora::Complex);
    auto matrixDevice = matrixHost.toDeviceMatrix();
    auto result1 = Aurora::ifft_symmetric(matrixHost,18);
    auto result2 = Aurora::ifft_symmetric(matrixDevice,18).toHostMatrix();
    EXPECT_FLOAT_EQ(result1.getDataSize(),result2.getDataSize());
    for(unsigned int i=0; i<result1.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(result1[i], result2[i]);
    }
}

TEST_F(Function2D_Cuda_Test, std){
    float *dataMA=  new float [9]{1, 2, 3, 2, 2, 6, 3, 3, 6};
    auto A = Aurora::Matrix::fromRawData(dataMA,3,3);
    auto D= Aurora::std(A.toDeviceMatrix());
    EXPECT_FLOAT_EQ(1.0, D.getValue(0));
    EXPECT_FLOAT_EQ(2.3094, fourDecimalRound(D.getValue(1)));
    EXPECT_FLOAT_EQ(1.7321, fourDecimalRound(D.getValue(2)));

}
