// C port of KeywordSpotter (see phoneme_engine/decoder.py for the
// reference semantics and design rationale).
#include "pww_decoder.h"

#include <string.h>

#include "model_int8.h"

#define NEG_INF (-1e30f)
// 200 ms/phone cap: looser caps let garbage alignments crawl across
// continuous background speech (see decoder.py)
#define MAX_FRAMES_PER_PHONE 10

// Confusable sets, mirroring decoder.py CONFUSABLE. Ids are 1-based
// (0 = CTC blank), matching PWW_PHONES order in model_int8.h:
// AA=1 AE=2 AH=3 AO=4 AW=5 AY=6 B=7 CH=8 D=9 DH=10 EH=11 ER=12 EY=13
// F=14 G=15 HH=16 IH=17 IY=18 JH=19 K=20 L=21 M=22 N=23 NG=24 OW=25
// OY=26 P=27 R=28 S=29 SH=30 T=31 TH=32 UH=33 UW=34 V=35 W=36 Y=37
// Z=38 ZH=39
static int confusable(uint8_t id, uint8_t out[PWW_DEC_MAX_ALLOWED]) {
    switch (id) {
        case 3:  out[0] = 3;  out[1] = 17; out[2] = 12; return 3; // AH
        case 17: out[0] = 17; out[1] = 3;  return 2;              // IH
        case 12: out[0] = 12; out[1] = 3;  return 2;              // ER
        case 4:  out[0] = 4;  out[1] = 1;  return 2;              // AO
        case 2:  out[0] = 2;  out[1] = 1;  return 2;              // AE
        case 33: out[0] = 33; out[1] = 34; return 2;              // UH
        case 20: out[0] = 20; out[1] = 15; return 2;              // K
        case 15: out[0] = 15; out[1] = 20; return 2;              // G
        case 31: out[0] = 31; out[1] = 9;  return 2;              // T
        case 9:  out[0] = 9;  out[1] = 31; return 2;              // D
        default: out[0] = id; return 1;
    }
}

int pww_spotter_init(pww_spotter_t *sp, const uint8_t *phone_ids,
                     int n_phones, float threshold) {
    if (n_phones < 2 || n_phones > PWW_DEC_MAX_PHONES) return -1;
    memset(sp, 0, sizeof(*sp));
    sp->n_phones = n_phones;
    sp->threshold = threshold;
    sp->strong_margin = 1.5f;
    sp->strong_ratio = 0.5f;
    sp->refractory = 50;

    int n = 0;
    // leading blank
    sp->labels[n] = 0;
    sp->preds[n][0] = 0; sp->preds[n][1] = -1;
    sp->entry[n] = 1;
    n++;
    for (int i = 0; i < n_phones; i++) {
        uint8_t pid = phone_ids[i];
        int a = n;
        if (i == 0) {
            sp->labels[n] = pid;
            sp->preds[n][0] = 0; sp->preds[n][1] = -1;
            sp->entry[n] = 1;
            n++;
        } else {
            sp->labels[n] = pid;
            sp->preds[n][0] = (int8_t)(a - 1);
            // skip the blank between different phones
            sp->preds[n][1] = (sp->labels[a - 2] != pid)
                                  ? (int8_t)(a - 2) : -1;
            sp->entry[n] = 0;
            n++;
        }
        sp->labels[n] = pid;                     // B state, only from A
        sp->preds[n][0] = (int8_t)a; sp->preds[n][1] = -1;
        sp->entry[n] = 0;
        n++;
        sp->labels[n] = 0;                       // blank after phone
        sp->preds[n][0] = (int8_t)(a + 1); sp->preds[n][1] = -1;
        sp->entry[n] = 0;
        n++;
    }
    sp->n_states = n;
    for (int s = 0; s < n; s++)
        sp->n_allowed[s] = (sp->labels[s] == 0)
            ? (sp->allowed[s][0] = 0, 1)
            : (uint8_t)confusable(sp->labels[s], sp->allowed[s]);
    sp->min_dur = 3 * n_phones;
    sp->max_dur = MAX_FRAMES_PER_PHONE * n_phones;
    pww_spotter_reset(sp);
    return 0;
}

void pww_spotter_reset(pww_spotter_t *sp) {
    for (int s = 0; s < sp->n_states; s++) {
        sp->rel[s] = NEG_INF;
        sp->start[s] = 0;
        sp->pframes[s] = 0;
        sp->strong[s] = 0;
    }
    sp->cooldown = 0;
    sp->t = 0;
    sp->best_seen = NEG_INF;
}

int pww_spotter_step(pww_spotter_t *sp, const float *logits,
                     float *score_out, int *dur_out) {
    float filler = logits[0];
    for (int c = 1; c < PWW_NUM_CLASSES; c++)
        if (logits[c] > filler) filler = logits[c];

    float rel[PWW_DEC_MAX_STATES];
    int32_t start[PWW_DEC_MAX_STATES], pf[PWW_DEC_MAX_STATES],
        strong[PWW_DEC_MAX_STATES];

    for (int s = 0; s < sp->n_states; s++) {
        float best = sp->rel[s];
        int32_t b_start = sp->start[s], b_pf = sp->pframes[s],
                b_sf = sp->strong[s];
        for (int q = 0; q < 2; q++) {
            int8_t p = sp->preds[s][q];
            if (p >= 0 && sp->rel[p] > best) {
                best = sp->rel[p];
                b_start = sp->start[p];
                b_pf = sp->pframes[p];
                b_sf = sp->strong[p];
            }
        }
        if (sp->entry[s] && 0.0f >= best) {
            best = 0.0f; b_start = sp->t; b_pf = 0; b_sf = 0;
        }
        float emit = logits[sp->allowed[s][0]];
        for (int a = 1; a < sp->n_allowed[s]; a++) {
            float v = logits[sp->allowed[s][a]];
            if (v > emit) emit = v;
        }
        int is_phone = sp->labels[s] != 0;
        rel[s] = best + emit - filler;
        start[s] = b_start;
        pf[s] = b_pf + (is_phone ? 1 : 0);
        strong[s] = b_sf +
            ((is_phone && emit >= filler - sp->strong_margin) ? 1 : 0);
    }
    memcpy(sp->rel, rel, sizeof(float) * sp->n_states);
    memcpy(sp->start, start, sizeof(int32_t) * sp->n_states);
    memcpy(sp->pframes, pf, sizeof(int32_t) * sp->n_states);
    memcpy(sp->strong, strong, sizeof(int32_t) * sp->n_states);
    sp->t++;

    float norm_best = NEG_INF;
    int norm_dur = 0, have = 0;
    for (int f = 0; f < 2; f++) {
        int s = sp->n_states - 1 - f;
        int32_t dur = sp->t - sp->start[s];
        int32_t p = sp->pframes[s];
        if (dur < sp->min_dur || dur > sp->max_dur) continue;
        if (p < 2 * sp->n_phones) continue;
        if ((float)sp->strong[s] < sp->strong_ratio * (float)p) continue;
        float norm = sp->rel[s] / (float)p;
        if (!have || norm > norm_best) {
            norm_best = norm; norm_dur = (int)dur; have = 1;
        }
    }
    if (have && norm_best > sp->best_seen) sp->best_seen = norm_best;
    if (sp->cooldown > 0) { sp->cooldown--; return 0; }
    if (have && norm_best > sp->threshold) {
        sp->cooldown = sp->refractory;
        for (int s = 0; s < sp->n_states; s++) sp->rel[s] = NEG_INF;
        *score_out = norm_best;
        *dur_out = norm_dur;
        return 1;
    }
    return 0;
}