// Part of JoeyLib, Copyright (c) 2026 Scott Duensing, under the MIT License.
// See LICENSE.md for license information.
// SPDX-License-Identifier: MIT
//
// gateFft.c - numpy's spectral arithmetic for the gate tools (gateFft.h).
//
// The transform is decimation in time over the length's prime factors, each
// stage's butterfly a direct DFT of its radix (two by sums and differences),
// with the roots of unity taken from the first octant by symmetry so that
// the quarter turns are exact. A length with a prime factor past
// FFT_RADIX_MAX goes through Bluestein's chirp instead: the transform as a
// convolution, done on a power-of-two transform at least twice as long.
// numpy's own transform (pocketfft) differs from this one in the last bits
// of each bin only; nothing a grade prints can tell them apart.
//
// numpy's sum is not a running total: it adds a block of up to 8192 values
// (its buffer) pairwise, and the blocks' sums in turn. gateSum does the
// same, so means and levels come out to the bit.

#define _DEFAULT_SOURCE

#include <math.h>
#include <stdlib.h>
#include <string.h>

#include "gateFft.h"
#include "toolFile.h"

// The largest prime a stage takes directly; a length with a larger factor
// is transformed by Bluestein's chirp.
#define FFT_RADIX_MAX      61u
#define FFT_RADIX_TWO      2u
#define FFT_FIRST_ODD      3u
#define OCTANTS            8u
#define HALF_TURN_OCTANTS  4u
#define QUARTER_OCTANTS    2u
// Bluestein's chirp angles are k^2 over twice the length.
#define CHIRP_SPAN         2u
// numpy's sum: blocks of its buffer size, pairwise within them below a
// block of PAIRWISE_LEAF values unrolled eight ways.
#define SUM_BLOCK          8192u
#define PAIRWISE_LEAF      128u
#define PAIRWISE_WAYS      8u

static void   fftAlloc(double **re, double **im, size_t count);
static void   fftLoad(GateFftT *fft, const double *in, size_t inLen);
static void   fftStage(const GateFftT *fft, const double *inRe, const double *inIm, size_t inStride, double *outRe, double *outIm, size_t n, size_t factorAt);
static void   fftTransform(GateFftT *fft);
static double pairwiseSum(const double *x, size_t n);
static void   unitRoot(uint64_t k, uint64_t n, double *re, double *im);


// Two arrays of count doubles, zeroed.
static void fftAlloc(double **re, double **im, size_t count) {
    *re = (double *)calloc(count, sizeof(double));
    *im = (double *)calloc(count, sizeof(double));
    if (*re == NULL || *im == NULL) {
        toolFail("out of memory");
    }
}


// numpy.fft.rfft's input: in[0..inLen) zero-padded (or cut) to n real
// points, transformed into fft->out.
static void fftLoad(GateFftT *fft, const double *in, size_t inLen) {
    size_t count;

    count = inLen < fft->n ? inLen : fft->n;
    memcpy(fft->inRe, in, count * sizeof(double));
    memset(fft->inRe + count, 0, (fft->n - count) * sizeof(double));
    memset(fft->inIm, 0, fft->n * sizeof(double));
    fftTransform(fft);
}


