#include <gtest/gtest.h>

#include <NBData/FFT.hpp>

using namespace nb;

template<typename A, typename B>
void COMPLEX_EQ(const std::complex<A>& a, const std::complex<B>& b) {
    ASSERT_EQ(a.real(), b.real());
    ASSERT_EQ(a.imag(), b.imag());
}

template<typename A, typename B>
void COMPLEX_FLOAT_EQ(
    const std::complex<A>& a,
    const std::complex<B>& b,
    const long double epsilon=0
) {
    if (epsilon) {
        ASSERT_NEAR(a.real(), b.real(), epsilon);
        ASSERT_NEAR(a.imag(), b.imag(), epsilon);
    } else {
        ASSERT_FLOAT_EQ(a.real(), b.real());
        ASSERT_FLOAT_EQ(a.imag(), b.imag());
    }
}

template<typename A, typename B>
void COMPLEX_EQ_VEC(
    const std::vector<std::complex<A>>& a,
    const std::vector<std::complex<A>>& b
) {
    if (a.size() != b.size()) { FAIL(); }
    for (int i = 0; i < a.size(); ++i) {
        COMPLEX_EQ(a[i], b[i]);
    }
}

template<typename A, typename B>
void COMPLEX_FLOAT_EQ_VEC(
    const std::vector<std::complex<A>>& a,
    const std::vector<std::complex<B>>& b,
    const long double epsilon
) {
    if (a.size() != b.size()) { FAIL(); }
    for (int i = 0; i < a.size(); ++i) {
        COMPLEX_FLOAT_EQ(a[i], b[i], epsilon);
    }
}

TEST(FFTTest, TestFFT) {
    std::vector<std::vector<std::complex<double>>> input = {
        {1.0, 0.5, 0},
        {1,1,1,1,1,1,1,1},

    };
    std::vector<std::vector<std::complex<double>>> expected = {
        {1.5,1.f-0.5if,0.5,1.f+0.5if},
        {8, 0, 0, 0, 0, 0, 0, 0}
    };
    for (int i = 0; i < input.size(); ++i) {
        COMPLEX_FLOAT_EQ_VEC(
            nb::fft2(input[i]),
            expected[i],
             1e-10
        );
    }
}