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

// fft_utils.hpp
#ifndef SANGI_FFT_UTILS_HPP
#define SANGI_FFT_UTILS_HPP

#include "fft.hpp"
#include <vector>
#include <math/core/Complex.hpp>
#include <span>
#include <optional>

namespace sangi {

/**
 * @brief Polynomial multiplication using FFT
 * @tparam T Coefficient type
 * @param a Coefficients of polynomial 1
 * @param b Coefficients of polynomial 2
 * @return Coefficients of the product polynomial
 */
template<typename T>
std::vector<T> polynomial_multiply(std::span<const T> a, std::span<const T> b) {
    return FFT<T>::convolve(a, b);
}

/**
 * @brief Polynomial evaluation
 * @tparam T Coefficient type
 * @tparam U Type of the evaluation point
 * @param poly Coefficients of the polynomial
 * @param x Evaluation point
 * @return Value of the polynomial
 */
template<typename T, typename U>
auto evaluate_polynomial(std::span<const T> poly, const U& x) {
    using result_type = decltype(std::declval<T>() * std::declval<U>());
    
    if (poly.empty()) {
        return result_type{};
    }
    
    // Horner's method
    result_type result = poly.back();
    for (int i = static_cast<int>(poly.size()) - 2; i >= 0; --i) {
        result = result * x + poly[i];
    }
    
    return result;
}

/**
 * @brief Find a primitive root for the Number Theoretic Transform
 * @tparam P Modulus
 * @param n Transform size
 * @return The value of a suitable primitive root if one exists
 */
template<int P>
std::optional<ModularInt<P>> find_ntt_primitive_root(int n) {
    if ((n & (n - 1)) != 0) {
        // The size must be a power of two
        return std::nullopt;
    }

    // Check whether P-1 is divisible by 2 to the required power
    int64_t temp = P - 1;
    int required_power = 0;
    while (n > 1) {
        n /= 2;
        required_power++;
        
        if (temp % 2 != 0) {
            return std::nullopt;
        }
        temp /= 2;
    }
    
    // Search for candidate primitive roots
    for (int g = 2; g < P; ++g) {
        ModularInt<P> candidate(g);
        
        // Check at (P-1)/2
        if (candidate.pow((P - 1) / 2) == ModularInt<P>(1) && g != 1) {
            continue;
        }
        
        // Compute a suitable primitive root for n
        return candidate.pow((P - 1) >> required_power);
    }
    
    return std::nullopt;
}

/**
 * @brief Compute the amplitude of a frequency spectrum
 * @param spectrum Complex spectrum
 * @return Amplitude spectrum
 */
template<typename Real>
std::vector<Real> amplitude_spectrum(std::span<const Complex<Real>> spectrum) {
    std::vector<Real> amplitude(spectrum.size());
    
    for (size_t i = 0; i < spectrum.size(); ++i) {
        amplitude[i] = std::abs(spectrum[i]);
    }
    
    return amplitude;
}

/**
 * @brief Compute the phase of a frequency spectrum
 * @param spectrum Complex spectrum
 * @return Phase spectrum (radians, principal value $(-\pi, \pi]$)
 *
 * The return value is the principal value of arg(), so it is wrapped.
 * Apply unwrap_phase() when a continuous phase response is required.
 */
template<typename Real>
std::vector<Real> phase_spectrum(std::span<const Complex<Real>> spectrum) {
    std::vector<Real> phase(spectrum.size());

    for (size_t i = 0; i < spectrum.size(); ++i) {
        phase[i] = arg(spectrum[i]);
    }

    return phase;
}

/**
 * @brief Unwrap a phase array (1D)
 * @param phase Input phase (arbitrary -π..π principal values)
 * @param tol  Threshold for jump detection (default: π)
 * @return Unwrapped continuous phase
 *
 * When the difference between adjacent elements exceeds ±tol, add a multiple
 * of 2π to make it continuous.
 * Equivalent to NumPy's np.unwrap / MATLAB's unwrap.
 * Used for group-delay computation, minimum-phase estimation, complex cepstrum, etc.
 */
template<typename Real>
std::vector<Real> unwrap_phase(std::span<const Real> phase,
                               Real tol = std::numbers::pi_v<Real>) {
    std::vector<Real> unwrapped(phase.begin(), phase.end());
    if (unwrapped.size() < 2) return unwrapped;

    const Real two_pi = Real(2) * std::numbers::pi_v<Real>;
    Real offset = 0;
    for (size_t i = 1; i < unwrapped.size(); ++i) {
        Real diff = phase[i] - phase[i - 1];
        if (diff > tol)       offset -= two_pi;
        else if (diff < -tol) offset += two_pi;
        unwrapped[i] = phase[i] + offset;
    }
    return unwrapped;
}

/**
 * @brief Reconstruct a complex spectrum from amplitude and phase
 * @param amplitude Amplitude spectrum
 * @param phase Phase spectrum (radians)
 * @return Complex spectrum
 */
template<typename Real>
std::vector<Complex<Real>> reconstruct_spectrum(
    std::span<const Real> amplitude, 
    std::span<const Real> phase) {
    
    if (amplitude.size() != phase.size()) {
        throw MathError("Amplitude and phase must have the same size");
    }
    
    std::vector<Complex<Real>> spectrum(amplitude.size());
    
    for (size_t i = 0; i < amplitude.size(); ++i) {
        spectrum[i] = polar(amplitude[i], phase[i]);
    }
    
    return spectrum;
}

// ================================================================
// Discrete Cosine Transform (DCT-II) — FFT-based O(n log n)
// ================================================================
// Normalized DCT-II: X[k] = c(k) · Σ_{i=0}^{N-1} x[i] · cos(π(i+0.5)k/N)
// c(0) = √(1/N), c(k) = √(2/N) (k>0)
// Makhoul's method: reorder the input and compute with an N-point FFT

namespace detail {

/// O(n²) direct DCT-II (fallback for sizes unsuitable for FFT)
template<typename Real>
std::vector<Real> dct_direct(std::span<const Real> data) {
    const int n = static_cast<int>(data.size());
    const Real pi_over_n = std::numbers::pi_v<Real> / n;
    std::vector<Real> result(n);
    for (int k = 0; k < n; ++k) {
        Real sum = 0;
        Real scale = (k == 0) ? std::sqrt(Real(1) / n) : std::sqrt(Real(2) / n);
        for (int i = 0; i < n; ++i)
            sum += data[i] * std::cos(pi_over_n * (i + Real(0.5)) * k);
        result[k] = scale * sum;
    }
    return result;
}

/// O(n²) direct IDCT (DCT-III)
template<typename Real>
std::vector<Real> idct_direct(std::span<const Real> coeffs) {
    const int n = static_cast<int>(coeffs.size());
    const Real pi_over_n = std::numbers::pi_v<Real> / n;
    std::vector<Real> result(n);
    for (int i = 0; i < n; ++i) {
        Real sum = coeffs[0] * std::sqrt(Real(1) / n);
        for (int k = 1; k < n; ++k)
            sum += coeffs[k] * std::sqrt(Real(2) / n) * std::cos(pi_over_n * (i + Real(0.5)) * k);
        result[i] = sum;
    }
    return result;
}

/// O(n²) direct DST-II
template<typename Real>
std::vector<Real> dst_direct(std::span<const Real> data) {
    const int n = static_cast<int>(data.size());
    const Real pi_over_n = std::numbers::pi_v<Real> / n;
    std::vector<Real> result(n);
    for (int k = 0; k < n; ++k) {
        Real sum = 0;
        Real scale = (k == n - 1) ? std::sqrt(Real(1) / n) : std::sqrt(Real(2) / n);
        for (int i = 0; i < n; ++i)
            sum += data[i] * std::sin(pi_over_n * (i + Real(0.5)) * (k + 1));
        result[k] = scale * sum;
    }
    return result;
}

/// O(n²) direct IDST (DST-III)
template<typename Real>
std::vector<Real> idst_direct(std::span<const Real> coeffs) {
    const int n = static_cast<int>(coeffs.size());
    const Real pi_over_n = std::numbers::pi_v<Real> / n;
    std::vector<Real> result(n);
    for (int i = 0; i < n; ++i) {
        Real sum = Real(0);
        for (int k = 0; k < n; ++k) {
            Real c = (k == n - 1) ? std::sqrt(Real(1) / n) : std::sqrt(Real(2) / n);
            sum += c * coeffs[k] * std::sin(pi_over_n * (i + Real(0.5)) * (k + 1));
        }
        result[i] = sum;
    }
    return result;
}

} // namespace detail

/**
 * @brief Discrete Cosine Transform (DCT-II)
 *
 * If N is an FFT-supported size (2^a·3^b·5^c), uses Makhoul's method O(n log n).
 * Otherwise falls back to direct computation O(n²).
 *
 * @param data Input data (size N)
 * @return DCT coefficients (size N)
 */
template<typename Real>
std::vector<Real> dct(std::span<const Real> data) {
    const int n = static_cast<int>(data.size());
    if (n == 0) return {};
    if (n == 1) return { data[0] };

    using complex_type = Complex<Real>;
    if (!FFT<complex_type, Real>::is_valid_size(n))
        return detail::dct_direct(data);

    // Makhoul's method: y[j] = x[2j] (even indices), y[N-1-j] = x[2j+1] (odd indices)
    std::vector<complex_type> buf(n);
    for (int j = 0; j < (n + 1) / 2; ++j)
        buf[j] = complex_type(data[2 * j], 0);
    for (int j = 0; j < n / 2; ++j)
        buf[n - 1 - j] = complex_type(data[2 * j + 1], 0);

    FFT<complex_type, Real>::fft(std::span<complex_type>(buf));

    // Multiply by the twiddle factor and take the real part
    // The FFT uses e^{-i2πkn/N} (negative sign), so the twist is -π/(2N)
    const Real pi_over_2n = std::numbers::pi_v<Real> / (2 * n);
    std::vector<Real> result(n);
    for (int k = 0; k < n; ++k) {
        Real scale = (k == 0) ? std::sqrt(Real(1) / n) : std::sqrt(Real(2) / n);
        Real angle = -pi_over_2n * k;
        complex_type W(std::cos(angle), std::sin(angle));
        result[k] = scale * (buf[k] * W).real();
    }
    return result;
}

/**
 * @brief Inverse Discrete Cosine Transform (IDCT, DCT-III)
 * @param coeffs DCT coefficients (size N)
 * @return Reconstructed data (size N)
 */
template<typename Real>
std::vector<Real> idct(std::span<const Real> coeffs) {
    const int n = static_cast<int>(coeffs.size());
    if (n == 0) return {};
    if (n == 1) return { coeffs[0] };

    using complex_type = Complex<Real>;
    if (!FFT<complex_type, Real>::is_valid_size(n))
        return detail::idct_direct(coeffs);

    // DCT-III (IDCT): multiply normalization coefficients by the twiddle factor, IFFT → inverse reorder
    // The IFFT uses e^{+i2πkn/N} (positive sign), so the twist is +π/(2N)
    const Real pi_over_2n = std::numbers::pi_v<Real> / (2 * n);
    std::vector<complex_type> buf(n);
    for (int k = 0; k < n; ++k) {
        Real c = (k == 0) ? std::sqrt(Real(1) / n) : std::sqrt(Real(2) / n);
        Real angle = pi_over_2n * k;
        complex_type W(std::cos(angle), std::sin(angle));
        buf[k] = complex_type(c * coeffs[k]) * W;
    }

    FFT<complex_type, Real>::ifft(std::span<complex_type>(buf));

    // The IFFT is already 1/N normalized → multiply by N to restore
    // Inverse reorder: y → x
    std::vector<Real> result(n);
    for (int j = 0; j < (n + 1) / 2; ++j)
        result[2 * j] = buf[j].real() * n;
    for (int j = 0; j < n / 2; ++j)
        result[2 * j + 1] = buf[n - 1 - j].real() * n;

    return result;
}

/**
 * @brief fast_dct — alias for dct (backward compatibility)
 */
template<typename Real>
std::vector<Real> fast_dct(std::span<const Real> data) {
    return dct(data);
}

// ================================================================
// Discrete Sine Transform (DST-II)
// ================================================================
// Normalized DST-II: X[k] = c(k) · Σ_{i=0}^{N-1} x[i] · sin(π(i+0.5)(k+1)/N)
// c(N-1) = √(1/N), c(k) = √(2/N) (k<N-1)

/**
 * @brief Discrete Sine Transform (DST-II)
 *
 * Exploits the relationship with DCT-II: reverse the input → DCT-II → alternately negate the coefficients.
 * O(n log n) (when N is an FFT-supported size).
 *
 * @param data Input data (size N)
 * @return DST coefficients (size N)
 */
template<typename Real>
std::vector<Real> dst(std::span<const Real> data) {
    const int n = static_cast<int>(data.size());
    if (n == 0) return {};
    if (n == 1) return { data[0] };

    // DST-II(x)[k] = (-1)^k · DCT-II(x_reversed)[N-1-k]
    // However, the DST normalization is c(N-1)=√(1/N) while the DCT normalization is c(0)=√(1/N)
    // → fall back to direct computation
    return detail::dst_direct(data);
}

/**
 * @brief Inverse Discrete Sine Transform (IDST, DST-III)
 * @param coeffs DST coefficients (size N)
 * @return Reconstructed data (size N)
 */
template<typename Real>
std::vector<Real> idst(std::span<const Real> coeffs) {
    const int n = static_cast<int>(coeffs.size());
    if (n == 0) return {};
    if (n == 1) return { coeffs[0] };
    return detail::idst_direct(coeffs);
}

// ================================================================
// Discrete Hankel Transform (DHT) — equivalent to GSL gsl_dht
// ================================================================
// Hankel transform of arbitrary order ν:
//   F(k_m) = Σ_{n=0}^{N-1} T_{mn} · f(r_n)
// T_{mn} = (2/j_{ν,N+1}²) · J_ν(j_{ν,m+1}·j_{ν,n+1}/j_{ν,N+1})
//                          / |J_{ν+1}(j_{ν,m+1})·J_{ν+1}(j_{ν,n+1})|
// j_{ν,s} is the s-th zero of J_ν
//
// ν=0: fast path (McMahon asymptotic approximation + J₀/J₁ approximation)
// ν>0: besselJ(nu, x) + zero search via Brent's method

namespace detail {

/// Asymptotic approximation of the zeros of J₀ (McMahon expansion)
/// j_{0,s} ≈ β - 1/(8β) - 31/(384β³) - 3779/(15360β⁵)
/// β = (s - 1/4)·π
template<typename Real>
Real besselJ0_zero(int s) {
    // McMahon asymptotic initial estimate
    Real beta = (s - Real(0.25)) * std::numbers::pi_v<Real>;
    Real b_inv = Real(1) / beta;
    Real b_inv2 = b_inv * b_inv;
    Real x = beta - b_inv / Real(8) * (Real(1)
        + b_inv2 * (Real(31) / Real(48)
        + b_inv2 * Real(3779) / Real(1920)));
    // Newton-Raphson refinement: J₀(x)=0, J₀'(x) = -J₁(x)
    for (int iter = 0; iter < 10; ++iter) {
        Real j0 = static_cast<Real>(std::cyl_bessel_j(0.0, static_cast<double>(x)));
        Real j1 = static_cast<Real>(std::cyl_bessel_j(1.0, static_cast<double>(x)));
        if (std::abs(j1) < Real(1e-30)) break;
        Real dx = j0 / j1;  // Newton: x -= J₀/J₀' = J₀/(-J₁), i.e. x += J₀/J₁
        x += dx;
        if (std::abs(dx) < std::abs(x) * std::numeric_limits<Real>::epsilon()) break;
    }
    return x;
}

/// J₀(x) — Taylor series (|x| < 20) / asymptotic expansion (|x| ≥ 20)
template<typename Real>
Real besselJ0_approx(Real x) {
    Real ax = std::abs(x);
    if (ax < Real(1e-15)) return Real(1);
    if (ax < Real(20)) {
        // J₀(x) = Σ_{k=0}^∞ (-1)^k (x/2)^{2k} / (k!)²
        Real half_x = x / Real(2);
        Real half_x2 = half_x * half_x;
        Real term = Real(1);
        Real sum = Real(1);
        for (int k = 1; k <= 40; ++k) {
            term *= -half_x2 / (Real(k) * Real(k));
            sum += term;
            if (std::abs(term) < std::abs(sum) * std::numeric_limits<Real>::epsilon()) break;
        }
        return sum;
    }
    // Asymptotic expansion: J₀(x) ≈ √(2/(πx)) cos(x - π/4)
    Real sq = std::sqrt(Real(2) / (std::numbers::pi_v<Real> * ax));
    Real theta = ax - std::numbers::pi_v<Real> / Real(4);
    return sq * std::cos(theta);
}

/// J₁(x) — Taylor series (|x| < 20) / asymptotic expansion (|x| ≥ 20)
template<typename Real>
Real besselJ1_approx(Real x) {
    Real ax = std::abs(x);
    if (ax < Real(1e-15)) return x / Real(2);
    if (ax < Real(20)) {
        // J₁(x) = (x/2) · Σ_{k=0}^∞ (-1)^k (x/2)^{2k} / (k! · (k+1)!)
        Real half_x = x / Real(2);
        Real half_x2 = half_x * half_x;
        Real term = Real(1);
        Real sum = Real(1);
        for (int k = 1; k <= 40; ++k) {
            term *= -half_x2 / (Real(k) * Real(k + 1));
            sum += term;
            if (std::abs(term) < std::abs(sum) * std::numeric_limits<Real>::epsilon()) break;
        }
        return half_x * sum;
    }
    Real sq = std::sqrt(Real(2) / (std::numbers::pi_v<Real> * ax));
    Real theta = ax - Real(3) * std::numbers::pi_v<Real> / Real(4);
    Real sgn = (x < Real(0)) ? Real(-1) : Real(1);
    return sgn * sq * std::cos(theta);
}

/// Wrapper for J_ν(x) (uses std::cyl_bessel_j)
template<typename Real>
Real besselJnu(Real nu, Real x) {
    return static_cast<Real>(std::cyl_bessel_j(
        static_cast<double>(nu), static_cast<double>(x)));
}

/// Search for the s-th positive zero of J_ν (McMahon initial estimate + bisection refinement)
template<typename Real>
Real besselJnu_zero(Real nu, int s) {
    // ν=0 uses the fast path
    if (std::abs(nu) < Real(1e-12))
        return besselJ0_zero<Real>(s);

    // McMahon asymptotic initial estimate: j_{ν,s} ≈ (s + ν/2 - 1/4) · π
    Real beta = (Real(s) + nu / Real(2) - Real(0.25)) * std::numbers::pi_v<Real>;

    // Find a sign change within the search interval
    Real lo = std::max(Real(1e-6), beta - std::numbers::pi_v<Real>);
    Real hi = beta + std::numbers::pi_v<Real> / Real(2);
    Real flo = besselJnu(nu, lo);

    constexpr int maxSearch = 100;
    Real step = (hi - lo) / Real(maxSearch);
    for (int i = 1; i <= maxSearch; ++i) {
        Real x = lo + step * Real(i);
        Real fx = besselJnu(nu, x);
        if (flo * fx < Real(0)) {
            hi = x;
            break;
        }
        lo = x;
        flo = fx;
    }

    // Refine with bisection
    for (int iter = 0; iter < 60; ++iter) {
        Real mid = (lo + hi) * Real(0.5);
        Real fmid = besselJnu(nu, mid);
        if (std::abs(fmid) < std::numeric_limits<Real>::epsilon() * Real(10))
            return mid;
        if (flo * fmid < Real(0))
            hi = mid;
        else {
            lo = mid;
            flo = fmid;
        }
    }
    return (lo + hi) * Real(0.5);
}

} // namespace detail

/**
 * @brief Discrete Hankel Transform (supports arbitrary order ν, equivalent to GSL gsl_dht)
 *
 * ν=0: fast path (McMahon + J₀/J₁ approximation)
 * ν>0: besselJ(nu, x) + zero search
 *
 * Sample points:    r_n = j_{ν,n+1} · R / j_{ν,N+1}
 * Frequency points: k_m = j_{ν,m+1} / R
 * Transform matrix: T_{mn} = (2/j²_{ν,N+1}) · J_ν(...) / |J_{ν+1}(...)|²
 *
 * @param f Input data (size N)
 * @param R Radius of the spatial domain
 * @return Result of the Hankel transform (size N)
 */
template<typename Real>
struct HankelTransform {
    int N;
    Real R;
    Real nu;                     // order ν (default 0)
    std::vector<Real> r;         // sample points
    std::vector<Real> k;         // frequency points
    std::vector<std::vector<Real>> Tmat;  // transform matrix

