// Streaming log-mel frontend - C port of phoneme_engine/features.py.
// Frame t is centered at sample t*HOP (torchaudio center=True), so the
// stream carries WIN/2 = 200 samples (12.5 ms) of lookahead latency.
#include "pww_frontend.h"

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

#include "frontend_data.h"

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

// ---- 512-point iterative radix-2 complex FFT (input real) -------------
#define NFFT PWW_FE_NFFT

static float fft_re[NFFT], fft_im[NFFT];
static float tw_cos[NFFT / 2], tw_sin[NFFT / 2];
static int fft_init_done = 0;

static void fft_init(void) {
    for (int i = 0; i < NFFT / 2; i++) {
        tw_cos[i] = cosf(-2.0f * (float)M_PI * i / NFFT);
        tw_sin[i] = sinf(-2.0f * (float)M_PI * i / NFFT);
    }
    fft_init_done = 1;
}

static void fft512(void) {
    // bit-reversal permutation
    for (int i = 1, j = 0; i < NFFT; i++) {
        int bit = NFFT >> 1;
        for (; j & bit; bit >>= 1) j ^= bit;
        j ^= bit;
        if (i < j) {
            float tr = fft_re[i]; fft_re[i] = fft_re[j]; fft_re[j] = tr;
            float ti = fft_im[i]; fft_im[i] = fft_im[j]; fft_im[j] = ti;
        }
    }
    for (int len = 2; len <= NFFT; len <<= 1) {
        int half = len >> 1, step = NFFT / len;
        for (int i = 0; i < NFFT; i += len) {
            for (int k = 0; k < half; k++) {
                float wr = tw_cos[k * step], wi = tw_sin[k * step];
                int a = i + k, b = i + k + half;
                float xr = fft_re[b] * wr - fft_im[b] * wi;
                float xi = fft_re[b] * wi + fft_im[b] * wr;
                fft_re[b] = fft_re[a] - xr;
                fft_im[b] = fft_im[a] - xi;
                fft_re[a] += xr;
                fft_im[a] += xi;
            }
        }
    }
}

// ---- streaming state ---------------------------------------------------
// ring of raw samples; enough for one centered window plus slack
#define SBUF (PWW_FE_WIN + 4 * PWW_FE_HOP)

struct pww_frontend {
    float samples[SBUF];
    int n_samples;        // total pushed
    int64_t next_frame;   // next frame index to emit
    float ema[PWW_FE_NMELS];
    int ema_init;
};

static struct pww_frontend g_fe;

void pww_frontend_reset(void) {
    memset(&g_fe, 0, sizeof(g_fe));
    if (!fft_init_done) fft_init();
}

// Push samples; calls emit(mel_frame) for every completed 10 ms frame.
void pww_frontend_push(const float *x, int n,
                       void (*emit)(const float *mel, void *user),
                       void *user) {
    struct pww_frontend *fe = &g_fe;
    for (int i = 0; i < n; i++) {
        fe->samples[fe->n_samples % SBUF] = x[i];
        fe->n_samples++;
        // frame f is ready when its centered window [f*HOP-200, f*HOP+200)
        // is fully available
        long long center = fe->next_frame * PWW_FE_HOP;
        while (center + PWW_FE_WIN / 2 <= fe->n_samples) {
            float mel[PWW_FE_NMELS];
            long long start = center - PWW_FE_WIN / 2;
            memset(fft_re, 0, sizeof(fft_re));
            memset(fft_im, 0, sizeof(fft_im));
            int off = (NFFT - PWW_FE_WIN) / 2;   // window centered in FFT
            for (int w = 0; w < PWW_FE_WIN; w++) {
                long long s = start + w;
                // reflect padding at the stream start (center=True)
                if (s < 0) s = -s;
                float v = (s < fe->n_samples)
                              ? fe->samples[s % SBUF] : 0.f;
                fft_re[off + w] = v * PWW_FE_HANN[w];
            }
            fft512();
            for (int m = 0; m < PWW_FE_NMELS; m++) mel[m] = 0.f;
            for (int b = 0; b < PWW_FE_NFREQS; b++) {
                float p = fft_re[b] * fft_re[b] + fft_im[b] * fft_im[b];
                const float *fbrow = PWW_FE_MELFB + b;
                for (int m = 0; m < PWW_FE_NMELS; m++)
                    mel[m] += p * fbrow[(size_t)m * PWW_FE_NFREQS];
            }
            for (int m = 0; m < PWW_FE_NMELS; m++) {
                float lm = logf(mel[m] + 1e-6f);
                // causal EMA mean subtraction (features.py lfilter)
                float e = fe->ema_init
                              ? (1.f - PWW_FE_EMA_ALPHA) * fe->ema[m]
                                    + PWW_FE_EMA_ALPHA * lm
                              : PWW_FE_EMA_ALPHA * lm;
                fe->ema[m] = e;
                mel[m] = lm - e;
            }
            fe->ema_init = 1;
            emit(mel, user);
            fe->next_frame++;
            center = fe->next_frame * PWW_FE_HOP;
        }
    }
}