#include "Base/Math/FourierTransform.h"

#include "Device/Data/Datafield.h"
#include "Tests/GTestWrapper/google_test.h"

// Testing implementation of 1D FT with with low freuency centering
TEST(FourierTransform, fft1D)
{
    std::vector<double> signal_odd({0, 0, 1, 0, 0, 1, 1});     // odd # input size
    std::vector<double> signal_even({0, 0, 1, 0, 0, 1, 1, 2}); // even # input size

    FourierTransform ft;

    // shift, odd
    std::vector<double> signal_odd_shifted = ft.fftshift(signal_odd);
    signal_odd_shifted = ft.ifftshift(signal_odd_shifted);
    for (size_t i = 0; i < signal_odd.size(); ++i)
        EXPECT_NEAR(signal_odd_shifted[i], signal_odd[i], 1e-13);

    // shift, even
    std::vector<double> signal_even_shifted = ft.fftshift(signal_even);
    signal_even_shifted = ft.ifftshift(signal_even_shifted);
    for (size_t i = 0; i < signal_even.size(); ++i)
        EXPECT_NEAR(signal_even_shifted[i], signal_even[i], 1e-13);

    // odd, amplitude, shifted
    std::vector<double> result_amp_odd = ft.ramplitude(signal_odd);
    result_amp_odd = ft.fftshift(result_amp_odd);
    std::vector<double> expected_result_amp_odd(
        {0.55495813, 2.2469796, 0.80193774, 3., 0.80193774, 2.2469796, 0.55495813});

    EXPECT_EQ(signal_odd.size(), result_amp_odd.size());
    EXPECT_EQ(expected_result_amp_odd.size(), result_amp_odd.size());
    for (size_t i = 0; i < signal_odd.size(); ++i)
        EXPECT_NEAR(result_amp_odd[i], expected_result_amp_odd[i], 1e-6);

    // odd
    std::vector<complex_t> result_odd = ft.rfft(signal_odd);
    std::vector<complex_t> expected_result_odd({3. + 0. * I, 0.17844793 + 0.78183148 * I,
                                                -2.02445867 + 0.97492791 * I,
                                                0.34601074 + 0.43388374 * I});

    EXPECT_EQ(signal_odd.size() / 2 + 1, result_odd.size());
    EXPECT_EQ(expected_result_odd.size(), result_odd.size());
    for (size_t i = 0; i < result_odd.size(); ++i)
        EXPECT_NEAR(abs(result_odd[i] - expected_result_odd[i]), 0, abs(result_odd[i]) * 1e-8);

    // odd, inverse
    std::vector<double> result_inv_odd = ft.irfft(result_odd, signal_odd.size());

    EXPECT_EQ(signal_odd.size(), result_inv_odd.size());
    for (size_t i = 0; i < signal_odd.size(); ++i)
        EXPECT_NEAR(result_inv_odd[i], signal_odd[i], 1e-8);

    // even, amplitude, shifted
    std::vector<double> result_amp_even = ft.ramplitude(signal_even);
    result_amp_even = ft.fftshift(result_amp_even);
    std::vector<double> expected_result_amp_even(
        {1., 2.23606798, 2.23606798, 2.23606798, 5., 2.23606798, 2.23606798, 2.23606798});

    EXPECT_EQ(signal_even.size(), result_amp_even.size());
    EXPECT_EQ(expected_result_amp_even.size(), result_amp_even.size());
    for (size_t i = 0; i < signal_even.size(); ++i)
        EXPECT_NEAR(result_amp_even[i], expected_result_amp_even[i], 1e-6);

    // even
    std::vector<complex_t> result_even = ft.rfft(signal_even);
    std::vector<complex_t> expected_result_even({5. + 0. * I, 0.70710678 + 2.12132034 * I,
                                                 -2. + 1. * I, -0.70710678 + 2.12132034 * I,
                                                 -1. + 0. * I});

    EXPECT_EQ(signal_even.size() / 2 + 1, result_even.size());
    EXPECT_EQ(expected_result_even.size(), result_even.size());
    for (size_t i = 0; i < result_even.size(); ++i)
        EXPECT_NEAR(abs(result_even[i] - expected_result_even[i]), 0, abs(result_even[i]) * 1e-8);

    // even, inverse
    std::vector<double> result_inv_even = ft.irfft(result_even, signal_even.size());

    EXPECT_EQ(signal_even.size(), result_inv_even.size());
    for (size_t i = 0; i < signal_even.size(); ++i)
        EXPECT_NEAR(result_inv_even[i], signal_even[i], 1e-8);
}

