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

void tmk_linear_vec(const uint8_t* codes, int step, int m, int n, const float* tlut, const uint8_t* su,
                    const uint8_t* sv, float wscale, const float* x, float* y, float* tmp);
void tmk_gelu(float* x, long n);
void tmk_silu(float* x, long n);

typedef struct { const uint8_t* codes; const uint8_t* su; const uint8_t* sv; float ws; int m, n; } tmk_chunk;
typedef struct { const tmk_chunk* ch; int nch, step; const float* tlut; const float* bias; const float* dense; int m, n; } tmk_lin;

#define MAXL 8
typedef struct {
    int D, NL, NH, HD, FF, Pmax;
    float eps;
    const float* inv_freq;
    const float* ln[MAXL][4];
    tmk_lin lin[MAXL][6];
    const float *lnf_w, *lnf_b;
    tmk_lin in_proj;
    const float* in_add;
} tmk_lm;

typedef struct {
    int D, NL, LD;
    float eps;
    tmk_lin inp, cond, ada[MAXL], fc1[MAXL], fc2[MAXL], fada, fproj;
    const float* norm[MAXL][2];
    const float* time;
} tmk_flow;

static void lin(const tmk_lin* L, const float* x, float* y, float* tmp) {
    if (L->nch == 0) {
        for (int i = 0; i < L->m; i++) {
            const float* w = L->dense + (size_t)i * L->n;
            float a = 0.f;
            for (int j = 0; j < L->n; j++) a += w[j] * x[j];
            y[i] = a;
        }
    } else {
        int r0 = 0;
        for (int c = 0; c < L->nch; c++) {
            const tmk_chunk* C = L->ch + c;
            tmk_linear_vec(C->codes, L->step, C->m, C->n, L->tlut, C->su, C->sv, C->ws, x, y + r0, tmp);
            r0 += C->m;
        }
    }
    if (L->bias) for (int i = 0; i < L->m; i++) y[i] += L->bias[i];
}

static void layer_norm(const float* x, float* y, int n, const float* w, const float* b, float eps) {
    float mu = 0.f, var = 0.f;
    for (int i = 0; i < n; i++) mu += x[i];
    mu /= n;
    for (int i = 0; i < n; i++) { float d = x[i] - mu; var += d * d; }
    float r = 1.0f / sqrtf(var / n + eps);
    for (int i = 0; i < n; i++) { float v = (x[i] - mu) * r; y[i] = w ? v * w[i] + b[i] : v; }
}

static void rope(float* x, int H, int hd, int pos, const float* inv) {
    const int h2 = hd / 2;
    for (int i = 0; i < h2; i++) {
        float f = (float)pos * inv[i], c = cosf(f), s = sinf(f);
        for (int h = 0; h < H; h++) {
            float* v = x + h * hd;
            float a = v[i], b = v[i + h2];
            v[i] = a * c - b * s; v[i + h2] = b * c + a * s;
        }
    }
}

size_t tmk_lm_work(const tmk_lm* M) { return 6 * (size_t)M->D + 2 * (size_t)M->FF + M->Pmax + 4096; }

void tmk_lm_step(const tmk_lm* M, const float* lat, const float* x0, int pos, float* const* Kc, float* const* Vc,
                 float* hout, float* work) {
    const int D = M->D, H = M->NH, hd = M->HD;
    float *x = work, *h = x + D, *q = h + D, *k = q + D, *v = k + D, *a = v + D, *u = a + D, *s = u + 2 * M->FF,
          *tmp = s + M->Pmax;
    if (lat) {
        lin(&M->in_proj, lat, x, tmp);
        if (M->in_add) for (int i = 0; i < D; i++) x[i] += M->in_add[i];
    } else memcpy(x, x0, sizeof(float) * D);
    const float scale = 1.0f / sqrtf((float)hd);
    for (int l = 0; l < M->NL; l++) {
        layer_norm(x, h, D, M->ln[l][0], M->ln[l][1], M->eps);
        lin(&M->lin[l][0], h, q, tmp); lin(&M->lin[l][1], h, k, tmp); lin(&M->lin[l][2], h, v, tmp);
        rope(q, H, hd, pos, M->inv_freq); rope(k, H, hd, pos, M->inv_freq);
        for (int hh = 0; hh < H; hh++) {
            float* Kh = Kc[l] + (size_t)hh * M->Pmax * hd;
            float* Vh = Vc[l] + (size_t)hh * M->Pmax * hd;
            memcpy(Kh + (size_t)pos * hd, k + hh * hd, sizeof(float) * hd);
            memcpy(Vh + (size_t)pos * hd, v + hh * hd, sizeof(float) * hd);
            const float* qh = q + hh * hd;
            float mx = -3.0e38f;
            for (int j = 0; j <= pos; j++) {
                const float* kj = Kh + (size_t)j * hd;
                float d = 0.f;
                for (int e = 0; e < hd; e++) d += qh[e] * kj[e];
                s[j] = d * scale;
                mx = s[j] > mx ? s[j] : mx;
            }
            float sum = 0.f;
            for (int j = 0; j <= pos; j++) { s[j] = expf(s[j] - mx); sum += s[j]; }
            float* ah = a + hh * hd;
            for (int e = 0; e < hd; e++) ah[e] = 0.f;
            for (int j = 0; j <= pos; j++) {
                const float* vj = Vh + (size_t)j * hd;
                const float w = s[j] / sum;
                for (int e = 0; e < hd; e++) ah[e] += w * vj[e];
            }
        }
        lin(&M->lin[l][3], a, h, tmp);
        for (int i = 0; i < D; i++) x[i] += h[i];
        layer_norm(x, h, D, M->ln[l][2], M->ln[l][3], M->eps);
        lin(&M->lin[l][4], h, u, tmp);
        tmk_gelu(u, M->FF);
        lin(&M->lin[l][5], u, h, tmp);
        for (int i = 0; i < D; i++) x[i] += h[i];
    }
    layer_norm(x, hout, D, M->lnf_w, M->lnf_b, M->eps);
}

void tmk_flow_step(const tmk_flow* F, const float* hcond, const float* noise, float* out, float* work) {
    const int D = F->D;
    float *x = work, *sc = x + D, *ada = sc + D, *m = ada + 3 * D, *u = m + D, *t = u + D, *tmp = t + D;
    lin(&F->inp, noise, x, tmp);
    lin(&F->cond, hcond, sc, tmp);
    for (int i = 0; i < D; i++) sc[i] += F->time[i];
    tmk_silu(sc, D);
    for (int l = 0; l < F->NL; l++) {
        lin(&F->ada[l], sc, ada, tmp);
        const float *shift = ada, *scale = ada + D, *gate = ada + 2 * D;
        layer_norm(x, m, D, F->norm[l][0], F->norm[l][1], F->eps);
        for (int i = 0; i < D; i++) m[i] = m[i] * (1.f + scale[i]) + shift[i];
        lin(&F->fc1[l], m, u, tmp);
        tmk_silu(u, D);
        lin(&F->fc2[l], u, t, tmp);
        for (int i = 0; i < D; i++) x[i] += gate[i] * t[i];
    }
    lin(&F->fada, sc, ada, tmp);
    layer_norm(x, m, D, NULL, NULL, F->eps);
    for (int i = 0; i < D; i++) m[i] = m[i] * (1.f + ada[D + i]) + ada[i];
    lin(&F->fproj, m, out, tmp);
    for (int i = 0; i < F->LD; i++) out[i] += noise[i];
}