/* Copyright (C) 2025 Joel Yliluoma — Material for https://youtu.be/C4tHxouarGc - https://iki.fi/bisqwit/ */
/* MIT license follows:
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software,
and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of
the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS
OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT
OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/
#include <unordered_map>
#include <cassert>
#include <numbers>
#include <chrono>
#include <cstdio>
#include <format>
#include <memory>
#include <string>
#include <array>
#include <mutex>
#include <map>
#include <set>
#include "fft.hh"

template<bool unique, std::size_t limit>
static auto get_factors_gen(std::size_t N) // Builds a list of factors of N, up to limit. Returns the number of primes written.
{
    std::array<std::size_t,limit> target;
    static class lore /* This structure caches prime-related information. */
    {
        std::unordered_map<std::size_t, std::size_t> smallest_factors{ {0,0}, {1,1}, {2,2} };
        std::vector<signed char> is_prime{1,1,1}; // 0=not known, -1=no, 1=yes
        std::size_t              lastprime{2};    // extent of contiguous knowledge
        std::set<std::size_t>    primes{2};       // list of known primes
        std::mutex lock{};
    public:
        auto get_lock() { return std::unique_lock<std::mutex>(lock); }
        std::size_t find(std::size_t N, std::unique_lock<std::mutex>&)
        {
            if(auto i = smallest_factors.find(N); i != smallest_factors.end()) return i->second;
            std::size_t solution = N;
            for(auto p: primes)
            {
                if(N % p == 0) { solution = p; break; }
                if(p*p >= N) { break; }
            }
            for(std::size_t p = lastprime|1; p*p <= N; p+=2) if(N % p == 0) { solution = p; break; }
            if(is_prime.size() <= N) is_prime.resize(N+1);
            for(is_prime[N] = (solution==N) ? 1 : -1; lastprime+1 < is_prime.size() && is_prime[lastprime+1] != 0; ) ++lastprime;
            if(solution == N) primes.insert(solution);
            return smallest_factors.emplace(N, solution).first->second;
        }
        lore() { }
    } prime_lore;
    auto lock = prime_lore.get_lock();
    for(std::size_t orig=N,count=0,prev=0,f; ; N/=f)
        if(f=prime_lore.find(N, lock); count >= limit || f <= 1 || f == orig) { return std::pair(target,count); }
        else if(f != prev || !unique) { target[count++] = prev = f; }
}

// log2 of the product of first 16 primes ≈ 64.82. If size_t ≤ 64 bits, then size 16 array is enough.
auto& get_factors_unique = get_factors_gen<true, 16>;

// The smallest prime is 2. The largest size_t is 2^64, which has 64 factors of 2.
auto& get_factors_all    = get_factors_gen<false, 64>;

static const float pi = std::numbers::pi_v<float>;
static complex w(std::size_t k, std::size_t N) { return std::polar(1.0f, (-2 * pi) * k / N); }

#define BEGIN_DFT_METHOD(name,is_inplace,tpl,getname) \
    tpl struct DFT_##name: public DFTbase \
    { \
        DFT_##name(long n, std::monostate={}): DFTbase(n,is_inplace){} \
        virtual std::string get_name() const override { return getname; } \
        virtual void xform_many(const complex* in, std::size_t istep,std::size_t istep2, \
                                complex* out, std::size_t ostep,std::size_t ostep2, std::size_t num) override \
        { \
            if(!num) [[unlikely]] return; \
            if(num==1)      xform_impl<false>(in,istep,istep2,out,ostep,ostep2,num); \
            else [[likely]] xform_impl<true>(in,istep,istep2,out,ostep,ostep2,num); \
        } \
        template<bool many> \
        void xform_impl(const complex* in, std::size_t istep,std::size_t istep2, \
                        complex* out, std::size_t ostep,std::size_t ostep2, std::size_t num) \
        { \
            if constexpr(!many) { assert(istep==1); assert(ostep==1); istep=ostep=1; istep2=ostep2=0; num=1; }
#define END_DFT_METHOD \
    } \
}
#define DFT_CONSTRUCT(name, params, ...) DFT_##name params: __VA_ARGS__ {
#define DFT_PROPERTIES }