// Testing implementation of 2D FT with low freuency centering for the following:

// 3x5 input of all zeros
TEST(FourierTransform, fft2DTest1)
{
    double2d_t signal({{0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}});
    double2d_t expected_result({{0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}});

    FourierTransform ft;
    double2d_t result = ft.ramplitude(signal);
    result = ft.fftshift(result);

    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(result[i][j], expected_result[i][j], 1e-6);
}

// 4x5 input of all zeros except for 1 element
TEST(FourierTransform, fft2DTest2)
{
    double2d_t signal({{0, 0, 0, 0, 0}, {0, 0, 2, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}});
    double2d_t expected_result(
        {{2, 2, 2, 2, 2}, {2, 2, 2, 2, 2}, {2, 2, 2, 2, 2}, {2, 2, 2, 2, 2}});

    FourierTransform ft;
    double2d_t result = ft.ramplitude(signal);
    result = ft.fftshift(result);

    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(result[i][j], expected_result[i][j], 1e-6);
}

// 6x6 input of all ones except for 1 element
TEST(FourierTransform, fft2DTest3)
{
    double2d_t signal({{1, 1, 1, 1, 1, 1},
                       {1, 1, 1, 1, 1, 1},
                       {1, 1, 1, 1, 1, 1},
                       {1, 1, 1, 1, 1, 1},
                       {1, 1, 1, 1, 1, 32},
                       {1, 1, 1, 1, 1, 1}});
    double2d_t expected_result({{31, 31, 31, 31, 31, 31},
                                {31, 31, 31, 31, 31, 31},
                                {31, 31, 31, 31, 31, 31},
                                {31, 31, 31, 67, 31, 31},
                                {31, 31, 31, 31, 31, 31},
                                {31, 31, 31, 31, 31, 31}});

    FourierTransform ft;
    double2d_t result = ft.ramplitude(signal);
    result = ft.fftshift(result);

    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(result[i][j], expected_result[i][j], 1e-6);
}

// 3x5 input with 1 row of all zeros
TEST(FourierTransform, fft2DTest4)
{
    double2d_t signal({{1, 88, 0, 1, 0}, {0, 1, 1, 1, 0}, {0, 0, 0, 0, 0}});

    FourierTransform ft;

    // shift
    double2d_t signal_shifted = ft.fftshift(signal);
    signal_shifted = ft.ifftshift(signal_shifted);
    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(signal_shifted[i][j], signal[i][j], 1e-13);

    // amplitude, shifted
    double2d_t expected_result_amp(
        {{87.56947917, 85.92017286, 88.53812738, 88.59651195, 86.95382846},
         {88.02055461, 88.00785173, 93., 88.00785173, 88.02055461},
         {86.95382846, 88.59651195, 88.53812738, 85.92017286, 87.56947917}});

    double2d_t result_amp = ft.ramplitude(signal);
    result_amp = ft.fftshift(result_amp);

    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(result_amp[i][j], expected_result_amp[i][j], 1e-6);

    // complex spectrum
    std::vector<std::vector<complex_t>> expected_result(
        {{93. + 0. * I, 26.07546152 - 84.0562447 * I, -70.07546152 - 53.26394397 * I},
         {88.5 - 2.59807621 * I, 27.2153479 - 81.49601795 * I, -70.29802397 - 52.21686996 * I},
         {88.5 + 2.59807621 * I, 28.86262611 - 83.76330189 * I, -69.27995005 - 52.54766223 * I}});
    std::vector<std::vector<complex_t>> result = ft.rfft(signal);

    EXPECT_EQ(signal.size(), result.size());
    EXPECT_EQ(signal[0].size() / 2 + 1, result[0].size());
    for (size_t i = 0; i < expected_result.size(); ++i)
        for (size_t j = 0; j < expected_result[0].size(); ++j)
            EXPECT_NEAR(abs(result[i][j] - expected_result[i][j]), 0, abs(result[i][j]) * 1e-8);

    // inverse
    double2d_t result_inv = ft.irfft(result, signal[0].size());

    EXPECT_EQ(signal[0].size(), result_inv[0].size());
    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(result_inv[i][j], signal[i][j], 1e-13);
}