    /// Construct a discrete Hankel transform of arbitrary order ν
    HankelTransform(int n, Real radius, Real order = Real(0))
        : N(n), R(radius), nu(order), r(n), k(n),
          Tmat(n, std::vector<Real>(n))
    {
        // Compute the zeros of J_ν
        std::vector<Real> zeros(n + 1);
        for (int s = 1; s <= n + 1; ++s)
            zeros[s - 1] = detail::besselJnu_zero<Real>(nu, s);

        Real jN1 = zeros[n];  // j_{ν,N+1}

        // Sample points and frequency points
        for (int i = 0; i < n; ++i) {
            r[i] = zeros[i] * R / jN1;
            k[i] = zeros[i] / R;
        }

        // Forward transform matrix:
        // T[m][n] = (2 / j_{N+1}) · J_ν(j_m·j_n / j_{N+1}) / J_{ν+1}(j_n)²
        // Inverse transform matrix (for T⁻¹[n][m] the denominator changes to J_{ν+1}(j_m)²):
        // T⁻¹[n][m] = (2 / j_{N+1}) · J_ν(j_m·j_n / j_{N+1}) / J_{ν+1}(j_m)²
        Real scale = Real(2) / jN1;
        std::vector<Real> Jnup1_sq(n);
        for (int i = 0; i < n; ++i) {
            Real j = detail::besselJnu(nu + Real(1), zeros[i]);
            Jnup1_sq[i] = j * j;
        }

        Tinv_.resize(n, std::vector<Real>(n));
        for (int m = 0; m < n; ++m) {
            for (int nn = 0; nn < n; ++nn) {
                Real arg = zeros[m] * zeros[nn] / jN1;
                Real Jval = detail::besselJnu(nu, arg);
                Tmat[m][nn]  = scale * Jval / Jnup1_sq[nn];
                Tinv_[nn][m] = scale * Jval / Jnup1_sq[m];
            }
        }
    }

