// Spectral fingerprint. MUST stay numerically identical to tools/make_template.py.
#include <math.h>
#include <string.h>
#include "esp_log.h"
#include "dsps_fft2r.h"
#include "fingerprint.h"

static const char *TAG = "features";

static float s_window[FRAME_LEN];
static int s_edges[NUM_BANDS + 1];
static float s_inv_width[NUM_BANDS];
static float s_sxx;                 // sum of squared centred band indices (detrend)
static bool s_ready;

esp_err_t features_init(void)
{
    if (s_ready) {
        return ESP_OK;
    }
    // Periodic Hann: w[n] = 0.5 - 0.5 cos(2 pi n / N)
    for (int n = 0; n < FRAME_LEN; n++) {
        s_window[n] = 0.5f - 0.5f * cosf(2.0f * (float)M_PI * (float)n / (float)FRAME_LEN);
    }

    // Log-spaced band edges, see feature_params.h. Double math + floor(x+0.5) so
    // C and numpy can never disagree on rounding of exact .5 values.
    const double ratio = (double)F_MAX_HZ / (double)F_MIN_HZ;
    for (int b = 0; b <= NUM_BANDS; b++) {
        double f = (double)F_MIN_HZ * pow(ratio, (double)b / (double)NUM_BANDS);
        int bin = (int)floor(f * (double)FRAME_LEN / (double)FS_HZ + 0.5);
        if (b > 0 && bin < s_edges[b - 1] + 1) {
            bin = s_edges[b - 1] + 1;       // guarantee >= 1 bin per band
        }
        s_edges[b] = bin;
    }
    if (s_edges[NUM_BANDS] > FRAME_LEN / 2) {
        ESP_LOGE(TAG, "band edges exceed Nyquist (%d > %d): lower NUM_BANDS or raise F_MIN_HZ",
                 s_edges[NUM_BANDS], FRAME_LEN / 2);
        return ESP_ERR_INVALID_ARG;
    }
    for (int b = 0; b < NUM_BANDS; b++) {
        s_inv_width[b] = 1.0f / (float)(s_edges[b + 1] - s_edges[b]);
    }
    s_sxx = 0.0f;
    for (int b = 0; b < NUM_BANDS; b++) {
        float xc = (float)b - 0.5f * (float)(NUM_BANDS - 1);
        s_sxx += xc * xc;
    }

    esp_err_t err = dsps_fft2r_init_fc32(NULL, FRAME_LEN);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "dsps_fft2r_init_fc32: %s", esp_err_to_name(err));
        return err;
    }
    s_ready = true;
    return ESP_OK;
}

const int *features_band_edges(void) { return s_edges; }

bool features_l2_normalize(float *v, int n)
{
    float s = 0.0f;
    for (int i = 0; i < n; i++) s += v[i] * v[i];
    float nrm = sqrtf(s);
    if (!(nrm > 1e-3f)) {               // flat or NaN (silence leaves ~1e-5 float residue)
        memset(v, 0, (size_t)n * sizeof(float));
        return false;
    }
    float inv = 1.0f / nrm;
    for (int i = 0; i < n; i++) v[i] *= inv;
    return true;
}

float features_dot(const float *a, const float *b)
{
    float s = 0.0f;
    for (int i = 0; i < NUM_BANDS; i++) s += a[i] * b[i];
    if (s > 1.0f) s = 1.0f;
    if (s < -1.0f) s = -1.0f;
    return s;
}

bool features_compute(features_ctx_t *ctx, const int16_t *pcm, float fp[NUM_BANDS], float *rms_out)
{
    // 1) DC removal per frame (exact integer sum, then float).
    int32_t isum = 0;
    for (int i = 0; i < FRAME_LEN; i++) isum += pcm[i];
    const float mean = (float)isum / (float)FRAME_LEN;

    // 2) RMS of the DC-free frame (int16 units) + windowed, +-1.0-scaled FFT input.
    float sq = 0.0f;
    for (int i = 0; i < FRAME_LEN; i++) {
        float d = (float)pcm[i] - mean;
        sq += d * d;
        ctx->fft[2 * i] = d * (1.0f / 32768.0f) * s_window[i];
        ctx->fft[2 * i + 1] = 0.0f;
    }
    const float rms = sqrtf(sq / (float)FRAME_LEN);
    if (rms_out) *rms_out = rms;

    // 3) FFT (decimation-in-frequency + bit reversal) -> bins 0..FRAME_LEN/2-1.
    if (dsps_fft2r_fc32(ctx->fft, FRAME_LEN) != ESP_OK ||
        dsps_bit_rev_fc32(ctx->fft, FRAME_LEN) != ESP_OK) {
        memset(fp, 0, NUM_BANDS * sizeof(float));
        return false;
    }

    // 4) Mean power per band -> natural log (eps guards log(0)).
    float mean_log = 0.0f;
    for (int b = 0; b < NUM_BANDS; b++) {
        float acc = 0.0f;
        for (int k = s_edges[b]; k < s_edges[b + 1]; k++) {
            float re = ctx->fft[2 * k], im = ctx->fft[2 * k + 1];
            acc += re * re + im * im;
        }
        fp[b] = logf(acc * s_inv_width[b] + LOG_EPS);
        mean_log += fp[b];
    }
    // 5) Subtract the mean across bands: removes overall level (gain, distance).
    mean_log /= (float)NUM_BANDS;
    for (int b = 0; b < NUM_BANDS; b++) fp[b] -= mean_log;

#if FP_DETREND
    // 6) Remove the least-squares line over band index. A mic's frequency-response
    // tilt looks like a linear slope in log-spectrum, so this cancels it. Vector is
    // already zero-mean, so the fit is slope = <v, xc> / <xc, xc>.
    float sxv = 0.0f;
    for (int b = 0; b < NUM_BANDS; b++) {
        sxv += fp[b] * ((float)b - 0.5f * (float)(NUM_BANDS - 1));
    }
    const float slope = sxv / s_sxx;
    for (int b = 0; b < NUM_BANDS; b++) {
        fp[b] -= slope * ((float)b - 0.5f * (float)(NUM_BANDS - 1));
    }
#endif

    // 7) L2 normalise so cosine similarity is a plain dot product.
    return features_l2_normalize(fp, NUM_BANDS);
}