#include <cstddef>
#include <gtest/gtest.h>
#include <limits>

#include "CudaMatrix.h"
#include "Function1D.h"
#include "Function3D.h"
#include "MatlabReader.h"
#include "MatlabWriter.h"
#include "Matrix.h"
#include "Sparse.h"
#include "common/getMeasurementMetaData.h"
#include "config/config.h"
#include "reflectionReconstruction/preprocessData/preprocessAScanBlockForReflection.h"
#include "transmissionReconstruction/reconstruction/buildMatrix/DGradient.h"
#include "transmissionReconstruction/reconstruction/buildMatrix/FMM.h"
#include "transmissionReconstruction/reconstruction/buildMatrix/buildMatrix.h"

#include "transmissionReconstruction/reconstruction/buildMatrix/buildMatrix.cuh"

#include "transmissionReconstruction/reconstruction/reconstruction.h"
#include "transmissionReconstruction/reconstruction/solvingEquationSystem/TVAL/TVAL.h"

#include "reflectionReconstruction/preprocessData/preprocessTransmissionReconstructionForReflection.h"
#include "reflectionReconstruction/preprocessData/precalcImageParameters.h"
#include "reflectionReconstruction/reconstructionSAFT/reconstructionSAFT.h"

#include "reflectionReconstruction/preprocessData/determineOptimalPulse.h"

#include <cusparse.h>
#include <mkl_spblas.h>




inline float fourDecimalRound(float src){
    return round(src*10000.0)/10000.0;
}

#define EXPECT_FLOAT_AE(valueA,valueB)\
    EXPECT_FLOAT_EQ(fourDecimalRound(valueA),fourDecimalRound(valueB))

    #define ASSERT_FLOAT_AE(valueA,valueB)\
    ASSERT_FLOAT_EQ(fourDecimalRound(valueA),fourDecimalRound(valueB))

class Reconstruction_Test : public ::testing::Test {
protected:
    static void SetUpReconstructionTester() {

    }

    static void TearDownTestCase() {
    }

    void SetUp() {
    }

    void TearDown() {
    }
};

TEST_F(Reconstruction_Test, GPUMajor) {
	int cudaDevice;
    cudaGetDevice(&cudaDevice);
    cudaDeviceProp prop;
	cudaGetDeviceProperties(&prop, cudaDevice);
    EXPECT_GT(prop.major,2)<<"GPU major:"<<prop.major;
}

TEST_F(Reconstruction_Test, determineOptimalPulse) {
    Recon::reflectParams::imageResolution = 8.925316499483377e-04;
    Recon::reflectParams::optPulseFactor = 48;
    auto result = Recon::determineOptimalPulse(1.0000e-07,4096);
    EXPECT_EQ(4096, result.getDataSize());
    MatlabReader m2("/home/krad/TestData/sincPeakft.mat");
    auto f1 = m2.read("sincPeak_ft");
    for(size_t i=0; i<f1.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(f1[i], 0)<<"index:"<<i;
    }
}

TEST_F(Reconstruction_Test, reconstructionSAFT) {

     Recon::initalizeConfig("");
     MatlabReader m("/home/sun/testData/recontructSAFT.mat");
     auto Ascans = m.read("AScans");
     auto AttenuationMap3D = m.read("AttenuationMap3D");
     auto begin_TransMap = m.read("begin_TransMap");
     auto delta_TransMap = m.read("delta_TransMap");
     auto EnvOut = m.read("Env");
     auto mpIndex = m.read("mpIndex");
     auto receiverPos = m.read("receiverPos");
     auto SAFT_mode = m.read("SAFT_mode");
     auto senderPos = m.read("senderPos");
     auto SpeedMap3D = m.read("SpeedMap3D");
     Recon::TransRecos transRecos;
     transRecos.speedMap3D = SpeedMap3D;
     transRecos.attenuationMap3D = AttenuationMap3D;
     transRecos.beginTransMap = begin_TransMap;
     transRecos.deltaTransMap = delta_TransMap[0];
     MatlabReader m1("/home/sun/testData/resampleTransmissionVolume.mat");
     auto maxEmitter = m1.read("maxEmitter");
     auto maxReceiver = m1.read("maxReceiver");
     auto minEmitter = m1.read("minEmitter");
     auto minReceiver = m1.read("minReceiver");
     Recon::GeometryInfo geom;
     geom.maxEmitter = maxEmitter;
     geom.minEmitter = minEmitter;
     geom.maxReceiver = maxReceiver;
     geom.minReceiver = minReceiver;
     geom.minSize = m1.read("minSize");
     geom.maxSize = m1.read("maxSize");
     precalcImageParameters(geom);
     Aurora::Matrix Env = Aurora::zeros((int)Recon::reflectParams::imageXYZ[0],(int)Recon::reflectParams::imageXYZ[1],(int)Recon::reflectParams::imageXYZ[2]);
     auto result = Recon::recontructSAFT(Ascans, senderPos, receiverPos, mpIndex, 2, transRecos, Env);
    for(size_t i=0; i<result.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(result[i], EnvOut[i])<<"index:"<<i;
    }

}