// 4x4 input
TEST(FourierTransform, fft2DTest5)
{
    double2d_t signal({{1, 0, 0, 5}, {1, 0, 0, 0}, {0, 1, 1, 1}, {0, 1, 1, 1}});

    FourierTransform ft;

    // shift
    double2d_t signal_shifted = ft.fftshift(signal);
    signal_shifted = ft.ifftshift(signal_shifted);
    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(signal_shifted[i][j], signal[i][j], 1e-13);

    // amplitude, shifted
    double2d_t expected_result_amp({{5., 5., 5., 5.},
                                    {3.60555128, 3.60555128, 3.60555128, 7.28010989},
                                    {5., 5., 13., 5.},
                                    {3.60555128, 7.28010989, 3.60555128, 3.60555128}});

    double2d_t result_amp = ft.ramplitude(signal);
    result_amp = ft.fftshift(result_amp);

    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(result_amp[i][j], expected_result_amp[i][j], 1e-6);

    // complex spectrum
    std::vector<std::vector<complex_t>> expected_result({{13. + 0. * I, 0. + 5. * I, -5. + 0. * I},
                                                         {3. + 2. * I, 2. + 3. * I, -3. - 2. * I},
                                                         {5. + 0. * I, 0. + 5. * I, -5. + 0. * I},
                                                         {3. - 2. * I, 2. + 7. * I, -3. + 2. * I}});
    std::vector<std::vector<complex_t>> result = ft.rfft(signal);

    EXPECT_EQ(signal.size(), result.size());
    EXPECT_EQ(signal[0].size() / 2 + 1, result[0].size());
    for (size_t i = 0; i < expected_result.size(); ++i)
        for (size_t j = 0; j < expected_result[0].size(); ++j)
            EXPECT_NEAR(abs(result[i][j] - expected_result[i][j]), 0, abs(result[i][j]) * 1e-8);

    // inverse
    double2d_t result_inv = ft.irfft(result, signal[0].size());

    EXPECT_EQ(signal[0].size(), result_inv[0].size());
    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(result_inv[i][j], signal[i][j], 1e-13);
}

// 7x7 input
TEST(FourierTransform, fft2DTest6)
{
    double2d_t signal{{1., 0., 0., 0., 0., 0., 0.}, {1., 0., 0., 0., 0., 0., 0.},
                      {1., 0., 5., 0., 0., 0., 0.}, {1., 0., 0., 0., 0., 0., 0.},
                      {1., 0., 0., 0., 0., 5., 0.}, {1., 0., 0., 0., 0., 0., 0.},
                      {1., 0., 0., 0., 0., 0., 0.}};

    double2d_t expected_result(
        {{9.00968868, 6.23489802, 6.23489802, 9.00968868, 2.22520934, 10., 2.22520934},
         {9.00968868, 2.22520934, 10., 2.22520934, 9.00968868, 6.23489802, 6.23489802},
         {2.22520934, 9.00968868, 6.23489802, 6.23489802, 9.00968868, 2.22520934, 10.},
         {13.23489802, 2.00968868, 4.77479066, 17., 4.77479066, 2.00968868, 13.23489802},
         {10., 2.22520934, 9.00968868, 6.23489802, 6.23489802, 9.00968868, 2.22520934},
         {6.23489802, 6.23489802, 9.00968868, 2.22520934, 10., 2.22520934, 9.00968868},
         {2.22520934, 10., 2.22520934, 9.00968868, 6.23489802, 6.23489802, 9.00968868}});

    FourierTransform ft;
    double2d_t result = ft.ramplitude(signal);
    result = ft.fftshift(result);

    for (size_t i = 0; i < signal.size(); ++i)
        for (size_t j = 0; j < signal[0].size(); ++j)
            EXPECT_NEAR(result[i][j], expected_result[i][j], 1e-8);
}
