// learn / adapt / selfcheck / test / background / autothresh.
// A session is started by a console task, fed by the processing task via
// calibration_on_frame(), and finished (mode -> NORMAL) by the processing task, which
// then signals the waiting console task. Session data is guarded by s_mx.
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"
#include "calibration.h"
#include "config_store.h"
#include "fingerprint.h"
#include "matcher.h"

#define LEARN_SECS 60
#define TEST_SECS 10
#define BG_SECS 60
#define SECS_TO_FRAMES(s) ((uint32_t)(s) * FS_HZ / HOP_LEN)
#define NOISE_FRAMES (FS_HZ / HOP_LEN)          // first second -> noise floor
#define MIN_LEARN_FRAMES 20
#define MAX_SCORES 2048
_Static_assert(BG_SECS * FS_HZ / HOP_LEN <= MAX_SCORES, "score buffer too small for background run");


static SemaphoreHandle_t s_mx;
static SemaphoreHandle_t s_done;
static volatile app_mode_t s_mode = MODE_NORMAL;

static struct {
    uint32_t frame_idx, total_frames;
    float min_rms, threshold, noise_mult;
    // learn / adapt
    float noise[NOISE_FRAMES];
    float gate;
    float fp_sum[NUM_BANDS];
    uint32_t count;
    // test / background
    float scores[MAX_SCORES];
    uint32_t n_scores;
    uint32_t n_gated, n_above, n_above_gated;
} S;

static score_stats_t s_bg, s_buzz;

static int cmp_f(const void *a, const void *b)
{
    float x = *(const float *)a, y = *(const float *)b;
    return (x > y) - (x < y);
}

static float pct(const float *sorted, uint32_t n, float p)
{
    // Linear interpolation, same as numpy.percentile default.
    float pos = p / 100.0f * (float)(n - 1);
    uint32_t lo = (uint32_t)pos;
    uint32_t hi = lo + 1 < n ? lo + 1 : lo;
    float fr = pos - (float)lo;
    return sorted[lo] * (1.0f - fr) + sorted[hi] * fr;
}

static void compute_stats(float *a, uint32_t n, score_stats_t *o)
{
    memset(o, 0, sizeof(*o));
    if (n == 0) return;
    qsort(a, n, sizeof(float), cmp_f);
    double sum = 0;
    for (uint32_t i = 0; i < n; i++) sum += a[i];
    o->valid = true;
    o->n = n;
    o->min = a[0];
    o->max = a[n - 1];
    o->mean = (float)(sum / n);
    o->p5 = pct(a, n, 5);
    o->p50 = pct(a, n, 50);
    o->p95 = pct(a, n, 95);
    o->p99 = pct(a, n, 99);
}

static void print_stats(const char *label, const score_stats_t *s)
{
    printf("%s (n=%lu): min=%.3f p5=%.3f p50=%.3f mean=%.3f p95=%.3f p99=%.3f max=%.3f\n",
           label, (unsigned long)s->n, s->min, s->p5, s->p50, s->mean, s->p95, s->p99, s->max);
}

esp_err_t calibration_init(void)
{
    s_mx = xSemaphoreCreateMutex();
    s_done = xSemaphoreCreateBinary();
    return (s_mx && s_done) ? ESP_OK : ESP_ERR_NO_MEM;
}

app_mode_t calibration_get_mode(void) { return s_mode; }

const char *calibration_mode_name(app_mode_t m)
{
    switch (m) {
    case MODE_NORMAL: return "NORMAL";
    case MODE_LEARNING: return "LEARNING";
    case MODE_ADAPTING: return "ADAPTING";
    case MODE_BACKGROUND_CAL: return "BACKGROUND_CAL";
    case MODE_TEST_CAL: return "TEST_CAL";
    }
    return "?";
}

bool calibration_get_background(score_stats_t *o) { *o = s_bg; return s_bg.valid; }
bool calibration_get_buzz(score_stats_t *o) { *o = s_buzz; return s_buzz.valid; }

// ---- frame feed (processing task) -------------------------------------------