TEST_F(Reconstruction_Test, preprocessAScanBlockForReflection) {
    Recon::initalizeConfig("");
    MatlabReader m2("/home/krad/TestData/preprocessRefC.mat");
    auto blockedAScans = m2.read("blockedAScans");
    auto blockedMP = m2.read("blockedMP");
    auto blockedSL = m2.read("blockedSL");
    auto blockedSN = m2.read("blockedSN");
    auto blockedRL = m2.read("blockedRL");
    auto blockedRN = m2.read("blockedRN");
    auto senderPositionBlock = m2.read("blockedSenderPosition");
    auto receiverPositionBlock = m2.read("blockedReceiverPosition");
    auto channelBlock = m2.read("blockedChannels");
    auto gainBlock = m2.read("blockedGain");
    Recon::MeasurementInfo info;
    info.expectedAScanLength= 4096;
    info.Wavelength=3;
    Recon::PreComputes preComputes;
    preComputes.measuredCEUsed = true;
    preComputes.offset = 5.2000e-07;
    preComputes.timeInterval = 1.0000e-07;
    preComputes.matchedFilter =  m2.read("matchedFilter");
    MatlabReader m("/home/krad/TestData/sincPeakft.mat");
    preComputes.sincPeak_ft = m.read("sincPeak_ft");
    auto result = Recon::preprocessAScanBlockForReflection(blockedAScans, blockedMP, blockedSL, blockedSN, blockedRL,
     blockedRN,senderPositionBlock, receiverPositionBlock, gainBlock, channelBlock, info, preComputes);

    size_t size = result.AscanBlock.getDataSize();
    // for(size_t i=0; i<result.AscanBlock.getDataSize(); ++i)
    // {
    //     ASSERT_FLOAT_AE(f1[i], result[i])<<"index:"<<i;
    // }
}

TEST_F(Reconstruction_Test, reconstructArt) {
    MatlabReader m("/home/sun/testData/reconstructArt.mat");
    auto data  = m.read("data");
    auto dataAtt  = m.read("dataAtt");
    auto dims = m.read("dims");
    auto receiverList  = m.read("receiverList");
    auto res = m.read("res");
    auto senderList  = m.read("senderList");
    auto SOS_IN_WATER  = m.read("SOS_IN_WATER");
    
    MatlabReader m2("/home/krad/TestData/gpuresult.mat");
    auto f1 = m2.read("out");
    auto result = Recon::reconstructArt(data, dataAtt, dims, senderList, receiverList, res, SOS_IN_WATER[0]);
    for(size_t i=0; i<f1.getDataSize(); ++i)
    {
        ASSERT_FLOAT_AE(f1[i], result.outSOS[i])<<"index:"<<i;
    }
}

TEST_F(Reconstruction_Test, checkAndScale) {
    MatlabReader m("/home/sun/testData/buildMatrix.mat");
    auto i  = m.read("i1");
    auto j  = m.read("j1");
    auto s = m.read("s1");
    MatlabReader m2("/home/sun/testData/checkAndScale.mat");
    auto b = m2.read("b");

    Aurora::Sparse sparse(i,j,s,709613,1196032);
    Recon::checkAndScale(sparse,b,1196032);
}