// One decimation-in-time stage: the n points inStride apart from in, as
// factors[factorAt] interleaved transforms of n / p points (each by the
// stages after this one), then combined by radix-p butterflies into out.
static void fftStage(const GateFftT *fft, const double *inRe, const double *inIm, size_t inStride, double *outRe, double *outIm, size_t n, size_t factorAt) {
    size_t p;
    size_t m;
    size_t q;
    size_t k;
    size_t r;
    size_t rootStride;
    size_t radixStride;
    size_t root;
    double re;
    double im;
    double sumRe;
    double sumIm;

    p           = fft->factors[factorAt];
    m           = n / p;
    rootStride  = fft->n / n;
    radixStride = fft->n / p;
    for (q = 0u; q < p; q++) {
        if (m == 1u) {
            outRe[q] = inRe[q * inStride];
            outIm[q] = inIm[q * inStride];
        } else {
            fftStage(fft, inRe + q * inStride, inIm + q * inStride, inStride * p, outRe + q * m, outIm + q * m, m, factorAt + 1u);
        }
    }
    for (k = 0u; k < m; k++) {
        for (q = 0u; q < p; q++) {
            re = outRe[q * m + k];
            im = outIm[q * m + k];
            if (q == 0u || k == 0u) {
                fft->spokeRe[q] = re;
                fft->spokeIm[q] = im;
            } else {
                root            = q * k * rootStride;
                fft->spokeRe[q] = re * fft->rootRe[root] - im * fft->rootIm[root];
                fft->spokeIm[q] = re * fft->rootIm[root] + im * fft->rootRe[root];
            }
        }
        if (p == FFT_RADIX_TWO) {
            outRe[k]     = fft->spokeRe[0] + fft->spokeRe[1];
            outIm[k]     = fft->spokeIm[0] + fft->spokeIm[1];
            outRe[k + m] = fft->spokeRe[0] - fft->spokeRe[1];
            outIm[k + m] = fft->spokeIm[0] - fft->spokeIm[1];
            continue;
        }
        for (r = 0u; r < p; r++) {
            sumRe = fft->spokeRe[0];
            sumIm = fft->spokeIm[0];
            for (q = 1u; q < p; q++) {
                root   = ((q * r) % p) * radixStride;
                sumRe += fft->spokeRe[q] * fft->rootRe[root] - fft->spokeIm[q] * fft->rootIm[root];
                sumIm += fft->spokeRe[q] * fft->rootIm[root] + fft->spokeIm[q] * fft->rootRe[root];
            }
            outRe[r * m + k] = sumRe;
            outIm[r * m + k] = sumIm;
        }
    }
}


// fft->in into fft->out.
static void fftTransform(GateFftT *fft) {
    GateFftT *inner;
    size_t    k;
    size_t    m;
    double    re;
    double    im;

    if (fft->inner == NULL) {
        fftStage(fft, fft->inRe, fft->inIm, 1u, fft->outRe, fft->outIm, fft->n, 0u);
        return;
    }
    // Bluestein: X[j] = w[j] * sum(x[k] w[k] conj(w[j - k])), the sum a
    // convolution with the filter; the inverse transform is the forward
    // one of the conjugate, conjugated and scaled.
    inner = fft->inner;
    m     = inner->n;
    for (k = 0u; k < m; k++) {
        if (k < fft->n) {
            inner->inRe[k] = fft->inRe[k] * fft->chirpRe[k] - fft->inIm[k] * fft->chirpIm[k];
            inner->inIm[k] = fft->inRe[k] * fft->chirpIm[k] + fft->inIm[k] * fft->chirpRe[k];
        } else {
            inner->inRe[k] = 0.0;
            inner->inIm[k] = 0.0;
        }
    }
    fftTransform(inner);
    for (k = 0u; k < m; k++) {
        re             = inner->outRe[k] * fft->filterRe[k] - inner->outIm[k] * fft->filterIm[k];
        im             = inner->outRe[k] * fft->filterIm[k] + inner->outIm[k] * fft->filterRe[k];
        inner->inRe[k] = re;
        inner->inIm[k] = -im;
    }
    fftTransform(inner);
    for (k = 0u; k < fft->n; k++) {
        re            = inner->outRe[k] / (double)m;
        im            = -inner->outIm[k] / (double)m;
        fft->outRe[k] = re * fft->chirpRe[k] - im * fft->chirpIm[k];
        fft->outIm[k] = re * fft->chirpIm[k] + im * fft->chirpRe[k];
    }
}


