// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// fft_simd.hpp
//
// SIMD-optimized Radix-2 FFT (float / double)
//
// Optimization points:
//   1. precomputed twiddle table (cos/sin lookup, avoids recomputing every stage)
//   2. AVX2 butterfly: float 4 complex (8-wide), double 2 complex (4-wide)
//   3. SSE2 butterfly: float 2 complex (4-wide), double 1 complex (2-wide)
//   4. compile-time SIMD level selection (#if dispatch)
//
// Usage:
//   std::vector<Complex<float>> data(1024);
//   sangi::simd_fft::fft(data.data(), 1024);   // forward transform
//   sangi::simd_fft::ifft(data.data(), 1024);  // inverse transform
//
//   std::vector<Complex<double>> data_d(1024);
//   sangi::simd_fft::fft(data_d.data(), 1024);

#ifndef SANGI_FFT_SIMD_HPP
#define SANGI_FFT_SIMD_HPP

#include <math/core/Complex.hpp>
#include <cmath>
#include <cstdint>
#include <numbers>
#include <vector>
#include <type_traits>

// SIMD detection (same pattern as simd_backend.hpp)
#if defined(__AVX2__) || (defined(_MSC_VER) && defined(__AVX2__))
#define SANGI_FFT_HAS_AVX2 1
#endif

#if defined(__AVX__) || (defined(_MSC_VER) && defined(__AVX__))
#define SANGI_FFT_HAS_AVX 1
#endif

#if defined(__SSE2__) || defined(_M_X64) || defined(_M_AMD64) || (defined(_M_IX86_FP) && _M_IX86_FP >= 2)
#define SANGI_FFT_HAS_SSE2 1
#endif

#if defined(SANGI_FFT_HAS_AVX2) || defined(SANGI_FFT_HAS_AVX) || defined(SANGI_FFT_HAS_SSE2)
#include <immintrin.h>
#endif

namespace sangi {
namespace simd_fft {

// =============================================================================
// Bit-reversal reordering
// =============================================================================

inline void bit_reversal(Complex<float>* data, int n)
{
    for (int i = 0, j = 0; i < n; ++i) {
        if (i < j) std::swap(data[i], data[j]);
        int mask = n;
        do { mask >>= 1; j ^= mask; } while (mask && (j & mask) == 0);
    }
}

inline void bit_reversal(Complex<double>* data, int n)
{
    for (int i = 0, j = 0; i < n; ++i) {
        if (i < j) std::swap(data[i], data[j]);
        int mask = n;
        do { mask >>= 1; j ^= mask; } while (mask && (j & mask) == 0);
    }
}

// =============================================================================
// Twiddle table (precomputed cos/sin)
// =============================================================================

template<typename Real>
struct TwiddleTable {
    // twiddles[s] = twiddle factors of stage s (len = 2^(s+1))
    // twiddles[s][j] = exp(-2πi·j / 2^(s+1)),  j = 0..2^s - 1
    std::vector<std::vector<Complex<Real>>> twiddles;

    void build(int n) {
        int stages = 0;
        for (int tmp = n; tmp > 1; tmp >>= 1) ++stages;
        twiddles.resize(stages);

        for (int s = 0; s < stages; ++s) {
            int len = 2 << s;        // 2, 4, 8, ...
            int half = len / 2;      // 1, 2, 4, ...
            twiddles[s].resize(half);
            Real angle = Real(-2) * std::numbers::pi_v<Real> / Real(len);
            for (int j = 0; j < half; ++j) {
                Real a = angle * Real(j);
                twiddles[s][j] = Complex<Real>(std::cos(a), std::sin(a));
            }
        }
    }