TEST_F(Reconstruction_Test, buildMatrix) {
    Recon::transParams::bresenham = 1;
    MatlabReader m("/home/krad/TestData/buildmatrixFloat.mat");
    auto j1 = m.read("jCopy");
    auto i1 = m.read("iCopy");
    auto s1 = m.read("sCopy");
    auto dims = Aurora::Matrix::fromRawData(new float[3]{128, 128, 73}, 1, 3);;
    auto receiverList  = m.read("receiverList");
    auto potentialMap = m.read("potentialMap");
    auto senderList  = m.read("senderList"); 
    auto res  = m.read("res");
    auto result = Recon::buildMatrix(senderList, receiverList, res, dims, false, potentialMap);
    size_t count = 0;

    for(size_t i=0; i<i1.getDataSize(); ++i)
    {
        // EXPECT_FLOAT_AE(i1[i] - 1, result.M.getRowVector()[i])<<"ROW error,"<<" index:"<<i;
        if ((int)(i1[i] - 1)!= (int)result.M.getRowVector()[i]){
            count++;
        }
    }
    EXPECT_EQ(count, 0)<<"ROW error percent,"<<((double)count/(double)i1.getDataSize())*100.0;
    count=0;
    for(size_t i=0; i<j1.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(j1[i]- 1, result.M.getColVector()[i])<<"COL error,"<<" index:"<<i;
        if ((int)(j1[i] - 1)!= (int)result.M.getColVector()[i]){
            count++;
        }
    }
    EXPECT_EQ(count, 0)<<"COL error percent,"<<((double)count/(double)j1.getDataSize())*100.0;
    count=0;
    for(size_t i=0; i<s1.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(s1[i], result.M.getValVector()[i])<<"Value error,"<<" index:"<<i;
        if ((int)(s1[i] - 1)!= (int)result.M.getValVector()[i]){
            count++;
        }
    }
    EXPECT_EQ(count, 0)<<"Value error percent,"<<((double)count/(double)j1.getDataSize())*100.0;

}

TEST_F(Reconstruction_Test, buildMatrixCuda) {
    Recon::transParams::bresenham = 1;

    MatlabReader m("/home/krad/TestData/buildmatrixFloat.mat");
    auto j1 = m.read("jCopy");
    auto i1 = m.read("iCopy");
    auto s1 = m.read("sCopy");
    auto dims = Aurora::Matrix::fromRawData(new float[3]{128, 128, 73}, 1, 3).toDeviceMatrix();
    auto receiverList  = m.read("receiverList").toDeviceMatrix();
    auto potentialMap = m.read("potentialMap").toDeviceMatrix();
    auto senderList  = m.read("senderList").toDeviceMatrix(); 
    auto res  = m.read("res").toDeviceMatrix();
    auto result = Recon::buildMatrix(senderList, receiverList, res, dims, false, potentialMap);
    size_t count = 0;

    for(size_t i=0; i<i1.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(i1[i] - 1, result.M.getRowVector()[i])<<"ROW error,"<<" index:"<<i;
        if ((int)(i1[i] - 1)!= (int)result.M.getRowVector()[i]){
            count++;
        }
    }
    EXPECT_EQ(count, 0)<<"ROW error percent,"<<((double)count/(double)i1.getDataSize())*100.0;
    count=0;
    for(size_t i=0; i<j1.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(j1[i]- 1, result.M.getColVector()[i])<<"COL error,"<<" index:"<<i;
        if ((int)(j1[i] - 1)!= (int)result.M.getColVector()[i]){
            count++;
        }
    }
    EXPECT_EQ(count, 0)<<"COL error percent,"<<((double)count/(double)j1.getDataSize())*100.0;
    for(size_t i=0; i<s1.getDataSize(); ++i)
    {
        EXPECT_FLOAT_EQ(s1[i], result.M.getValVector()[i])<<"Value error,"<<" index:"<<i;
    }
    EXPECT_EQ(count, 0)<<"Value error percent,"<<((double)count/(double)j1.getDataSize())*100.0;

}

