#include "psrdada_cpp/effelsberg/edd/DadaBufferLayout.hpp"
#include "psrdada_cpp/effelsberg/edd/GatedSpectrometer.cuh"

#include "psrdada_cpp/dada_db.hpp"
#include "psrdada_cpp/dada_null_sink.hpp"
#include "psrdada_cpp/multilog.hpp"
#include "gtest/gtest.h"

#include "thrust/device_vector.h"
#include "thrust/extrema.h"

TEST(GatedSpectrometer, BitManipulationMacros) {
  for (int i = 0; i < 64; i++) {
    uint64_t v = 0;
    SET_BIT(v, i);

    for (int j = 0; j < 64; j++) {
      if (j == i)
        EXPECT_EQ(TEST_BIT(v, j), 1);
      else
        EXPECT_EQ(TEST_BIT(v, j), 0);
    }
  }
}


TEST(GatedSpectrometer, stokes_IQUV)
{
    float I,Q,U,V;
    // No field
    psrdada_cpp::effelsberg::edd::stokes_IQUV((float2){0.0f,0.0f}, (float2){0.0f,0.0f}, I, Q, U, V);
    EXPECT_FLOAT_EQ(I, 0);
    EXPECT_FLOAT_EQ(Q, 0);
    EXPECT_FLOAT_EQ(U, 0);
    EXPECT_FLOAT_EQ(V, 0);

    // For p1 = Ex, p2 = Ey
    // horizontal polarized
    psrdada_cpp::effelsberg::edd::stokes_IQUV((float2){1.0f,0.0f}, (float2){0.0f,0.0f}, I, Q, U, V);
    EXPECT_FLOAT_EQ(I, 1);
    EXPECT_FLOAT_EQ(Q, 1);
    EXPECT_FLOAT_EQ(U, 0);
    EXPECT_FLOAT_EQ(V, 0);

    // vertical polarized
    psrdada_cpp::effelsberg::edd::stokes_IQUV((float2){0.0f,0.0f}, (float2){1.0f,0.0f}, I, Q, U, V);
    EXPECT_FLOAT_EQ(I, 1);
    EXPECT_FLOAT_EQ(Q, -1);
    EXPECT_FLOAT_EQ(U, 0);
    EXPECT_FLOAT_EQ(V, 0);

    //linear +45 deg.
    psrdada_cpp::effelsberg::edd::stokes_IQUV((float2){1.0f/std::sqrt(2),0.0f}, (float2){1.0f/std::sqrt(2),0.0f}, I, Q, U, V);
    EXPECT_FLOAT_EQ(I, 1);
    EXPECT_FLOAT_EQ(Q, 0);
    EXPECT_FLOAT_EQ(U, 1);
    EXPECT_FLOAT_EQ(V, 0);

    //linear -45 deg.
    psrdada_cpp::effelsberg::edd::stokes_IQUV((float2){-1.0f/std::sqrt(2),0.0f}, (float2){1.0f/std::sqrt(2),0.0f}, I, Q, U, V);
    EXPECT_FLOAT_EQ(I, 1);
    EXPECT_FLOAT_EQ(Q, 0);
    EXPECT_FLOAT_EQ(U, -1);
    EXPECT_FLOAT_EQ(V, 0);

    //left circular
    psrdada_cpp::effelsberg::edd::stokes_IQUV((float2){.0f,1.0f/std::sqrt(2)}, (float2){1.0f/std::sqrt(2),.0f}, I, Q, U, V);
    EXPECT_FLOAT_EQ(I, 1);
    EXPECT_FLOAT_EQ(Q, 0);
    EXPECT_FLOAT_EQ(U, 0);
    EXPECT_FLOAT_EQ(V, -1);

    // right circular
    psrdada_cpp::effelsberg::edd::stokes_IQUV((float2){.0f,-1.0f/std::sqrt(2)}, (float2){1.0f/std::sqrt(2),.0f}, I, Q, U, V);
    EXPECT_FLOAT_EQ(I, 1);
    EXPECT_FLOAT_EQ(Q, 0);
    EXPECT_FLOAT_EQ(U, 0);
    EXPECT_FLOAT_EQ(V, 1);
}