    // for inverse transform: returns the conjugate
    Complex<Real> get_inv(int stage, int j) const {
        return conj(twiddles[stage][j]);
    }
};

// =============================================================================
// Scalar butterfly (fallback)
// =============================================================================

template<typename Real>
inline void butterfly_stage_scalar(
    Complex<Real>* data, int n, int half_len,
    const Complex<Real>* tw)
{
    for (int i = 0; i < n; i += half_len * 2) {
        for (int j = 0; j < half_len; ++j) {
            auto u = data[i + j];
            auto v = data[i + j + half_len] * tw[j];
            data[i + j]            = u + v;
            data[i + j + half_len] = u - v;
        }
    }
}

// =============================================================================
// SSE2 butterfly — double (1 complex = 2 double = 1 __m128d)
// =============================================================================

#if defined(SANGI_FFT_HAS_SSE2)

inline void butterfly_stage_sse2_double(
    Complex<double>* data, int n, int half_len,
    const Complex<double>* tw)
{
    // half_len < 1 cannot happen, but for safety
    for (int i = 0; i < n; i += half_len * 2) {
        int j = 0;
        for (; j < half_len; ++j) {
            // load data[k] and data[k+half]
            int k = i + j;
            __m128d u = _mm_loadu_pd(reinterpret_cast<double*>(&data[k]));
            __m128d d = _mm_loadu_pd(reinterpret_cast<double*>(&data[k + half_len]));
            __m128d w = _mm_loadu_pd(reinterpret_cast<const double*>(&tw[j]));

            // complex multiplication v = d * w
            // d = [dre, dim], w = [wre, wim]
            __m128d d_re = _mm_shuffle_pd(d, d, 0); // [dre, dre]
            __m128d d_im = _mm_shuffle_pd(d, d, 3); // [dim, dim]
            __m128d w_flip = _mm_shuffle_pd(w, w, 1); // [wim, wre]

            __m128d p1 = _mm_mul_pd(d_re, w);       // [dre*wre, dre*wim]
            __m128d p2 = _mm_mul_pd(d_im, w_flip);  // [dim*wim, dim*wre]

            // addsub: [p1[0]-p2[0], p1[1]+p2[1]] = [dre*wre-dim*wim, dre*wim+dim*wre]
            // SSE3 has _mm_addsub_pd, but for SSE2 we do manually:
            __m128d sign = _mm_set_pd(1.0, -1.0);
            __m128d v = _mm_add_pd(p1, _mm_mul_pd(p2, sign));

            _mm_storeu_pd(reinterpret_cast<double*>(&data[k]),            _mm_add_pd(u, v));
            _mm_storeu_pd(reinterpret_cast<double*>(&data[k + half_len]), _mm_sub_pd(u, v));
        }
    }
}

// SSE2 butterfly — float (2 complex = 4 float = 1 __m128)
inline void butterfly_stage_sse2_float(
    Complex<float>* data, int n, int half_len,
    const Complex<float>* tw)
{
    const __m128 sign_mask = _mm_set_ps(1.0f, -1.0f, 1.0f, -1.0f);

    for (int i = 0; i < n; i += half_len * 2) {
        int j = 0;
        // process 2 complex at a time
        for (; j + 1 < half_len; j += 2) {
            int k = i + j;
            // load 2 complex numbers (4 float)
            __m128 u = _mm_loadu_ps(reinterpret_cast<float*>(&data[k]));
            __m128 d = _mm_loadu_ps(reinterpret_cast<float*>(&data[k + half_len]));
            __m128 w = _mm_loadu_ps(reinterpret_cast<const float*>(&tw[j]));

            // d = [dre0, dim0, dre1, dim1], w = [wre0, wim0, wre1, wim1]
            __m128 d_re = _mm_shuffle_ps(d, d, 0xA0); // [dre0,dre0,dre1,dre1]
            __m128 d_im = _mm_shuffle_ps(d, d, 0xF5); // [dim0,dim0,dim1,dim1]
            __m128 w_flip = _mm_shuffle_ps(w, w, 0xB1); // [wim0,wre0,wim1,wre1]

            __m128 p1 = _mm_mul_ps(d_re, w);
            __m128 p2 = _mm_mul_ps(d_im, w_flip);
            __m128 v  = _mm_add_ps(p1, _mm_mul_ps(p2, sign_mask));

            _mm_storeu_ps(reinterpret_cast<float*>(&data[k]),            _mm_add_ps(u, v));
            _mm_storeu_ps(reinterpret_cast<float*>(&data[k + half_len]), _mm_sub_ps(u, v));
        }
        // remainder (when half_len is odd)
        for (; j < half_len; ++j) {
            int k = i + j;
            auto uu = data[k];
            auto vv = data[k + half_len] * tw[j];
            data[k]            = uu + vv;
            data[k + half_len] = uu - vv;
        }
    }
}

#endif // SANGI_FFT_HAS_SSE2

// =============================================================================
// AVX2 butterfly — double (2 complex = 4 double = 1 __m256d)
// =============================================================================

#if defined(SANGI_FFT_HAS_AVX2)

inline void butterfly_stage_avx2_double(
    Complex<double>* data, int n, int half_len,
    const Complex<double>* tw)
{
    const __m256d sign_mask = _mm256_set_pd(1.0, -1.0, 1.0, -1.0);

    for (int i = 0; i < n; i += half_len * 2) {
        int j = 0;
        // 2 complex at a time
        for (; j + 1 < half_len; j += 2) {
            int k = i + j;
            __m256d u = _mm256_loadu_pd(reinterpret_cast<double*>(&data[k]));
            __m256d d = _mm256_loadu_pd(reinterpret_cast<double*>(&data[k + half_len]));
            __m256d w = _mm256_loadu_pd(reinterpret_cast<const double*>(&tw[j]));

            // complex multiplication: v = d * w
            __m256d d_re = _mm256_shuffle_pd(d, d, 0x0); // [dre0,dre0,dre1,dre1]
            __m256d d_im = _mm256_shuffle_pd(d, d, 0xF); // [dim0,dim0,dim1,dim1]
            __m256d w_flip = _mm256_shuffle_pd(w, w, 0x5); // [wim0,wre0,wim1,wre1]

            __m256d p1 = _mm256_mul_pd(d_re, w);
            __m256d p2 = _mm256_mul_pd(d_im, w_flip);
            __m256d v  = _mm256_add_pd(p1, _mm256_mul_pd(p2, sign_mask));

            _mm256_storeu_pd(reinterpret_cast<double*>(&data[k]),            _mm256_add_pd(u, v));
            _mm256_storeu_pd(reinterpret_cast<double*>(&data[k + half_len]), _mm256_sub_pd(u, v));
        }
        // remainder
        for (; j < half_len; ++j) {
            int k = i + j;
            auto uu = data[k];
            auto vv = data[k + half_len] * tw[j];
            data[k]            = uu + vv;
            data[k + half_len] = uu - vv;
        }
    }
}

// AVX2 butterfly — float (4 complex = 8 float = 1 __m256)
inline void butterfly_stage_avx2_float(
    Complex<float>* data, int n, int half_len,
    const Complex<float>* tw)
{
    const __m256 sign_mask = _mm256_set_ps(
        1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -1.0f);

    for (int i = 0; i < n; i += half_len * 2) {
        int j = 0;
        // 4 complex at a time
        for (; j + 3 < half_len; j += 4) {
            int k = i + j;
            __m256 u = _mm256_loadu_ps(reinterpret_cast<float*>(&data[k]));
            __m256 d = _mm256_loadu_ps(reinterpret_cast<float*>(&data[k + half_len]));
            __m256 w = _mm256_loadu_ps(reinterpret_cast<const float*>(&tw[j]));

            // d = [re0,im0, re1,im1, re2,im2, re3,im3]
            __m256 d_re = _mm256_shuffle_ps(d, d, 0xA0);   // [re,re, re,re, ...]
            __m256 d_im = _mm256_shuffle_ps(d, d, 0xF5);   // [im,im, im,im, ...]
            __m256 w_flip = _mm256_shuffle_ps(w, w, 0xB1); // [im,re, im,re, ...]

            __m256 p1 = _mm256_mul_ps(d_re, w);
            __m256 p2 = _mm256_mul_ps(d_im, w_flip);
            __m256 v  = _mm256_add_ps(p1, _mm256_mul_ps(p2, sign_mask));

            _mm256_storeu_ps(reinterpret_cast<float*>(&data[k]),            _mm256_add_ps(u, v));
            _mm256_storeu_ps(reinterpret_cast<float*>(&data[k + half_len]), _mm256_sub_ps(u, v));
        }
        // 2 complex with SSE2
#if defined(SANGI_FFT_HAS_SSE2)
        {
            const __m128 sign4 = _mm_set_ps(1.0f, -1.0f, 1.0f, -1.0f);
            for (; j + 1 < half_len; j += 2) {
                int k = i + j;
                __m128 u = _mm_loadu_ps(reinterpret_cast<float*>(&data[k]));
                __m128 d = _mm_loadu_ps(reinterpret_cast<float*>(&data[k + half_len]));
                __m128 w = _mm_loadu_ps(reinterpret_cast<const float*>(&tw[j]));

                __m128 d_re = _mm_shuffle_ps(d, d, 0xA0);
                __m128 d_im = _mm_shuffle_ps(d, d, 0xF5);
                __m128 w_flip = _mm_shuffle_ps(w, w, 0xB1);

                __m128 p1 = _mm_mul_ps(d_re, w);
                __m128 p2 = _mm_mul_ps(d_im, w_flip);
                __m128 v  = _mm_add_ps(p1, _mm_mul_ps(p2, sign4));

                _mm_storeu_ps(reinterpret_cast<float*>(&data[k]),            _mm_add_ps(u, v));
                _mm_storeu_ps(reinterpret_cast<float*>(&data[k + half_len]), _mm_sub_ps(u, v));
            }
        }
#endif
        // scalar remainder
        for (; j < half_len; ++j) {
            int k = i + j;
            auto uu = data[k];
            auto vv = data[k + half_len] * tw[j];
            data[k]            = uu + vv;
            data[k + half_len] = uu - vv;
        }
    }
}

#endif // SANGI_FFT_HAS_AVX2

// =============================================================================
// Dispatch: select the optimal kernel for each stage
// =============================================================================

template<typename Real>
inline void butterfly_stage(
    Complex<Real>* data, int n, int half_len,
    const Complex<Real>* tw)
{
    if constexpr (std::is_same_v<Real, double>) {
#if defined(SANGI_FFT_HAS_AVX2)
        if (half_len >= 2) {
            butterfly_stage_avx2_double(data, n, half_len, tw);
            return;
        }
#elif defined(SANGI_FFT_HAS_SSE2)
        if (half_len >= 1) {
            butterfly_stage_sse2_double(data, n, half_len, tw);
            return;
        }
#endif
        butterfly_stage_scalar(data, n, half_len, tw);
    }
    else if constexpr (std::is_same_v<Real, float>) {
#if defined(SANGI_FFT_HAS_AVX2)
        if (half_len >= 4) {
            butterfly_stage_avx2_float(data, n, half_len, tw);
            return;
        }
#endif
#if defined(SANGI_FFT_HAS_SSE2)
        if (half_len >= 2) {
            butterfly_stage_sse2_float(data, n, half_len, tw);
            return;
        }
#endif
        butterfly_stage_scalar(data, n, half_len, tw);
    }
    else {
        butterfly_stage_scalar(data, n, half_len, tw);
    }
}

// =============================================================================
// Main FFT / IFFT (Radix-2, powers of 2 only)
// =============================================================================

/// forward transform (in-place)
template<typename Real>
inline void fft(Complex<Real>* data, int n)
{
    if (n <= 1) return;

    // build the twiddle table (thread_local cache)
    thread_local TwiddleTable<Real> table;
    thread_local int cached_n = 0;
    if (cached_n != n) {
        table.build(n);
        cached_n = n;
    }

    bit_reversal(data, n);

    int stage = 0;
    for (int len = 2; len <= n; len <<= 1) {
        int half = len / 2;
        butterfly_stage(data, n, half, table.twiddles[stage].data());
        ++stage;
    }
}

/// inverse transform (in-place, with 1/N normalization)
template<typename Real>
inline void ifft(Complex<Real>* data, int n)
{
    if (n <= 1) return;

    // take conjugate → FFT → conjugate → 1/N
    for (int i = 0; i < n; ++i) data[i] = conj(data[i]);
    fft(data, n);
    Real inv = Real(1) / Real(n);
    for (int i = 0; i < n; ++i) data[i] = conj(data[i]) * inv;
}

/// magnitude only (equivalent to performFrequencyOnlyForwardTransform)
/// input: n real numbers (float* / double*)
/// output: magnitudes in the first n/2 elements (in-place, rest undefined)
template<typename Real>
inline void magnitude_spectrum(Real* realData, int n)
{
    // convert real → complex (needs a 2n buffer)
    // assumes realData has a region of 2n elements (JUCE compatible)
    auto* cdata = reinterpret_cast<Complex<Real>*>(realData);

    // copy from back to front (safe in-place)
    for (int i = n - 1; i >= 0; --i) {
        cdata[i] = Complex<Real>(realData[i], Real(0));
    }

    fft(cdata, n);

    // magnitude computation
    int half = n / 2;
    for (int i = 0; i < half; ++i) {
        realData[i] = std::abs(cdata[i]);
    }
}

// =============================================================================
// Compile-time SIMD level information
// =============================================================================

constexpr const char* simd_level_name()
{
#if defined(SANGI_FFT_HAS_AVX2)
    return "AVX2";
#elif defined(SANGI_FFT_HAS_AVX)
    return "AVX";
#elif defined(SANGI_FFT_HAS_SSE2)
    return "SSE2";
#else
    return "Scalar";
#endif
}

} // namespace simd_fft
} // namespace sangi

#endif // SANGI_FFT_SIMD_HPP