TEST_F(Reconstruction_Test, calculateMinimalMaximalTransducerPositions) {
    auto senderList  = Aurora::Matrix::fromRawData(new float[6]{1, 2, 3, 1, 2, 4}, 3, 2);
    auto receiverList  = Aurora::Matrix::fromRawData(new float[6]{1, 8, 3, 1, 2, 1}, 3, 2);
    auto result = Recon::calculateMinimalMaximalTransducerPositions(senderList,receiverList);
    EXPECT_FLOAT_EQ(1.0,result.getData()[0]);
    EXPECT_FLOAT_EQ(2,result.getData()[1]);
    EXPECT_FLOAT_EQ(1.0,result.getData()[2]);
    EXPECT_FLOAT_EQ(1.0,result.getData()[3]);
    EXPECT_FLOAT_EQ(8,result.getData()[4]);
    EXPECT_FLOAT_EQ(4,result.getData()[5]); 
}

TEST_F(Reconstruction_Test, calculateResolution) {
    auto ddims   = Aurora::Matrix::fromRawData(new float[6]{-0.1296,-0.1296,0.0185,0.1296,0.1295,0.1682}, 1, 6);
    auto dims  = Aurora::Matrix::fromRawData(new float[3]{128,128,74}, 1, 3);
    auto result  = Recon::calculateResolution(ddims, dims);
    EXPECT_FLOAT_AE(0.0020,result[0]);
    EXPECT_FLOAT_AE(0.0020,result[1]);
    EXPECT_FLOAT_AE(0.0021,result[2]);
}

TEST_F(Reconstruction_Test, getDimensions) {
    auto ddims = Aurora::Matrix::fromRawData(new float[6]{-0.1296,-0.1296,0.0185,0.1296,0.1295,0.1682}, 1, 6);
    float numPixelXY = 128;
    auto result  = Recon::getDimensions(numPixelXY,ddims);
    EXPECT_FLOAT_AE(128,result[0]);
    EXPECT_FLOAT_AE(128,result[1]);
    EXPECT_FLOAT_AE(74,result[2]);
}

TEST_F(Reconstruction_Test, discretizePositions) {
    MatlabReader m("/home/sun/testData/discretizePositions.mat");
    auto senderList  = m.read("senderList");
    auto receiverList  = m.read("receiverList");
    auto senderListResult = m.read("senderListResult");
    auto receiverListResult  = m.read("receiverListResult");
    auto dims = m.read("dims");
    auto ddims  = m.read("ddims");
    auto res  = m.read("res");
    auto result  = Recon::discretizePositions(senderList,receiverList,Recon::transParams::numPixelXY);

    
    EXPECT_FLOAT_AE(senderListResult.getDataSize(), result.senderCoordList.getDataSize());
    for(size_t i=0; i<senderListResult.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(senderListResult[i], result.senderCoordList[i]);
    }

    EXPECT_FLOAT_AE(receiverListResult.getDataSize(), result.receiverCoordList.getDataSize());
    for(size_t i=0; i<receiverListResult.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(receiverListResult[i], result.receiverCoordList[i]);
    }

    EXPECT_FLOAT_AE(dims.getDataSize(), result.dims.getDataSize());
    for(size_t i=0; i<dims.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(dims[i], result.dims[i]);
    }  

    EXPECT_FLOAT_AE(ddims.getDataSize(), result.ddims.getDataSize());
    for(size_t i=0; i<ddims.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(ddims[i], result.ddims[i]);
    } 

    EXPECT_FLOAT_AE(res.getDataSize(), result.res.getDataSize());
    for(size_t i=0; i<res.getDataSize(); ++i)
    {
        EXPECT_FLOAT_AE(res[i], result.res[i]);
    }           
}

TEST_F(Reconstruction_Test, DGradient) {
    MatlabReader m("/home/krad/TestData/DGradient.mat");

    auto x = m.read("x");
    auto y = m.read("y");
    // x.forceReshape(x.getDataSize(), 1, 1);
    auto result = Recon::DGradient(x,1,2);
    for (size_t i = 0; i < result.getDataSize(); i++)
    {
        EXPECT_FLOAT_AE(result.getData()[i],y.getData()[i])<<"index:"<<i;
    }
}

