// Streaming int8 phoneme engine. See pww_engine.h for the contract.
#include "pww_engine.h"

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

#include "model_int8.h"

// Debug: when >= 0, pww_engine_step returns layer N's snapped int8
// output column instead of the final logits.
int pww_dump_layer = -1;

// Ring buffer per layer holding the int8 input history a causal conv
// needs: (k-1)*dilation + 1 columns of in_c values.
typedef struct {
    int8_t *buf;      // hist * in_c, column-major by time step
    int hist;         // number of columns
    int pos;          // next write slot
    int primed;       // columns written so far (zeros before that)
} ring_t;

struct pww_engine {
    ring_t rings[PWW_NUM_LAYERS];
    // weights copied out of memory-mapped flash at create() time: flash
    // cache misses were 20x slower than the arithmetic
    pww_layer_t layers[PWW_NUM_LAYERS];
    // block input snapshot (int8 col + its scale) for residual adds
    int8_t block_in[512];
    float block_in_scale;
    float scratch_f[512];
    int8_t col_a[512], col_b[512];
};

// Weight allocation. On ESP32 the engine is memory-latency bound, not
// MAC bound: internal SRAM is several times faster to walk than octal
// PSRAM, so fill the internal budget first and spill the rest to PSRAM.
#ifdef ESP_PLATFORM
#include "esp_heap_caps.h"
#ifndef PWW_INTERNAL_WEIGHT_BUDGET
#define PWW_INTERNAL_WEIGHT_BUDGET (160 * 1024)
#endif
static size_t s_internal_used = 0;
static void *weights_alloc(size_t n) {
    if (s_internal_used + n <= PWW_INTERNAL_WEIGHT_BUDGET) {
        void *p = heap_caps_malloc(n, MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
        if (p) { s_internal_used += n; return p; }
    }
    void *p = heap_caps_malloc(n, MALLOC_CAP_SPIRAM);
    if (!p) p = malloc(n);
    return p;
}
size_t pww_internal_weight_bytes(void) { return s_internal_used; }
// engine state (ring buffers, scratch) is touched every frame - never
// let it land in PSRAM
static void *state_alloc(size_t n) {
    void *p = heap_caps_calloc(1, n, MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
    return p ? p : calloc(1, n);
}
#else
static void *weights_alloc(size_t n) { return malloc(n); }
static void *state_alloc(size_t n) { return calloc(1, n); }
#endif

static int8_t quant_clamp(float v, float inv_scale) {
    float q = roundf(v * inv_scale);
    if (q > 127.f) q = 127.f;
    if (q < -127.f) q = -127.f;
    return (int8_t)q;
}

pww_engine_t *pww_engine_create(void) {
    pww_engine_t *e = (pww_engine_t *)state_alloc(sizeof(pww_engine_t));
    if (!e) return NULL;
    for (int li = 0; li < PWW_NUM_LAYERS; li++) {
        e->layers[li] = PWW_LAYERS[li];
        pww_layer_t *M = &e->layers[li];
        size_t wn = (size_t)(M->is_dw ? M->out_c : M->in_c * M->out_c)
                    * M->k;
        int8_t *wcopy = (int8_t *)weights_alloc(wn);
        float *scopy = (float *)malloc(sizeof(float) * M->out_c);
        float *bcopy = (float *)malloc(sizeof(float) * M->out_c);
        if (!wcopy || !scopy || !bcopy) { pww_engine_destroy(e); return NULL; }
        memcpy(wcopy, M->w, wn);
        memcpy(scopy, M->s, sizeof(float) * M->out_c);
        memcpy(bcopy, M->b, sizeof(float) * M->out_c);
        M->w = wcopy; M->s = scopy; M->b = bcopy;
        const pww_layer_t *L = &PWW_LAYERS[li];
        int hist = (L->k - 1) * L->dilation + 1;
        // stride-2 stem consumes 2 input columns per step
        if (L->stride == 2) hist += 1;
        ring_t *r = &e->rings[li];
        r->hist = hist;
        r->buf = (int8_t *)state_alloc((size_t)hist * L->in_c);
        if (!r->buf) { pww_engine_destroy(e); return NULL; }
    }
    pww_engine_reset(e);
    return e;
}

void pww_engine_destroy(pww_engine_t *e) {
    if (!e) return;
    for (int li = 0; li < PWW_NUM_LAYERS; li++) {
        free(e->rings[li].buf);
        if (e->layers[li].w && e->layers[li].w != PWW_LAYERS[li].w) {
            free((void *)e->layers[li].w);
            free((void *)e->layers[li].s);
            free((void *)e->layers[li].b);
        }
    }
    free(e);
}

void pww_engine_reset(pww_engine_t *e) {
    for (int li = 0; li < PWW_NUM_LAYERS; li++) {
        ring_t *r = &e->rings[li];
        memset(r->buf, 0, (size_t)r->hist * PWW_LAYERS[li].in_c);
        r->pos = 0;
        r->primed = 0;
    }
}

static void ring_push(ring_t *r, const int8_t *col, int in_c) {
    memcpy(r->buf + (size_t)r->pos * in_c, col, (size_t)in_c);
    r->pos = (r->pos + 1) % r->hist;
    if (r->primed < r->hist) r->primed++;
}

// column at "delay" steps in the past (0 = newest)
static const int8_t *ring_at(const ring_t *r, int delay, int in_c) {
    int idx = r->pos - 1 - delay;
    while (idx < 0) idx += r->hist;
    return r->buf + (size_t)idx * in_c;
}

// Runs one layer on the newest ring content, writing float pre-snap
// output to out_f (out_c values). For stride-2 layers the newest TWO
// columns have been pushed before calling.
static void layer_forward(const pww_layer_t *L, const ring_t *r,
                          float *out_f) {
    int k = L->k, d = L->dilation, in_c = L->in_c, out_c = L->out_c;
    // stride-2 layers get two pushes per step but output frame t only
    // consumes up to input column 2t: the newest tap sits one column back
    int off = (L->stride == 2) ? 1 : 0;
    if (L->is_dw) {
        for (int c = 0; c < out_c; c++) out_f[c] = 0.f;
        for (int i = 0; i < k; i++) {
            // tap i is the newest at i == k-1
            const int8_t *col = ring_at(r, (k - 1 - i) * d + off, in_c);
            const int8_t *w = L->w + i;  // w layout: (ch, 1, k)
            for (int c = 0; c < out_c; c++)
                out_f[c] += (float)((int32_t)col[c] * (int32_t)w[c * k]);
        }
        for (int c = 0; c < out_c; c++)
            out_f[c] = out_f[c] * L->s[c] + L->b[c];
    } else if (k == 1) {
        // 1x1 conv = matrix-vector; 85% of all MACs live here.
        // Contiguous rows, 4-way unroll, single ring lookup.
        const int8_t *restrict col = ring_at(r, off, in_c);
        for (int c = 0; c < out_c; c++) {
            const int8_t *restrict w = L->w + (size_t)c * in_c;
            int32_t a0 = 0, a1 = 0, a2 = 0, a3 = 0;
            int j = 0;
            for (; j + 4 <= in_c; j += 4) {
                a0 += (int32_t)col[j] * (int32_t)w[j];
                a1 += (int32_t)col[j + 1] * (int32_t)w[j + 1];
                a2 += (int32_t)col[j + 2] * (int32_t)w[j + 2];
                a3 += (int32_t)col[j + 3] * (int32_t)w[j + 3];
            }
            int32_t acc = a0 + a1 + a2 + a3;
            for (; j < in_c; j++)
                acc += (int32_t)col[j] * (int32_t)w[j];
            out_f[c] = (float)acc * L->s[c] + L->b[c];
        }
    } else {
        // general conv: hoist the k column pointers out of the c loop
        const int8_t *cols[8];
        for (int i = 0; i < k; i++)
            cols[i] = ring_at(r, (k - 1 - i) * d + off, in_c);
        for (int c = 0; c < out_c; c++) {
            const int8_t *restrict w = L->w + (size_t)c * in_c * k;
            int32_t acc = 0;
            for (int i = 0; i < k; i++) {
                const int8_t *restrict col = cols[i];
                const int8_t *restrict wk = w + i;
                for (int j = 0; j < in_c; j++)
                    acc += (int32_t)col[j] * (int32_t)wk[(size_t)j * k];
            }
            out_f[c] = (float)acc * L->s[c] + L->b[c];
        }
    }
}

int pww_engine_step(pww_engine_t *e, const float *mel0, const float *mel1,
                    float *logits_out) {
    // quantize the two input mel columns to the input scale
    float inv_in = 1.0f / PWW_INPUT_SCALE;
    for (int j = 0; j < PWW_MELS; j++)
        e->col_a[j] = quant_clamp(mel0[j], inv_in);
    ring_push(&e->rings[0], e->col_a, PWW_MELS);
    for (int j = 0; j < PWW_MELS; j++)
        e->col_a[j] = quant_clamp(mel1[j], inv_in);
    ring_push(&e->rings[0], e->col_a, PWW_MELS);

    int8_t *cur = e->col_a;   // int8 column flowing between layers
    float cur_scale = PWW_INPUT_SCALE;
    (void)cur_scale;

    for (int li = 0; li < PWW_NUM_LAYERS; li++) {
        const pww_layer_t *L = &e->layers[li];
        ring_t *r = &e->rings[li];
        if (li > 0) ring_push(r, cur, L->in_c);

        if (L->is_dw) {
            // save the block input column + scale for the residual 2
            // layers later (dw -> pw(residual))
            memcpy(e->block_in, ring_at(r, 0, L->in_c), (size_t)L->in_c);
            e->block_in_scale =
                (li == 0) ? PWW_INPUT_SCALE : e->layers[li - 1].out_scale;
        }

        layer_forward(L, r, e->scratch_f);

        if (L->residual) {
            for (int c = 0; c < L->out_c; c++)
                e->scratch_f[c] += (float)e->block_in[c] * e->block_in_scale;
        }
        if (L->relu) {
            for (int c = 0; c < L->out_c; c++)
                if (e->scratch_f[c] < 0.f) e->scratch_f[c] = 0.f;
        }
        if (li == PWW_NUM_LAYERS - 1) {
            // final logits: snap to grid to mirror the simulator, then
            // return as float
            float s = L->out_scale;
            for (int c = 0; c < L->out_c; c++) {
                int8_t q = quant_clamp(e->scratch_f[c], 1.0f / s);
                logits_out[c] = (float)q * s;
            }
            return 0;
        }
        float inv = 1.0f / L->out_scale;
        int8_t *nxt = (cur == e->col_a) ? e->col_b : e->col_a;
        for (int c = 0; c < L->out_c; c++)
            nxt[c] = quant_clamp(e->scratch_f[c], inv);
        if (pww_dump_layer == li) {
            for (int c = 0; c < L->out_c && c < 256; c++)
                logits_out[c] = (float)nxt[c];
            return 0;
        }
        cur = nxt;
    }
    return -1;  // unreachable
}