// Discrete Fourier Transform: Xₙ = ∑ₖ₌₀ᴺ⁻¹ (xₖ · exp(-2iπnk/N))
BEGIN_DFT_METHOD(DFT,false,, "DFT("+std::to_string(N)+')')
    std::vector<complex> tmp(num);
    for(std::size_t a=0; a<N; ++a)
    {
        for(std::size_t n=0; n<num; ++n) tmp[n] = in[0*istep + n*istep2];
        for(std::size_t b=1; b<N; ++b)
        {
            complex s = w(a*b, N);
            for(std::size_t n=0; n<num; ++n) tmp[n] += in[b*istep + n*istep2] * s;
        }
        for(std::size_t n=0; n<num; ++n) out[a*ostep + n*ostep2] = tmp[n];
    }
END_DFT_METHOD;

// Cooley-Tukey — A fast method for radix-P sizes (the original FFT!)
BEGIN_DFT_METHOD(RadixP,false, template<std::size_t P>, std::format("RadixP<{}>({})", P,N))
    std::size_t Q = N/P; // P*Q = N
    /* P transforms of Q size each */
    DFTbase* DFTq = FindDFT(Q, P*istep,istep, ostep,ostep*Q, P);
    for(std::size_t n=0; n<num; ++n)
        DFTq->xform_many(in + n*istep2,P*istep,istep,  out + n*ostep2,ostep,ostep*Q, P);
    for(std::size_t b=1; b<P; ++b)
        for(std::size_t a=1; a<Q; ++a)
        {
            complex s = w(a*b, N);
            for(std::size_t n=0; n<num; ++n)
                out[ostep*(b*Q+a) + n*ostep2] *= s;
        }
    /* Q transforms of P size each */
    DFTbase* DFTp = FindDFT(P, ostep*Q,ostep, ostep*Q,ostep,  Q);
    for(std::size_t n=0; n<num; ++n)
        DFTp->xform_inplace_many(out + n*ostep2,ostep*Q,ostep, Q);
END_DFT_METHOD;

// Rader — It makes FFT possible on prime sizes! Fast, on small values of N.
BEGIN_DFT_METHOD(Rader,true,, std::format("Rader({})", N))
    DFTbase* DFTn1 = FindDFT(N-1, 1,N-1, 1,N-1, num); 
    std::vector<complex> buf(num + (N-1) * num*2);
    for(std::size_t n=0; n<num; ++n) buf[n] = in[0*istep + n*istep2];
    // Permutation of input
    for(std::size_t gₚ=1, k=0; k<N-1; ++k, gₚ = (gₚ*g)%N)
        for(std::size_t n=0; n<num; ++n)
            buf[num+k+n*(N-1)] = in[gₚ*istep + n*istep2];
    // DFT
    DFTn1->xform_many(&buf[num],1,N-1, &buf[num+(N-1)*num],1,N-1, num);
    // Set output DC component:
    for(std::size_t n=0; n<num; ++n) out[0*ostep + n*ostep2] = buf[n] + buf[num+(N-1)*(n+num)];
    // Multiply by ω:
    for(std::size_t n=0; n<num; ++n)
        for(std::size_t k=0; k<N-1; ++k)
            buf[num+k+(N-1)*(n+num)] = std::conj( buf[num+k+(N-1)*(n+num)] * ω[k] + (k==0 ? buf[n] : 0) );
    // IDFT 
    DFTn1->xform_many(&buf[num+(N-1)*num],1,N-1, &buf[num],1,N-1, num);
    // Inverse permutation
    for(std::size_t gₚ=1, k=0; k<N-1; ++k, gₚ = (gₚ*gᵢₙᵥ)%N)
        for(std::size_t n=0; n<num; ++n)
            out[gₚ*ostep + n*ostep2] = std::conj(buf[num + k + n*(N-1)]);
  DFT_PROPERTIES
    std::size_t g, gᵢₙᵥ;
    std::vector<complex> ω;
    /* Compute n^m mod p, where m >= 0 and p > 0. */
    static std::size_t powermod(std::size_t n, std::size_t m, std::size_t p)
    {
        if(m == 0) return 1;
        if(m % 2 == 0) { n = powermod(n, m/2, p); return (n*n) % p; }
        return (n * powermod(n, m-1, p)) % p;
    }
  DFT_CONSTRUCT(Rader, (std::size_t N), DFTbase{N,true}, g{2}, ω(N-1))
    auto [factors, count] = get_factors_unique(N-1);
    for(std::size_t m=0; m<count; )
        if(powermod(g, (N-1) / factors[m++], N) == 1)
            { ++g; m = 0; }
    gᵢₙᵥ = powermod(g,N-2,N);
    for(std::size_t gₚ=1, i=0; i < N-1; ++i, gₚ = (gₚ*gᵢₙᵥ)%N)
        { ω[i] = w(gₚ, N); }
    FindDFT(N-1, 1,1, 1,1, 1)->xform_inplace(&ω[0], 1);
    for(std::size_t i=0; i<N-1; ++i) ω[i] /= float(N-1);