TEST_F(Reconstruction_Test, correctPath) {
    auto path  = Aurora::Matrix::fromRawData(new float[6]{1,1,10,9,1,1},2,3);
    auto startPt = Aurora::Matrix::fromRawData(new float[3]{1, .5, 1}, 1, 3);
    auto endPt = Aurora::Matrix::fromRawData(new float[3]{1, 10, 1}, 1, 3);

    Recon::correctPath(path,startPt,endPt);
    for (size_t i = 0; i < 10; i++)
    {
        EXPECT_FLOAT_EQ(path[i],1.0);
        EXPECT_FLOAT_EQ(path[i+10],1.0+i);
        EXPECT_FLOAT_EQ(path[i+20],1.0);
    }
}

TEST_F(Reconstruction_Test,getPixelLengthApproximation){
    auto startPt = Aurora::Matrix::fromRawData(new float[3]{1, 1, 1}, 1, 3);
    auto endPt = Aurora::Matrix::fromRawData(new float[3]{4, 5, 6}, 1, 3);
    auto res = Aurora::Matrix::fromRawData(new float[3]{.1, .1, .1}, 1, 3);
    auto weight = Recon::getPixelLengthApproximation(startPt, endPt, res, 10);
    EXPECT_FLOAT_AE(weight[0],.0707);
}

TEST_F(Reconstruction_Test,traceLine3D){
    auto p1 = Aurora::Matrix::fromRawData(new float[3]{1, 10, 12}, 1, 3);
    auto p2 = Aurora::Matrix::fromRawData(new float[3]{20, 2, 13}, 1, 3);
    auto discretization = Aurora::Matrix::fromRawData(new float[3]{1, 1, 1}, 1, 3);
    auto result = Recon::traceLine3D(p1, p2, discretization);
    MatlabReader m("/home/krad/TestData/traceLine3D.mat");

    auto path = m.read("path");
    auto ds = m.read("ds");
    for (size_t i = 0; i < result.path.getDataSize(); i++)
    {
        EXPECT_FLOAT_AE(result.path.getData()[i],path[i])<<"index:"<<i;
    }
    for (size_t i = 0; i < result.ds.getDataSize(); i++)
    {
        EXPECT_FLOAT_AE(result.ds.getData()[i],ds[i])<<"index:"<<i;
    }
}

TEST_F(Reconstruction_Test,traceStraightRay){
    auto p1 = Aurora::Matrix::fromRawData(new float[3]{1, 1, 1}, 1, 3);
    auto p2 = Aurora::Matrix::fromRawData(new float[3]{5, 6, 7}, 1, 3);
    auto res = Aurora::Matrix::fromRawData(new float[3]{.1, .1, .1}, 1, 3);
    auto dmis = Aurora::Matrix::fromRawData(new float[3]{5, 10, 20}, 1, 3);
    auto result = Recon::traceStraightRay(p1, p2, res, dmis);
    EXPECT_EQ(13, result.weighting.getDataSize());
    EXPECT_FLOAT_AE(result.weighting[0],0);
    EXPECT_FLOAT_AE(result.weighting[1],0.1462);
    EXPECT_FLOAT_AE(result.weighting[11],0.0292);
    EXPECT_FLOAT_AE(result.weighting[3],0.0439);
    EXPECT_EQ(13, result.path.getDimSize(0));
    EXPECT_EQ(3, result.path.getDimSize(1));
    EXPECT_FLOAT_AE(result.path[0],1);
    EXPECT_FLOAT_AE(result.path[15],2);
    EXPECT_FLOAT_AE(result.path[30],3);
    EXPECT_EQ(13, result.pathLen);
}