void gateFftAbs(GateFftT *fft, const double *in, size_t inLen, double *magnitude) {
    size_t k;

    fftLoad(fft, in, inLen);
    for (k = 0u; k <= fft->n / 2u; k++) {
        magnitude[k] = hypot(fft->outRe[k], fft->outIm[k]);
    }
}


void gateFftFree(GateFftT *fft) {
    if (fft->inner != NULL) {
        gateFftFree(fft->inner);
        free(fft->inner);
    }
    free(fft->rootRe);
    free(fft->rootIm);
    free(fft->inRe);
    free(fft->inIm);
    free(fft->outRe);
    free(fft->outIm);
    free(fft->spokeRe);
    free(fft->spokeIm);
    free(fft->chirpRe);
    free(fft->chirpIm);
    free(fft->filterRe);
    free(fft->filterIm);
    memset(fft, 0, sizeof(*fft));
}


void gateFftInit(GateFftT *fft, size_t n) {
    GateFftT *inner;
    size_t    rest;
    size_t    p;
    size_t    k;
    size_t    m;
    size_t    largest;

    memset(fft, 0, sizeof(*fft));
    fft->n = n;
    rest   = n;
    // The prime factors, rising (one factor of 1 for a single point).
    while (rest > 1u && rest % FFT_RADIX_TWO == 0u) {
        fft->factors[fft->factorCount++] = FFT_RADIX_TWO;
        rest /= FFT_RADIX_TWO;
    }
    for (p = FFT_FIRST_ODD; p * p <= rest; p += FFT_RADIX_TWO) {
        while (rest % p == 0u) {
            fft->factors[fft->factorCount++] = p;
            rest /= p;
        }
    }
    if (rest > 1u || fft->factorCount == 0u) {
        fft->factors[fft->factorCount++] = rest;
    }
    largest = fft->factors[fft->factorCount - 1u];
    fftAlloc(&fft->inRe, &fft->inIm, n);
    fftAlloc(&fft->outRe, &fft->outIm, n);
    if (largest <= FFT_RADIX_MAX) {
        fftAlloc(&fft->rootRe, &fft->rootIm, n);
        fftAlloc(&fft->spokeRe, &fft->spokeIm, largest);
        for (k = 0u; k < n; k++) {
            unitRoot(k, n, &fft->rootRe[k], &fft->rootIm[k]);
        }
        return;
    }
    // Bluestein: the chirp w[k] = exp(-i pi k^2 / n), and the filter, the
    // transform of conj(w) laid out circularly over a power of two m of at
    // least 2n - 1 points.
    m = 1u;
    while (m < CHIRP_SPAN * n - 1u) {
        m *= FFT_RADIX_TWO;
    }
    inner = (GateFftT *)malloc(sizeof(GateFftT));
    if (inner == NULL) {
        toolFail("out of memory");
    }
    gateFftInit(inner, m);
    fft->inner = inner;
    fftAlloc(&fft->chirpRe, &fft->chirpIm, n);
    fftAlloc(&fft->filterRe, &fft->filterIm, m);
    for (k = 0u; k < n; k++) {
        unitRoot(((uint64_t)k * k) % ((uint64_t)CHIRP_SPAN * n), (uint64_t)CHIRP_SPAN * n, &fft->chirpRe[k], &fft->chirpIm[k]);
    }
    for (k = 0u; k < n; k++) {
        inner->inRe[k] = fft->chirpRe[k];
        inner->inIm[k] = -fft->chirpIm[k];
        if (k > 0u) {
            inner->inRe[m - k] = fft->chirpRe[k];
            inner->inIm[m - k] = -fft->chirpIm[k];
        }
    }
    fftTransform(inner);
    memcpy(fft->filterRe, inner->outRe, m * sizeof(double));
    memcpy(fft->filterIm, inner->outIm, m * sizeof(double));
}