void calibration_on_frame(const float fp[NUM_BANDS], bool fp_valid, float rms, float score)
{
    xSemaphoreTake(s_mx, portMAX_DELAY);
    app_mode_t m = s_mode;
    if (m == MODE_NORMAL) {
        xSemaphoreGive(s_mx);
        return;
    }
    uint32_t i = S.frame_idx++;

    switch (m) {
    case MODE_LEARNING:
    case MODE_ADAPTING:
        if (i < NOISE_FRAMES) {
            S.noise[i] = rms;
            if (i == NOISE_FRAMES - 1) {
                float tmp[NOISE_FRAMES];
                memcpy(tmp, S.noise, sizeof(tmp));
                qsort(tmp, NOISE_FRAMES, sizeof(float), cmp_f);
                float floor_rms = tmp[NOISE_FRAMES / 2];            // median
                float g = floor_rms * S.noise_mult;
                S.gate = g > S.min_rms ? g : S.min_rms;             // never below the detector gate
            }
        } else if (fp_valid && rms >= S.gate) {
            for (int b = 0; b < NUM_BANDS; b++) S.fp_sum[b] += fp[b];
            S.count++;
        }
        break;
    case MODE_TEST_CAL:
        if (fp_valid && rms >= S.min_rms && S.n_scores < MAX_SCORES) {
            S.scores[S.n_scores++] = score;
        }
        break;
    case MODE_BACKGROUND_CAL:
        if (S.n_scores < MAX_SCORES) S.scores[S.n_scores++] = fp_valid ? score : 0.0f;
        if (rms >= S.min_rms) S.n_gated++;
        if (score >= S.threshold) {
            S.n_above++;
            if (rms >= S.min_rms) S.n_above_gated++;
        }
        break;
    default:
        break;
    }

    if (S.frame_idx >= S.total_frames) {
        s_mode = MODE_NORMAL;
        xSemaphoreGive(s_done);
    }
    xSemaphoreGive(s_mx);
}

// ---- session control (console / button tasks) --------------------------------

static esp_err_t begin(app_mode_t m, uint32_t secs)
{
    app_config_t cfg;
    config_get(&cfg);
    xSemaphoreTake(s_mx, portMAX_DELAY);
    if (s_mode != MODE_NORMAL) {
        xSemaphoreGive(s_mx);
        printf("busy: a %s session is already running\n", calibration_mode_name(s_mode));
        return ESP_ERR_INVALID_STATE;
    }
    memset(&S, 0, sizeof(S));
    S.total_frames = SECS_TO_FRAMES(secs);
    S.min_rms = cfg.min_rms;
    S.threshold = cfg.threshold;
    S.noise_mult = cfg.noise_mult;
    xSemaphoreTake(s_done, 0);
    s_mode = m;                 // last: processing task starts feeding after this
    xSemaphoreGive(s_mx);
    return ESP_OK;
}

static esp_err_t wait_done(uint32_t secs)
{
    for (uint32_t t = 1; t <= secs + 10; t++) {
        if (xSemaphoreTake(s_done, pdMS_TO_TICKS(1000)) == pdTRUE) {
            printf("\n");
            return ESP_OK;
        }
        printf(" %lus", (unsigned long)t);
        fflush(stdout);
    }
    xSemaphoreTake(s_mx, portMAX_DELAY);
    s_mode = MODE_NORMAL;
    xSemaphoreGive(s_mx);
    printf("\nsession timed out: no audio frames arriving?\n");
    return ESP_ERR_TIMEOUT;
}

esp_err_t calibration_run_learn(bool adapt, float w_new)
{
    if (adapt && !matcher_has_template()) {
        printf("adapt needs an existing template; use `learn` first\n");
        return ESP_ERR_INVALID_STATE;
    }
    if (adapt && !(w_new >= 0.0f && w_new <= 1.0f)) {
        printf("adapt weight must be 0..1\n");
        return ESP_ERR_INVALID_ARG;
    }
    esp_err_t err = begin(adapt ? MODE_ADAPTING : MODE_LEARNING, LEARN_SECS);
    if (err != ESP_OK) return err;
    printf("%s: %d s. Stay QUIET for the first second, then press the buzzer a few times.\n",
           adapt ? "ADAPT" : "LEARN", LEARN_SECS);
    err = wait_done(LEARN_SECS);
    if (err != ESP_OK) return err;

    xSemaphoreTake(s_mx, portMAX_DELAY);
    uint32_t count = S.count;
    float gate = S.gate;
    float v[NUM_BANDS];
    for (int b = 0; b < NUM_BANDS; b++) v[b] = count ? S.fp_sum[b] / (float)count : 0.0f;
    xSemaphoreGive(s_mx);

    printf("noise-gate RMS %.1f, loud frames collected: %lu\n", gate, (unsigned long)count);
    if (count < MIN_LEARN_FRAMES) {
        printf("REFUSED: need >= %d loud frames. Press the buzzer longer / closer, or lower noise_mult/min_rms.\n",
               MIN_LEARN_FRAMES);
        return ESP_FAIL;
    }
    // |mean of unit vectors| = average cosine of the frames to their own mean direction.
    float s2 = 0;
    for (int b = 0; b < NUM_BANDS; b++) s2 += v[b] * v[b];
    printf("cohesion of collected frames: %.3f (close to 1 = consistent sound)\n", sqrtf(s2));
    if (!features_l2_normalize(v, NUM_BANDS)) {
        printf("REFUSED: collected fingerprint is degenerate (flat)\n");
        return ESP_FAIL;
    }
    if (adapt) {
        float cur[NUM_BANDS];
        matcher_get_template(cur);
        printf("similarity new vs current template: %.3f\n", features_dot(v, cur));
        for (int b = 0; b < NUM_BANDS; b++) v[b] = w_new * v[b] + (1.0f - w_new) * cur[b];
        if (!features_l2_normalize(v, NUM_BANDS)) {
            printf("REFUSED: blended template degenerate\n");
            return ESP_FAIL;
        }
    }
    err = matcher_set_template(v, count);
    if (err != ESP_OK) {
        printf("saving template failed: %s\n", esp_err_to_name(err));
        return err;
    }
    printf("%s done: template saved to NVS (source now: %s)%s\n", adapt ? "ADAPT" : "LEARN",
           matcher_source_name(), adapt ? "" : ". Run `selfcheck` to verify.");
    return ESP_OK;
}