void testStokesAccumulateParam(size_t nchans, size_t naccumulate){
    // Test the stokes accumulate kernel for different channel / accumualte
    // numbers

    CUDA_ERROR_CHECK(cudaDeviceSynchronize());

    thrust::device_vector<float2> P0(nchans * naccumulate);
    thrust::device_vector<float2> P1(nchans * naccumulate);
    thrust::fill(P0.begin(), P0.end(), (float2){0, 0});
    thrust::fill(P1.begin(), P1.end(), (float2){0, 0});
    thrust::device_vector<float> I(nchans);
    thrust::device_vector<float> Q(nchans);
    thrust::device_vector<float> U(nchans);
    thrust::device_vector<float> V(nchans);
    thrust::fill(I.begin(), I.end(), 0);
    thrust::fill(Q.begin(), Q.end(), 0);
    thrust::fill(U.begin(), U.end(), 0);
    thrust::fill(V.begin(), V.end(), 0);

    // This channel should be left circular polarized
    size_t idx0 = 23;
    for (size_t k = 0; k< naccumulate; k++)
    {
        size_t idx = idx0 + k * nchans;
        P0[idx] = (float2){0.0f, 1.0f/std::sqrt(2)};
        P1[idx] = (float2){1.0f/std::sqrt(2),0.0f};
    }

    psrdada_cpp::effelsberg::edd::stokes_accumulate<<<1024, 1024>>>(
          thrust::raw_pointer_cast(P0.data()),
          thrust::raw_pointer_cast(P1.data()),
          thrust::raw_pointer_cast(I.data()),
          thrust::raw_pointer_cast(Q.data()),
          thrust::raw_pointer_cast(U.data()),
          thrust::raw_pointer_cast(V.data()),
          nchans,
          naccumulate
            );

    CUDA_ERROR_CHECK(cudaDeviceSynchronize());
    thrust::pair<thrust::device_vector<float>::iterator, thrust::device_vector<float>::iterator> minmax;

    minmax = thrust::minmax_element(I.begin(), I.end());
    EXPECT_FLOAT_EQ(*minmax.first, 0);
    EXPECT_FLOAT_EQ(*minmax.second, naccumulate);

    minmax = thrust::minmax_element(Q.begin(), Q.end());
    EXPECT_FLOAT_EQ(*minmax.first, 0);
    EXPECT_FLOAT_EQ(*minmax.second, 0);

    minmax = thrust::minmax_element(U.begin(), U.end());
    EXPECT_FLOAT_EQ(*minmax.first, 0);
    EXPECT_FLOAT_EQ(*minmax.second, 0);

    minmax = thrust::minmax_element(V.begin(), V.end());
    EXPECT_FLOAT_EQ(*minmax.first, -1. * naccumulate);
    EXPECT_FLOAT_EQ(*minmax.second, 0);
};


TEST(GatedSpectrometer, stokes_accumulate)
{
    testStokesAccumulateParam(8 * 1024 * 1024 + 1, 5);
    testStokesAccumulateParam(8 * 1024 * 1024 + 1, 32);
    testStokesAccumulateParam(1024 + 1, 5);
    testStokesAccumulateParam(1024 + 1, 64);
}