TEST_F(Reconstruction_Test,traceStraightRayBresenham){
    auto p1 = Aurora::Matrix::fromRawData(new float[3]{1, 1, 1}, 1, 3);
    auto p2 = Aurora::Matrix::fromRawData(new float[3]{5, 6, 7}, 1, 3);
    auto res = Aurora::Matrix::fromRawData(new float[3]{.1, .1, .1}, 1, 3);
    auto result = Recon::traceStraightRayBresenham(p1, p2, res);
    EXPECT_EQ(1, result.weighting.getDataSize());
    EXPECT_FLOAT_AE(result.weighting[0],0.1462);
    EXPECT_EQ(6, result.path.getDimSize(0));
    EXPECT_EQ(3, result.path.getDimSize(1));
    EXPECT_FLOAT_AE(result.path[0],1);
    EXPECT_FLOAT_AE(result.path[15],4);
    EXPECT_FLOAT_AE(result.path[7],2);
    EXPECT_EQ(6, result.pathLen);
}

TEST_F(Reconstruction_Test,callTval3){
    TVALOptions opt;
    opt.nonneg = false;
    opt.bent = false;
    opt.tol = 1E-10;
    opt.maxit = 50;
    opt.TVnorm = 2;
    opt.disp = false;
    opt.mu0 = 100;
    opt.mu = 100;
    opt.beta = 1;
    opt.beta0 = 1;
    MatlabReader m("/home/krad/TestData/ijs.mat");
    auto j = m.read("j");   
    auto i = m.read("i");
    auto s = m.read("s");
    Aurora::Sparse M(i-1,j-1,s,1900092,1196032);
    MatlabReader m2("/home/sun/TVALGPU.mat");
    auto j1 = m2.read("colPtr");
    auto i1 = m2.read("rowPtr");
    auto s1 = m2.read("data");

    
    int * yIdxs = new int[M.getColVector().getDataSize()];
    std::copy(M.getColVector().getData(),M.getColVector().getData()+M.getColVector().getDataSize(),yIdxs);
    int * xIdxs = new int[M.getRowVector().getDataSize()];
    std::copy(M.getRowVector().getData(),M.getRowVector().getData()+M.getRowVector().getDataSize(),xIdxs);
    Aurora::Matrix values = M.getValVector();
    size_t rows = M.getM(), cols = M.getN();        
    int nz = std::max(M.getColVector().getDataSize(),std::max(M.getRowVector().getDataSize(),M.getValVector().getDataSize()));        
     sparse_matrix_t A;
    sparse_matrix_t csrA;
    sparse_matrix_t csrA_t;
    sparse_status_t status;
    status = mkl_sparse_s_create_coo(&A, sparse_index_base_t::SPARSE_INDEX_BASE_ZERO,  rows,cols,  nz, xIdxs,yIdxs,values.getData());     
  
    status = mkl_sparse_convert_csr(A, sparse_operation_t::SPARSE_OPERATION_TRANSPOSE, &csrA); 
    // status = mkl_sparse_order(csrA);  

    int n_rows = 0;
    int n_cols = 0;
    int *col_start,*col_end,*row_idx;
    float * csrValues;
    sparse_index_base_t index;


    status = mkl_sparse_s_export_csr(csrA, &index, &n_cols, &n_rows, &col_start, &col_end, &row_idx, &csrValues);      

        
    int *col_idx = new int[n_cols+1];
    std::copy(col_start,col_start+ n_cols,col_idx);
    col_idx[n_cols] = nz;
    
    printf("j1 count:%zu,i1 count:%zu,s1 count:%zu\r\n",j1.getDataSize(),i1.getDataSize(),s1.getDataSize());
    printf("n_rows:%d,n_cols:%d\r\n", n_rows, n_cols);
    for (size_t k = 0; k < j1.getDataSize(); k++)
    {
        ASSERT_FLOAT_EQ(j1.getData()[k],col_idx[k])<<"col error, index:"<<k;
    }

    for (size_t k = 0; k < i1.getDataSize(); k++)
    {
        ASSERT_FLOAT_EQ(i1.getData()[k],row_idx[k])<<"row error, index:"<<k;
    }
    for (size_t k = 0; k < s1.getDataSize(); k++)
    {
        ASSERT_FLOAT_EQ(s1.getData()[k],csrValues[k])<<"value error, index:"<<k;
    }
    mkl_sparse_destroy(A);
    mkl_sparse_destroy(csrA);
    delete [] xIdxs;
    delete [] yIdxs;
}