void gateFftReal(GateFftT *fft, const double *in, size_t inLen, double *re, double *im) {
    fftLoad(fft, in, inLen);
    memcpy(re, fft->outRe, (fft->n / 2u + 1u) * sizeof(double));
    memcpy(im, fft->outIm, (fft->n / 2u + 1u) * sizeof(double));
}


void gateHanning(double *window, size_t n) {
    size_t k;
    double size;

    if (n == 1u) {
        window[0] = 1.0;
        return;
    }
    // 0.5 + 0.5 cos(pi m / (n - 1)) for m = 1 - n, 3 - n, ... n - 1.
    size = (double)n;
    for (k = 0u; k < n; k++) {
        window[k] = 0.5 + 0.5 * cos(M_PI * ((1.0 - size) + 2.0 * (double)k) / (size - 1.0));
    }
}


double gateMean(const double *x, size_t n) {
    return gateSum(x, n) / (double)n;
}


double gateRfftFreq(size_t bin, size_t n, uint32_t rate) {
    double spacing;

    spacing = 1.0 / ((double)n * (1.0 / (double)rate));
    return (double)bin * spacing;
}


double gateSum(const double *x, size_t n) {
    double total;
    size_t at;
    size_t count;

    total = 0.0;
    for (at = 0u; at < n; at += SUM_BLOCK) {
        count  = n - at < SUM_BLOCK ? n - at : SUM_BLOCK;
        total += pairwiseSum(x + at, count);
    }
    return total;
}


// numpy's pairwise_sum (loops_utils.h).
static double pairwiseSum(const double *x, size_t n) {
    double sums[PAIRWISE_WAYS];
    double total;
    size_t i;
    size_t j;
    size_t half;

    if (n < PAIRWISE_WAYS) {
        total = 0.0;
        for (i = 0u; i < n; i++) {
            total += x[i];
        }
        return total;
    }
    if (n <= PAIRWISE_LEAF) {
        for (j = 0u; j < PAIRWISE_WAYS; j++) {
            sums[j] = x[j];
        }
        for (i = PAIRWISE_WAYS; i < n - n % PAIRWISE_WAYS; i += PAIRWISE_WAYS) {
            for (j = 0u; j < PAIRWISE_WAYS; j++) {
                sums[j] += x[i + j];
            }
        }
        total = ((sums[0] + sums[1]) + (sums[2] + sums[3])) + ((sums[4] + sums[5]) + (sums[6] + sums[7]));
        for (; i < n; i++) {
            total += x[i];
        }
        return total;
    }
    half  = n / 2u;
    half -= half % PAIRWISE_WAYS;
    return pairwiseSum(x, half) + pairwiseSum(x + half, n - half);
}


// exp(-2 pi i k / n) for k < n, from the first octant: the angle k / n of a
// turn is reflected (as eighths: 8k against n) into [0, 1/8], so the
// quarter turns come out exact and both halves of a turn agree.
static void unitRoot(uint64_t k, uint64_t n, double *re, double *im) {
    uint64_t eighths;
    double   angle;
    double   c;
    double   s;
    double   swap;
    bool     negSin;
    bool     negCos;
    bool     swapped;

    eighths = k * OCTANTS;
    negSin  = false;
    negCos  = false;
    swapped = false;
    if (eighths > HALF_TURN_OCTANTS * n) {
        eighths = OCTANTS * n - eighths;
        negSin  = true;
    }
    if (eighths > QUARTER_OCTANTS * n) {
        eighths = HALF_TURN_OCTANTS * n - eighths;
        negCos  = true;
    }
    if (eighths > n) {
        eighths = QUARTER_OCTANTS * n - eighths;
        swapped = true;
    }
    angle = 2.0 * M_PI * (double)eighths / (double)(OCTANTS * n);
    c     = cos(angle);
    s     = sin(angle);
    if (swapped) {
        swap = c;
        c    = s;
        s    = swap;
    }
    if (negCos) {
        c = -c;
    }
    if (negSin) {
        s = -s;
    }
    *re = c;
    *im = -s;
}