TEST(GatedSpectrometer, GatingKernel)
{
  const size_t blockSize = 1024;
  const size_t nBlocks = 16 * 1024;

  thrust::device_vector<float> G0(blockSize * nBlocks);
  thrust::device_vector<float> G1(blockSize * nBlocks);
  thrust::device_vector<uint64_t> _sideChannelData(nBlocks);
  thrust::device_vector<psrdada_cpp::effelsberg::edd::uint64_cu> _nG0(nBlocks);
  thrust::device_vector<psrdada_cpp::effelsberg::edd::uint64_cu> _nG1(nBlocks);
  thrust::device_vector<float> baseLineG0(1);
  thrust::device_vector<float> baseLineG1(1);

  thrust::device_vector<float> baseLineG0_update(1);
  thrust::device_vector<float> baseLineG1_update(1);
  thrust::fill(G0.begin(), G0.end(), 42);
  thrust::fill(G1.begin(), G1.end(), 23);
  thrust::fill(_sideChannelData.begin(), _sideChannelData.end(), 0);

  // everything to G0
  {
    thrust::fill(_nG0.begin(), _nG0.end(), 0);
    thrust::fill(_nG1.begin(), _nG1.end(), 0);
    baseLineG0[0] = -3;
    baseLineG1[0] = -4;
    baseLineG0_update[0] = 0;
    baseLineG1_update[0] = 0;

    const uint64_t *sideCD =
        (uint64_t *)(thrust::raw_pointer_cast(_sideChannelData.data()));
    psrdada_cpp::effelsberg::edd::gating<<<1024 , 1024>>>(
          thrust::raw_pointer_cast(G0.data()),
          thrust::raw_pointer_cast(G1.data()), sideCD,
          G0.size(), blockSize, 0, 1,
          0,
          thrust::raw_pointer_cast(baseLineG0.data()),
          thrust::raw_pointer_cast(baseLineG1.data()),
          thrust::raw_pointer_cast(baseLineG0_update.data()),
          thrust::raw_pointer_cast(baseLineG1_update.data()),
          thrust::raw_pointer_cast(_nG0.data()),
          thrust::raw_pointer_cast(_nG1.data())
          );

    thrust::pair<thrust::device_vector<float>::iterator, thrust::device_vector<float>::iterator> minmax;
    minmax = thrust::minmax_element(G0.begin(), G0.end());
    EXPECT_EQ(*minmax.first, 42);
    EXPECT_EQ(*minmax.second, 42);

    minmax = thrust::minmax_element(G1.begin(), G1.end());
    EXPECT_EQ(*minmax.first, -4.);
    EXPECT_EQ(*minmax.second, -4.);

    EXPECT_EQ(_nG0[0], G0.size());
    EXPECT_EQ(_nG1[0], 0u);

    EXPECT_FLOAT_EQ(42.f, baseLineG0_update[0] / (_nG0[0] + 1E-121));
    EXPECT_FLOAT_EQ(0.f, baseLineG1_update[0] / (_nG1[0] + 1E-121));
  }

  // everything to G1 // with baseline -5
  {
    thrust::fill(_nG0.begin(), _nG0.end(), 0);
    thrust::fill(_nG1.begin(), _nG1.end(), 0);
    baseLineG0[0] = 5.;
    baseLineG1[0] = -2;
    baseLineG0_update[0] = 0;
    baseLineG1_update[0] = 0;

    thrust::fill(_sideChannelData.begin(), _sideChannelData.end(), 1L);
    const uint64_t *sideCD =
        (uint64_t *)(thrust::raw_pointer_cast(_sideChannelData.data()));
    psrdada_cpp::effelsberg::edd::gating<<<1024, 1024>>>(
          thrust::raw_pointer_cast(G0.data()),
          thrust::raw_pointer_cast(G1.data()), sideCD,
          G0.size(), blockSize, 0, 1,
          0,
          thrust::raw_pointer_cast(baseLineG0.data()),
          thrust::raw_pointer_cast(baseLineG1.data()),
          thrust::raw_pointer_cast(baseLineG0_update.data()),
          thrust::raw_pointer_cast(baseLineG1_update.data()),
          thrust::raw_pointer_cast(_nG0.data()),
          thrust::raw_pointer_cast(_nG1.data())
          );
    thrust::pair<thrust::device_vector<float>::iterator, thrust::device_vector<float>::iterator> minmax;
    minmax = thrust::minmax_element(G0.begin(), G0.end());
    EXPECT_EQ(*minmax.first, 5.);
    EXPECT_EQ(*minmax.second, 5.);

    minmax = thrust::minmax_element(G1.begin(), G1.end());
    EXPECT_EQ(*minmax.first, 42);
    EXPECT_EQ(*minmax.second, 42);

    EXPECT_EQ(_nG0[0], 0u);
    EXPECT_EQ(_nG1[0], G1.size());

    EXPECT_FLOAT_EQ(0.f, baseLineG0_update[0] / (_nG0[0] + 1E-121));
    EXPECT_FLOAT_EQ(42.f, baseLineG1_update[0] / (_nG1[0] + 1E-121));
  }
}