END_DFT_METHOD;

// Bluestein — An FFT method for *all* sizes (not always fastest though)
BEGIN_DFT_METHOD(Bluestein,true,, std::format("Bluestein({})(nb={})", N, nb))
    DFTbase* DFTnb = FindDFT(nb,1,nb*2, 1,nb*2, num);
    std::vector<complex> buf(nb*2 * num);
    for(std::size_t n=0; n<num; ++n)
    for(std::size_t k=0; k<N; ++k) { buf[k+n*nb*2] = in[k*istep+n*istep2] * ω₀[k]; }
    DFTnb->xform_many(&buf[0],1,nb*2, &buf[nb],1,nb*2, num); // DFT 
    for(std::size_t n=0; n<num; ++n)
    for(std::size_t j=0; j<nb; ++j) { buf[nb+j+n*nb*2] = std::conj(buf[nb+j+n*nb*2]) * ω₁[j]; }
    DFTnb->xform_many(&buf[nb],1,nb*2, &buf[0],1,nb*2, num); // IDFT
    for(std::size_t n=0; n<num; ++n)
    for(std::size_t k=0; k<N; ++k) { out[k*ostep+n*ostep2] = std::conj(buf[k+n*nb*2]) * ω₀[k]; }
  DFT_PROPERTIES
    std::size_t nb;
    std::vector<complex> ω₀, ω₁;
  DFT_CONSTRUCT(Bluestein, (std::size_t N, std::size_t Nb), DFTbase{N,true}, nb{Nb}, ω₀(N), ω₁(nb))
    assert(nb >= 2*N-1);
    for(std::size_t k=0; k<N; ++k) { ω₀[k] = w(k*k, 2*N); }
    for(std::size_t k=0; k<N; ++k) { ω₁[k] = ω₀[k] / float(nb); }
    for(std::size_t k=1; k<N; ++k) { ω₁[nb-k] = ω₁[k]; }
    FindDFT(nb, 1,1, 1,1, 1)->xform_inplace(&ω₁[0],1);
END_DFT_METHOD;

DFTbase* FindDFT(std::size_t N, std::size_t istep,std::size_t istep2,
                                std::size_t ostep,std::size_t ostep2,
                                std::size_t num)
{
    static std::pair<std::mutex,std::map<std::size_t, std::unique_ptr<DFTbase>>> data;
    auto& [lock,table] = data;
    if(auto i = table.find(N); i != table.end()) return &*i->second;
    // Perform a second search under exclusive access
    std::unique_lock lk(lock);
    if(auto i = table.find(N); i != table.end()) return &*i->second;
    lk.unlock();

    DFTbase* res = [=]->DFTbase*{
        if(N <= 3 || N == 7) return new DFT_DFT(N);
        if(N>16 && N%16==0) return new DFT_RadixP<16>(N);
        if(N>4 && N%4==0) return new DFT_RadixP<4>(N);
        if(N>2 && N%2==0) return new DFT_RadixP<2>(N);
        if(N>3 && N%3==0) return new DFT_RadixP<3>(N);
        if(N>5 && N%5==0) return new DFT_RadixP<5>(N);
        if(N>7 && N%7==0) return new DFT_RadixP<7>(N);
        if(N>13 && N%13==0) return new DFT_RadixP<13>(N);
        if(N>17 && N%17==0) return new DFT_RadixP<17>(N);
        // Round N up to next power of two
        std::size_t nb = 2*N-1;
        for(unsigned a=1; a< 8*sizeof(nb); a<<=1) { nb = 1 + ((nb-1) | ((nb-1) >> a)); }
        // Prime, and unwieldly for Bluestein? Use Rader.
        auto [factors, count] = get_factors_unique(N);
        if(!count && nb > N*5/2) return new DFT_Rader(N);
        // Else fallback to Bluestein.
        return new DFT_Bluestein(N, nb);
    }();
    lk.lock();
    return &*(table[N] = std::unique_ptr<DFTbase>(res));
}