esp_err_t calibration_run_test(bool verdict)
{
    if (!matcher_has_template()) {
        printf("no template: run make_template.py or `learn` first\n");
        return ESP_ERR_INVALID_STATE;
    }
    esp_err_t err = begin(MODE_TEST_CAL, TEST_SECS);
    if (err != ESP_OK) return err;
    printf("%s: %d s. Press the real buzzer now (several presses).\n", verdict ? "SELFCHECK" : "TEST", TEST_SECS);
    err = wait_done(TEST_SECS);
    if (err != ESP_OK) return err;

    xSemaphoreTake(s_mx, portMAX_DELAY);
    uint32_t n = S.n_scores;
    float thr = S.threshold, gate = S.min_rms;
    score_stats_t st;
    compute_stats(S.scores, n, &st);
    xSemaphoreGive(s_mx);

    printf("loud frames (RMS >= %.0f): %lu\n", gate, (unsigned long)n);
    if (n < MIN_LEARN_FRAMES) {
        printf("WARNING: only %lu loud frames; too few for reliable stats (press longer or lower min_rms).\n",
               (unsigned long)n);
        if (n == 0) return ESP_FAIL;
    } else {
        s_buzz = st;      // used by autothresh
    }
    print_stats("score vs active template", &st);
    if (verdict && n >= MIN_LEARN_FRAMES) {
        if (st.p5 >= thr) {
            printf("VERDICT: GOOD match. 95%% of loud frames clear threshold %.2f.\n", thr);
        } else if (st.p50 >= thr) {
            printf("VERDICT: WEAK match. The median frame clears %.2f but the tail does not.\n"
                   "         Try `adapt` (blend with live mic) or `learn`, then `selfcheck` again.\n", thr);
        } else {
            printf("VERDICT: POOR match (median %.2f < threshold %.2f).\n"
                   "         Was that really the buzzer? If yes: `learn` on this mic, or see README about FP_DETREND\n"
                   "         (compile-time, feature_params.h) if mic tilt is the issue.\n", st.p50, thr);
        }
    }
    return ESP_OK;
}

esp_err_t calibration_run_background(void)
{
    if (!matcher_has_template()) {
        printf("no template: run make_template.py or `learn` first\n");
        return ESP_ERR_INVALID_STATE;
    }
    esp_err_t err = begin(MODE_BACKGROUND_CAL, BG_SECS);
    if (err != ESP_OK) return err;
    printf("BACKGROUND: %d s of normal ambient sound (do NOT press the buzzer).\n", BG_SECS);
    err = wait_done(BG_SECS);
    if (err != ESP_OK) return err;

    xSemaphoreTake(s_mx, portMAX_DELAY);
    uint32_t n = S.n_scores, gated = S.n_gated, above = S.n_above, above_g = S.n_above_gated;
    float thr = S.threshold;
    compute_stats(S.scores, n, &s_bg);
    xSemaphoreGive(s_mx);

    print_stats("background score vs active template", &s_bg);
    printf("frames passing the loudness gate: %lu / %lu\n", (unsigned long)gated, (unsigned long)n);
    printf("frames with score >= threshold %.2f: %lu (of which also loud enough to match: %lu)\n",
           thr, (unsigned long)above, (unsigned long)above_g);
    return ESP_OK;
}

esp_err_t calibration_autothresh(void)
{
    if (!s_bg.valid || !s_buzz.valid) {
        printf("need both: `background` and (`test` or `selfcheck`) with >= %d loud frames first\n", MIN_LEARN_FRAMES);
        return ESP_ERR_INVALID_STATE;
    }
    float lo = s_bg.p99, hi = s_buzz.p5;
    float thr = 0.5f * (lo + hi);
    if (thr < 0.50f) thr = 0.50f;
    if (thr > 0.98f) thr = 0.98f;
    printf("background p99 = %.3f, buzzer p5 = %.3f -> threshold %.3f\n", lo, hi, thr);
    if (hi <= lo) {
        printf("WARNING: distributions OVERLAP. Expect false alarms or missed buzzes. "
               "Improve the template (`learn`), move the mic, or revisit FP_DETREND.\n");
    }
    char msg[96], val[24];
    snprintf(val, sizeof(val), "%.4f", thr);
    if (config_set_param("threshold", val, msg, sizeof(msg)) != ESP_OK) {
        printf("could not set threshold: %s\n", msg);
        return ESP_FAIL;
    }
    printf("threshold set and saved.\n");
    return ESP_OK;
}