TEST(GatedSpectrometer, array_sum) {

  const size_t NBLOCKS = 16 * 32;
  const size_t NTHREADS = 1024;

  size_t inputLength = 1 << 22 + 1 ;
  size_t dataLength = inputLength;
  ////zero pad input array
  //if (inputLength % (2 * NTHREADS) != 0)
  //  dataLength = (inputLength / (2 * NTHREADS) + 1) * 2 * NTHREADS;
  thrust::device_vector<float> data(dataLength);
  thrust::fill(data.begin(), data.begin() + inputLength, 1);
  //thrust::fill(data.begin() + inputLength, data.end(), 0);
  thrust::device_vector<float> blr(NTHREADS * 2);

  thrust::fill(blr.begin(), blr.end(), 0);

  psrdada_cpp::effelsberg::edd::array_sum<<<NBLOCKS, NTHREADS, NTHREADS* sizeof(float)>>>(thrust::raw_pointer_cast(data.data()), data.size(), thrust::raw_pointer_cast(blr.data()));
  psrdada_cpp::effelsberg::edd::array_sum<<<1, NTHREADS, NTHREADS* sizeof(float)>>>(thrust::raw_pointer_cast(blr.data()), blr.size(), thrust::raw_pointer_cast(blr.data()));

  EXPECT_EQ(size_t(blr[0]), inputLength);
}


class GatedTestSinkSinglePol{
    private:
        size_t fft_length;
        size_t call_count;
        size_t nHeaps;
        size_t naccumulate;
    public:
        GatedTestSinkSinglePol(size_t fft_length, size_t nHeaps, size_t naccumulate): fft_length(fft_length), nHeaps(nHeaps), naccumulate(naccumulate), call_count(0) {};
    // Test the correctness of output of the processing test
        void init(psrdada_cpp::RawBytes&){
        };
        bool operator()(psrdada_cpp::RawBytes& buf)
        {
            const size_t number_of_sc_items = 2;
            const size_t nchans = fft_length / 2 + 1;
            EXPECT_EQ(buf.used_bytes(), (32 / 8 * nchans + number_of_sc_items* 64 / 8) * 2);

            float *G0 = reinterpret_cast<float*>(buf.ptr());
            float *G1 = reinterpret_cast<float*>(buf.ptr() + nchans * 32 /8);

            // Expected half number of samples per gate
            uint64_t *S = reinterpret_cast<uint64_t*>(buf.ptr() + 2 *nchans * 32 /8);
            EXPECT_EQ(S[0], nHeaps * 4096 / 2);
            EXPECT_EQ(S[2], nHeaps * 4096 / 2);

            // Correct number of overflowed samples
            EXPECT_EQ(S[1], 27uL);
            EXPECT_EQ(S[3], 23uL);  // First heap has 23 and bit set, thus G1

            call_count ++;
            return false;
        };

};

TEST(GatedSpectrometer, processingSinglePol)
{
    const size_t nbits = 8;
    const size_t nHeaps = 1024;
    const size_t fft_length = 1024 * 64;
    const size_t naccumulate = 4096 * nHeaps / fft_length;

    const size_t heapSize = 4096 * nbits / 8;
    const size_t inputBufferSize = nHeaps * (heapSize + 64 / 8);

    psrdada_cpp::DadaDB idbuffer(5, inputBufferSize, 1, 4096);
    idbuffer.create();

    psrdada_cpp::effelsberg::edd::DadaBufferLayout bufferLayout(idbuffer.key(), heapSize, 1);

    GatedTestSinkSinglePol sink(fft_length, nHeaps, naccumulate);
    psrdada_cpp::effelsberg::edd::GatedSpectrometer<
        GatedTestSinkSinglePol,
        psrdada_cpp::effelsberg::edd::SinglePolarizationInput,
        psrdada_cpp::effelsberg::edd::GatedPowerSpectrumOutput>
            spectrometer(bufferLayout,
                    0, 0,
                    fft_length,
                    naccumulate, nbits,
                    sink);

    char  *raw_buffer = new char[inputBufferSize];

    psrdada_cpp::RawBytes buff(raw_buffer, inputBufferSize, inputBufferSize);
    EXPECT_NO_THROW(spectrometer.init(buff));

    //// fill sci data
    uint64_t* sc_items = reinterpret_cast<uint64_t*>(raw_buffer + nHeaps * 4096);
    for (size_t i = 0; i < nHeaps; i+=2)
    {
        sc_items[i] = 0;
        sc_items[i+1] = 0;
        SET_BIT(sc_items[i], 0);
    }
    sc_items[0] |= (23UL) << 32;
    sc_items[1] |= (27UL) << 32;

    for (int i = 0; i < 12; i++)
    {
        EXPECT_NO_THROW(spectrometer(buff));
    }

    delete [] raw_buffer;
}