    /// Forward transform: f(r) → F(k)
    std::vector<Real> transform(std::span<const Real> f) const {
        std::vector<Real> result(N, Real(0));
        for (int m = 0; m < N; ++m)
            for (int n = 0; n < N; ++n)
                result[m] += Tmat[m][n] * f[n];
        return result;
    }

    /// Inverse transform: F(k) → f(r)
    std::vector<Real> inverse(std::span<const Real> F) const {
        std::vector<Real> result(N, Real(0));
        for (int n = 0; n < N; ++n)
            for (int m = 0; m < N; ++m)
                result[n] += Tinv_[n][m] * F[m];
        return result;
    }

private:
    std::vector<std::vector<Real>> Tinv_;
};

// ================================================================
// Discrete Hartley Transform
// ================================================================
// cas(θ) = cos(θ) + sin(θ)
// H[k] = Σ_{n=0}^{N-1} x[n] · cas(2πnk/N)
// Real input → real output. A real-valued alternative to the FFT. Self-inverse (1/N normalization).

/**
 * @brief Discrete Hartley Transform (DHartleyT)
 *
 * @param data Input data (size N)
 * @return Hartley transform coefficients (size N)
 */
template<typename Real>
[[nodiscard]]
std::vector<Real> hartley(std::span<const Real> data) {
    const std::size_t n = data.size();
    if (n == 0) return {};
    if (n == 1) return {data[0]};

    // Via FFT: H[k] = Re(X[k]) - Im(X[k])
    // where X = FFT(x)
    using complex_type = Complex<Real>;
    std::vector<complex_type> cx(n);
    for (std::size_t i = 0; i < n; ++i)
        cx[i] = complex_type(data[i], Real(0));

    FFT<complex_type, Real>::fft(std::span<complex_type>(cx));

    std::vector<Real> result(n);
    for (std::size_t k = 0; k < n; ++k)
        result[k] = cx[k].real() - cx[k].imag();

    return result;
}

/**
 * @brief Inverse Discrete Hartley Transform (iHartley)
 *
 * The Hartley transform is self-inverse: x[n] = (1/N) · H[H[k]]
 */
template<typename Real>
[[nodiscard]]
std::vector<Real> ihartley(std::span<const Real> coeffs) {
    auto result = hartley(coeffs);
    Real inv_n = Real(1) / static_cast<Real>(result.size());
    for (auto& v : result) v *= inv_n;
    return result;
}

// ================================================================
// Walsh-Hadamard Transform (Fast Walsh-Hadamard Transform)
// ================================================================
// In-place butterfly operations: O(n log n)
// Input size must be a power of two.
// No normalization (forward transform). The inverse is the same operation + 1/N normalization.

/**
 * @brief Walsh-Hadamard Transform (in-place)
 *
 * The size of data must be a power of two.
 */
template<typename Real>
void walshHadamardInplace(std::vector<Real>& data) {
    const std::size_t n = data.size();
    if (n <= 1) return;

    // Check whether n is a power of two
    if ((n & (n - 1)) != 0)
        throw std::invalid_argument("walshHadamard: size must be a power of 2");

    for (std::size_t len = 1; len < n; len <<= 1) {
        for (std::size_t i = 0; i < n; i += 2 * len) {
            for (std::size_t j = 0; j < len; ++j) {
                Real u = data[i + j];
                Real v = data[i + j + len];
                data[i + j]       = u + v;
                data[i + j + len] = u - v;
            }
        }
    }
}

/**
 * @brief Walsh-Hadamard Transform (copy version)
 *
 * @param data Input data (size = 2^k)
 * @return WHT coefficients
 */
template<typename Real>
[[nodiscard]]
std::vector<Real> walshHadamard(std::span<const Real> data) {
    std::vector<Real> result(data.begin(), data.end());
    walshHadamardInplace(result);
    return result;
}

/**
 * @brief Inverse Walsh-Hadamard Transform
 *
 * The WHT is self-inverse: x = (1/N) · WHT(WHT(x))
 */
template<typename Real>
[[nodiscard]]
std::vector<Real> iwalshHadamard(std::span<const Real> coeffs) {
    std::vector<Real> result(coeffs.begin(), coeffs.end());
    walshHadamardInplace(result);
    Real inv_n = Real(1) / static_cast<Real>(result.size());
    for (auto& v : result) v *= inv_n;
    return result;
}

// ================================================================
// Discrete Hilbert Transform
// ================================================================
// Constructs the analytic signal:
//   z[n] = x[n] + j·H{x}[n]
// where H{x} is the Hilbert transform of x.
//
// Via FFT: X = FFT(x), zero out negative frequencies, halve DC and Nyquist
//   Z[k] = { X[k]         k=0, k=N/2
//          { 2·X[k]       0 < k < N/2
//          { 0             N/2 < k < N
//   z = IFFT(Z)
// Im(z) is the Hilbert transform.

/**
 * @brief Compute the analytic signal
 *
 * @param data Real input signal (size N)
 * @return Analytic signal (complex, size N)
 */
template<typename Real>
[[nodiscard]]
std::vector<Complex<Real>> analyticSignal(std::span<const Real> data) {
    const std::size_t n = data.size();
    if (n == 0) return {};

    // FFT
    using complex_type = Complex<Real>;
    std::vector<complex_type> X(n);
    for (std::size_t i = 0; i < n; ++i)
        X[i] = complex_type(data[i], Real(0));

    FFT<complex_type, Real>::fft(std::span<complex_type>(X));

    // Zero out negative frequencies
    std::size_t half = n / 2;
    // k=0: leave as is
    // k=1..half-1: double
    for (std::size_t k = 1; k < half; ++k)
        X[k] *= Real(2);
    // k=half (Nyquist): leave as is (for even length)
    // k=half+1..n-1: zero
    for (std::size_t k = half + 1; k < n; ++k)
        X[k] = Complex<Real>(0, 0);

    // IFFT
    FFT<complex_type, Real>::ifft(std::span<complex_type>(X));

    return X;
}

/**
 * @brief Discrete Hilbert Transform
 *
 * @param data Real input signal (size N)
 * @return Hilbert transform result (real, size N)
 */
template<typename Real>
[[nodiscard]]
std::vector<Real> hilbertTransform(std::span<const Real> data) {
    auto z = analyticSignal(data);
    std::vector<Real> result(z.size());
    for (std::size_t i = 0; i < z.size(); ++i)
        result[i] = z[i].imag();
    return result;
}

/**
 * @brief Instantaneous amplitude / envelope
 *
 * |z[n]| = √(x[n]² + H{x}[n]²)
 */
template<typename Real>
[[nodiscard]]
std::vector<Real> instantaneousAmplitude(std::span<const Real> data) {
    auto z = analyticSignal(data);
    std::vector<Real> result(z.size());
    for (std::size_t i = 0; i < z.size(); ++i)
        result[i] = std::abs(z[i]);
    return result;
}

/**
 * @brief Instantaneous frequency
 *
 * ω[n] = d/dt arg(z[n]) ≈ (arg(z[n+1]) - arg(z[n])) · fs / (2π)
 */
template<typename Real>
[[nodiscard]]
std::vector<Real> instantaneousFrequency(std::span<const Real> data, Real sampleRate = Real(1)) {
    auto z = analyticSignal(data);
    const std::size_t n = z.size();
    if (n < 2) return {};

    std::vector<Real> freq(n - 1);
    Real inv2pi = sampleRate / (Real(2) * std::numbers::pi_v<Real>);
    for (std::size_t i = 0; i < n - 1; ++i) {
        Real dphi = arg(z[i + 1]) - arg(z[i]);
        // Phase unwrapping: normalize to [-π, π]
        while (dphi > std::numbers::pi_v<Real>)  dphi -= Real(2) * std::numbers::pi_v<Real>;
        while (dphi < -std::numbers::pi_v<Real>) dphi += Real(2) * std::numbers::pi_v<Real>;
        freq[i] = dphi * inv2pi;
    }
    return freq;
}

} // namespace sangi

#endif // SANGI_FFT_UTILS_HPP