class GatedTestSinkFullStokes{
    private:
        size_t fft_length;
        size_t call_count;
        size_t nHeaps;
        size_t naccumulate;
    public:
        GatedTestSinkFullStokes(size_t fft_length, size_t nHeaps, size_t naccumulate): fft_length(fft_length), nHeaps(nHeaps), naccumulate(naccumulate), call_count(0) {};
    // Test the correctness of output of the processing test
        void init(psrdada_cpp::RawBytes&){
        };
        bool operator()(psrdada_cpp::RawBytes& buf)
        {
            const size_t number_of_sc_items = 2;
            const size_t nchans = fft_length / 2 + 1;
            EXPECT_EQ(buf.used_bytes(), (32 / 8 * nchans + number_of_sc_items* 64 / 8) * 2 * 4); // 4 Spectra (IQUV), 2 Gates

            // Expected half number of samples per gate
            for(int i =0; i <4; i++)
            {
                uint64_t *S = reinterpret_cast<uint64_t*>(buf.ptr() + 8 *nchans * 32 /8 + i * 4 *sizeof(size_t));
                EXPECT_EQ(S[0], nHeaps * 4096) << "G0, S" << i; // expect samples from two polarizations
                EXPECT_EQ(S[2], nHeaps * 4096) << "G1, S" << i;

                // Correct number of overflowed samples
                EXPECT_EQ(S[1], uint64_t(27 + 7)) << "G0, S" << i;;
                EXPECT_EQ(S[3], uint64_t(23 + 3)) << "G1, S" << i;;  // First heap has 23+3 and bit set, thus G1
            }

            call_count ++;
            return false;
        };

};




TEST(GatedSpectrometer, processingFullStokes)
{
    const size_t nbits = 8;
    const size_t nHeaps = 1024;
    const size_t fft_length = 1024 * 64;
    const size_t naccumulate = 4096 * nHeaps / fft_length;

    const size_t heapSize = 4096 * nbits / 8;
    const size_t inputBufferSize = 2 * nHeaps * (heapSize + 64 / 8) ;

    psrdada_cpp::DadaDB idbuffer(5, inputBufferSize, 1, 4096);
    idbuffer.create();

    psrdada_cpp::effelsberg::edd::DadaBufferLayout bufferLayout(idbuffer.key(), heapSize, 1);

    GatedTestSinkFullStokes sink(fft_length, nHeaps, naccumulate);
    psrdada_cpp::effelsberg::edd::GatedSpectrometer<
        GatedTestSinkFullStokes,
        psrdada_cpp::effelsberg::edd::DualPolarizationInput,
        psrdada_cpp::effelsberg::edd::GatedFullStokesOutput>
            spectrometer(bufferLayout,
                    0, 0,
                    fft_length,
                    naccumulate, nbits,
                    sink);

    char  *raw_buffer = new char[inputBufferSize];

    psrdada_cpp::RawBytes buff(raw_buffer, inputBufferSize, inputBufferSize);
    EXPECT_NO_THROW(spectrometer.init(buff));

    //// fill sci data
    uint64_t* sc_items = reinterpret_cast<uint64_t*>(raw_buffer + 2*nHeaps * 4096);
    for (size_t i = 0; i < 2 * nHeaps; i+=4)
    {
        sc_items[i] = 0uL;
        sc_items[i+1] = 0uL;
        sc_items[i+2] = 0uL;
        sc_items[i+3] = 0uL;
        SET_BIT(sc_items[i], 0);
        SET_BIT(sc_items[i + 1], 0);
    }
    sc_items[0] |= (23UL) << 32;
    sc_items[1] |= (3UL) << 32;

    sc_items[2] |= (27UL) << 32;
    sc_items[3] |= (7UL) << 32;


    for (int i = 0; i < 12; i++)
    {
        EXPECT_NO_THROW(spectrometer(buff));
    }

    delete [] raw_buffer;
}









struct gated_params
{
    std::size_t fft_length;
    std::size_t naccumulate;
    std::size_t nbits;
    std::size_t nHeaps;
    std::string msg;
};


class ExecutionTests: public testing::TestWithParam<gated_params> {
// Test that the spectrometers execute without error in certain parameter
// settings

};


TEST_P(ExecutionTests, SinglePolOutput)
{
    gated_params params = GetParam();

    const size_t heapSize = 4096 * params.nbits / 8;
    const size_t inputBufferSize = params.nHeaps * (heapSize + 64 / 8);

    psrdada_cpp::DadaDB buffer(5, inputBufferSize, 1, 4096);
    buffer.create();

    psrdada_cpp::effelsberg::edd::DadaBufferLayout bufferLayout(buffer.key(), heapSize, 1);
    // Test buffer consistency with input parameters
    EXPECT_EQ(bufferLayout.getNHeaps(), params.nHeaps);

    psrdada_cpp::NullSink sink;

    psrdada_cpp::effelsberg::edd::GatedSpectrometer<
        psrdada_cpp::NullSink,
        psrdada_cpp::effelsberg::edd::SinglePolarizationInput,
        psrdada_cpp::effelsberg::edd::GatedPowerSpectrumOutput>
            spectrometer(bufferLayout,
                    0, 0,
                     params.fft_length,
                    params.naccumulate, params.nbits,
                    sink);

    char  *raw_buffer = new char[inputBufferSize];
    memset(raw_buffer, 0, inputBufferSize);

    psrdada_cpp::RawBytes buff(raw_buffer, inputBufferSize, inputBufferSize);

    EXPECT_NO_THROW(spectrometer.init(buff)) << params.msg;
    for (int i = 0; i < 5; i++)
    {
        EXPECT_NO_THROW(spectrometer(buff)) << params.msg;
    }

    delete [] raw_buffer;
}


TEST_P(ExecutionTests, FullStokesOutput)
{
    gated_params params = GetParam();

    const size_t heapSize = 4096 * params.nbits / 8;
    const size_t inputBufferSize = params.nHeaps * (heapSize + 64 / 8);

    psrdada_cpp::DadaDB buffer(5, inputBufferSize, 1, 4096);
    buffer.create();

    psrdada_cpp::effelsberg::edd::DadaBufferLayout bufferLayout(buffer.key(), heapSize, 1);
    // Test buffer consistency with input parameters
    EXPECT_EQ(bufferLayout.getNHeaps(), params.nHeaps);

    psrdada_cpp::NullSink sink;

    psrdada_cpp::effelsberg::edd::GatedSpectrometer<
        psrdada_cpp::NullSink,
        psrdada_cpp::effelsberg::edd::DualPolarizationInput,
        psrdada_cpp::effelsberg::edd::GatedFullStokesOutput>
            spectrometer(bufferLayout,
                    0, 0,
                     params.fft_length,
                    params.naccumulate, params.nbits,
                    sink);

    char  *raw_buffer = new char[inputBufferSize];
    memset(raw_buffer, 0, inputBufferSize);

    psrdada_cpp::RawBytes buff(raw_buffer, inputBufferSize, inputBufferSize);

    EXPECT_NO_THROW(spectrometer.init(buff)) << params.msg;
    for (int i = 0; i < 5; i++)
    {
        EXPECT_NO_THROW(spectrometer(buff)) << params.msg;
    }

    delete [] raw_buffer;
}


INSTANTIATE_TEST_CASE_P (GatedSpectrometer,
        ExecutionTests,
        testing::Values(

          // fft_length; naccumulate; nbits; nHeaps; msg;
          gated_params({2*1024, 2,  8, 1024, "1k Channel:  8 Bit"}),
          gated_params({2*1024, 2, 10, 1024, "1k Channel: 10 Bit"}),
          gated_params({2*1024, 2, 12, 1024, "1k Channel: 12 Bit"}),
          gated_params({2*1024 * 1024, 2,  8, 8*1024, "1M Channel, buffer 8x larger than spectrum; 8 bit"}),
          gated_params({2*1024 * 1024, 2, 10, 8*1024, "1M Channel, buffer 8x larger than spectrum; 10-bit"}),
          gated_params({2*1024 * 1024, 2, 12, 8*1024, "1M Channel, buffer 8x larger than spectrum; 12-bit"}),
          gated_params({2*1024 * 1024, 2048, 8, 1024, "1M Channel, integration larger than buffer"}),
          gated_params({2*1024, 4096, 8, 1024, "1k Channel, integration larger than buffer"})
            )
        );