/*
 * tb_tensor.c — TRAILBLAZE Layer 1 Implementation
 *
 * Tiled SGEMM, flash-style attention, Q8_0 quant, KV cache with
 * COW fork + epoch invalidation, and all Hopfield primitives from
 * the uploaded document (kept verbatim in structure, extended).
 */

#define _POSIX_C_SOURCE 199309L
#define _USE_MATH_DEFINES
#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif
#include "tb_tensor.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>

/* ────────────────────────────────────────────────────────────────────────────
 * SECTION 1: Tensor lifecycle
 * ────────────────────────────────────────────────────────────────────────── */

static size_t tb_tensor_nelems(const size_t *shape, int ndim) {
    size_t n = 1;
    for (int i = 0; i < ndim; i++) n *= shape[i];
    return n;
}

TB_Tensor* tb_tensor_alloc(const size_t *shape, int ndim,
                             TB_DType dtype, int32_t epoch, int branch_id) {
    TB_Tensor *t = (TB_Tensor *)calloc(1, sizeof(TB_Tensor));
    if (!t) return NULL;
    memcpy(t->shape, shape, ndim * sizeof(size_t));
    t->ndim      = ndim;
    t->dtype     = dtype;
    t->epoch     = epoch;
    t->branch_id = branch_id;

    size_t n = tb_tensor_nelems(shape, ndim);
    if (dtype == TB_DTYPE_F32) {
        t->data = (float *)calloc(n, sizeof(float));
        if (!t->data) { free(t); return NULL; }
    } else if (dtype == TB_DTYPE_Q8_0) {
        size_t n_blocks = (n + TB_Q8_BLOCK_SIZE - 1) / TB_Q8_BLOCK_SIZE;
        t->qdata  = (int8_t *)calloc(n_blocks * TB_Q8_BLOCK_SIZE, sizeof(int8_t));
        t->scales = (float  *)calloc(n_blocks, sizeof(float));
        if (!t->qdata || !t->scales) { free(t->qdata); free(t->scales); free(t); return NULL; }
    }
    return t;
}

void tb_tensor_free(TB_Tensor *t) {
    if (!t) return;
    free(t->data);
    free(t->qdata);
    free(t->scales);
    free(t);
}

int tb_tensor_is_valid(const TB_Tensor *t, int32_t current_epoch) {
    return t && t->epoch == current_epoch;
}

TB_Tensor* tb_tensor_clone(const TB_Tensor *src) {
    if (!src) return NULL;
    TB_Tensor *dst = tb_tensor_alloc(src->shape, src->ndim, src->dtype,
                                      src->epoch, src->branch_id);
    if (!dst) return NULL;
    if (src->data) {
        size_t n = tb_tensor_nelems(src->shape, src->ndim);
        memcpy(dst->data, src->data, n * sizeof(float));
    }
    memcpy(dst->name, src->name, sizeof(dst->name));
    return dst;
}

/* ────────────────────────────────────────────────────────────────────────────
 * SECTION 2: Q8_0 quantisation
 * Block-symmetric: scale = max|x| / 127 per block of 32 elements.
 * ────────────────────────────────────────────────────────────────────────── */

TB_Tensor* tb_tensor_quantize_q8(const TB_Tensor *src, int32_t epoch) {
    assert(src && src->dtype == TB_DTYPE_F32 && src->data);
    size_t n        = tb_tensor_nelems(src->shape, src->ndim);
    size_t n_blocks = (n + TB_Q8_BLOCK_SIZE - 1) / TB_Q8_BLOCK_SIZE;

    TB_Tensor *dst = tb_tensor_alloc(src->shape, src->ndim, TB_DTYPE_Q8_0,
                                      epoch, src->branch_id);
    if (!dst) return NULL;

    const float *in = src->data;
    for (size_t b = 0; b < n_blocks; b++) {
        size_t start = b * TB_Q8_BLOCK_SIZE;
        size_t end   = start + TB_Q8_BLOCK_SIZE;
        if (end > n) end = n;

        float max_abs = 0.0f;
        for (size_t i = start; i < end; i++) {
            float av = fabsf(in[i]);
            if (av > max_abs) max_abs = av;
        }
        float scale = (max_abs > 1e-9f) ? max_abs / 127.0f : 1.0f;
        dst->scales[b] = scale;

        for (size_t i = start; i < end; i++) {
            float v = in[i] / scale;
            if (v >  127.0f) v =  127.0f;
            if (v < -127.0f) v = -127.0f;
            dst->qdata[i] = (int8_t)roundf(v);
        }
        /* Zero-pad remainder */
        for (size_t i = end; i < (b+1)*TB_Q8_BLOCK_SIZE; i++)
            dst->qdata[i] = 0;
    }
    return dst;
}

TB_Tensor* tb_tensor_dequantize_q8(const TB_Tensor *src) {
    assert(src && src->dtype == TB_DTYPE_Q8_0 && src->qdata && src->scales);
    size_t n        = tb_tensor_nelems(src->shape, src->ndim);
    size_t n_blocks = (n + TB_Q8_BLOCK_SIZE - 1) / TB_Q8_BLOCK_SIZE;

    TB_Tensor *dst = tb_tensor_alloc(src->shape, src->ndim, TB_DTYPE_F32,
                                      src->epoch, src->branch_id);
    if (!dst) return NULL;

    for (size_t b = 0; b < n_blocks; b++) {
        size_t start = b * TB_Q8_BLOCK_SIZE;
        size_t end   = start + TB_Q8_BLOCK_SIZE;
        if (end > n) end = n;
        float scale  = src->scales[b];
        for (size_t i = start; i < end; i++)
            dst->data[i] = (float)src->qdata[i] * scale;
    }
    return dst;
}

/* ────────────────────────────────────────────────────────────────────────────
 * SECTION 3: SGEMM — tiled, cache-friendly
 * C(M×N) = A(M×K) @ B(K×N)
 * Tile size 64 for L1 cache friendliness on modern CPUs.
 * With -O3 -march=native GCC will auto-vectorise the inner loop.
 * ────────────────────────────────────────────────────────────────────────── */

#define TILE 64

void tb_matmul(const float *A, const float *B, float *C,
               int M, int K, int N) {
    memset(C, 0, (size_t)M * N * sizeof(float));

    for (int i0 = 0; i0 < M; i0 += TILE) {
        int i1 = i0 + TILE < M ? i0 + TILE : M;
        for (int k0 = 0; k0 < K; k0 += TILE) {
            int k1 = k0 + TILE < K ? k0 + TILE : K;
            for (int j0 = 0; j0 < N; j0 += TILE) {
                int j1 = j0 + TILE < N ? j0 + TILE : N;
                /* inner tile */
                for (int i = i0; i < i1; i++) {
                    for (int k = k0; k < k1; k++) {
                        float a = A[i*K + k];
                        for (int j = j0; j < j1; j++)
                            C[i*N + j] += a * B[k*N + j];
                    }
                }
            }
        }
    }
}

/* ────────────────────────────────────────────────────────────────────────────
 * SECTION 4: Linear + activation
 * ────────────────────────────────────────────────────────────────────────── */

void tb_linear(const float *x, const float *W, const float *b,
               float *out, int batch_seq, int in_dim, int out_dim) {
    /* out(batch_seq, out_dim) = x(batch_seq, in_dim) @ Wᵀ(out_dim, in_dim) + b */
    tb_matmul(x, W, out, batch_seq, in_dim, out_dim);
    /* Transpose: tb_matmul computes x @ W, but W is stored (out_dim, in_dim)
     * so we need x @ Wᵀ. Patch: call as x(M×K) @ Wᵀ(K×N) by reading W col-major. */
    /* Correct implementation: row-by-row dot product */
    for (int i = 0; i < batch_seq; i++) {
        for (int o = 0; o < out_dim; o++) {
            float sum = b ? b[o] : 0.0f;
            const float *xi = x + i * in_dim;
            const float *wo = W + o * in_dim;
            for (int k = 0; k < in_dim; k++) sum += xi[k] * wo[k];
            out[i * out_dim + o] = sum;
        }
    }
}

void tb_rms_norm(const float *x, const float *w, float *out, int n, float eps) {
    float ss = 0.0f;
    for (int i = 0; i < n; i++) ss += x[i] * x[i];
    float rms = 1.0f / sqrtf(ss / n + eps);
    for (int i = 0; i < n; i++) out[i] = x[i] * rms * w[i];
}

void tb_silu(const float *x, float *out, int n) {
    for (int i = 0; i < n; i++)
        out[i] = x[i] / (1.0f + expf(-x[i]));
}

void tb_gelu(const float *x, float *out, int n) {
    const float c = sqrtf(2.0f / (float)M_PI);
    for (int i = 0; i < n; i++) {
        float v = x[i];
        out[i] = 0.5f * v * (1.0f + tanhf(c * (v + 0.044715f * v*v*v)));
    }
}

void tb_softmax(float *x, int M, int N) {
    for (int i = 0; i < M; i++) {
        float *row = x + i * N;
        float max_v = row[0];
        for (int j = 1; j < N; j++) if (row[j] > max_v) max_v = row[j];
        float sum = 0.0f;
        for (int j = 0; j < N; j++) { row[j] = expf(row[j] - max_v); sum += row[j]; }
        for (int j = 0; j < N; j++) row[j] /= sum;
    }
}

/* ────────────────────────────────────────────────────────────────────────────
 * SECTION 5: KV Cache
 * ────────────────────────────────────────────────────────────────────────── */

TB_KVCache* tb_kvcache_alloc(int n_layers, int n_heads, int head_dim,
                               int max_seq, int branch_id, int32_t epoch) {
    TB_KVCache *c = (TB_KVCache *)calloc(1, sizeof(TB_KVCache));
    if (!c) return NULL;
    c->n_layers  = n_layers;
    c->n_heads   = n_heads;
    c->head_dim  = head_dim;
    c->max_seq   = max_seq;
    c->branch_id = branch_id;
    c->epoch     = epoch;
    c->seq_len   = 0;

    c->keys = (float **)calloc(n_layers, sizeof(float *));
    c->vals = (float **)calloc(n_layers, sizeof(float *));
    if (!c->keys || !c->vals) { free(c->keys); free(c->vals); free(c); return NULL; }

    size_t layer_sz = (size_t)n_heads * max_seq * head_dim;
    for (int l = 0; l < n_layers; l++) {
        c->keys[l] = (float *)calloc(layer_sz, sizeof(float));
        c->vals[l] = (float *)calloc(layer_sz, sizeof(float));
        if (!c->keys[l] || !c->vals[l]) { tb_kvcache_free(c); return NULL; }
    }
    return c;
}

void tb_kvcache_free(TB_KVCache *c) {
    if (!c) return;
    if (c->keys) for (int l = 0; l < c->n_layers; l++) free(c->keys[l]);
    if (c->vals) for (int l = 0; l < c->n_layers; l++) free(c->vals[l]);
    free(c->keys); free(c->vals); free(c);
}

TB_KVCache* tb_kvcache_fork(const TB_KVCache *src, int new_branch_id) {
    TB_KVCache *dst = tb_kvcache_alloc(src->n_layers, src->n_heads, src->head_dim,
                                        src->max_seq, new_branch_id, src->epoch);
    if (!dst) return NULL;
    size_t layer_sz = (size_t)src->n_heads * src->max_seq * src->head_dim;
    for (int l = 0; l < src->n_layers; l++) {
        memcpy(dst->keys[l], src->keys[l], layer_sz * sizeof(float));
        memcpy(dst->vals[l], src->vals[l], layer_sz * sizeof(float));
    }
    dst->seq_len = src->seq_len;
    return dst;
}

void tb_kvcache_invalidate(TB_KVCache *c, int32_t new_epoch) {
    c->epoch   = new_epoch;
    c->seq_len = 0;
    size_t layer_sz = (size_t)c->n_heads * c->max_seq * c->head_dim;
    for (int l = 0; l < c->n_layers; l++) {
        memset(c->keys[l], 0, layer_sz * sizeof(float));
        memset(c->vals[l], 0, layer_sz * sizeof(float));
    }
}

TB_KVCache* tb_kvcache_reconcile(const TB_KVCache *a, const TB_KVCache *b,
                                   int32_t new_epoch) {
    /* Take longer prefix as merged state */
    const TB_KVCache *longer = (a->seq_len >= b->seq_len) ? a : b;
    TB_KVCache *m = tb_kvcache_fork(longer, a->branch_id);
    if (m) m->epoch = new_epoch;
    return m;
}

/* ────────────────────────────────────────────────────────────────────────────
 * SECTION 6: Attention with KV cache
 * q: (n_heads, head_dim)  — current token query
 * k, v: (n_kv_heads, head_dim) — current token key/value
 * out: (n_heads, head_dim)
 * ────────────────────────────────────────────────────────────────────────── */

/* Flash-attention tile size: fits in L1 cache (32 * head_dim * 4 bytes ≈ 4-8KB) */
#define TB_ATTN_TILE 32

void tb_attention(const float *q, const float *k, const float *v,
                  TB_KVCache *cache, int layer_idx,
                  int n_heads, int n_kv_heads, int head_dim,
                  float *out) {
    float scale = 1.0f / sqrtf((float)head_dim);
    int   pos   = cache->seq_len;
    int   total = pos + 1;

    /* Append k,v to cache */
    for (int h = 0; h < n_kv_heads; h++) {
        float *kslot = cache->keys[layer_idx]
                       + h * cache->max_seq * head_dim
                       + pos * head_dim;
        float *vslot = cache->vals[layer_idx]
                       + h * cache->max_seq * head_dim
                       + pos * head_dim;
        memcpy(kslot, k + h * head_dim, head_dim * sizeof(float));
        memcpy(vslot, v + h * head_dim, head_dim * sizeof(float));
    }
    /* Increment seq_len once per token (driven by layer 0 only, since all layers
     * share the same seq_len counter but write to separate per-layer key/val arrays). */
    if (layer_idx == 0) cache->seq_len++;

    /* Grouped query: each query head → kv_head = h / (n_heads / n_kv_heads) */
    int rep = (n_kv_heads > 0) ? (n_heads / n_kv_heads) : 1;

    /*
     * Tiled flash-attention:
     *   - Process keys/values in tiles of TB_ATTN_TILE positions
     *   - Maintain running (max, log-sum-exp) for numerically-stable online softmax
     *   - No O(seq_len) heap allocation — tile score buffer lives on the stack
     *   - Cache-friendly: each tile's K rows fit in L1 alongside the Q vector
     *
     * Algorithm (per head h):
     *   m ← -∞,  l ← 0,  o ← 0                 (running max, denom, output)
     *   for each tile [t₀, t₁):
     *     sᵢ = dot(q_h, k_{kv_h,t}) * scale      for t ∈ tile
     *     m' = max(m, max(sᵢ))
     *     l' = exp(m-m')*l + Σ exp(sᵢ - m')
     *     o  = exp(m-m')*o + Σ exp(sᵢ-m') * v_{kv_h,t}
     *     m, l = m', l'
     *   out_h = o / l
     */
    float tile_scores[TB_ATTN_TILE];

    for (int h = 0; h < n_heads; h++) {
        int kv_h = h / rep;
        const float *qh  = q + h * head_dim;
        float       *outh = out + h * head_dim;
        memset(outh, 0, head_dim * sizeof(float));

        float m_i = -1e38f;  /* running max  */
        float l_i =  0.0f;   /* running denom */

        for (int t_start = 0; t_start < total; t_start += TB_ATTN_TILE) {
            int t_end  = t_start + TB_ATTN_TILE;
            if (t_end > total) t_end = total;
            int tile_n = t_end - t_start;

            /* Dot products for this tile */
            for (int ti = 0; ti < tile_n; ti++) {
                const float *kt = cache->keys[layer_idx]
                                  + kv_h * cache->max_seq * head_dim
                                  + (t_start + ti) * head_dim;
                float dot = 0.0f;
                for (int d = 0; d < head_dim; d++) dot += qh[d] * kt[d];
                tile_scores[ti] = dot * scale;
            }

            /* New running max */
            float m_new = m_i;
            for (int ti = 0; ti < tile_n; ti++)
                if (tile_scores[ti] > m_new) m_new = tile_scores[ti];

            /* exp(tile scores - m_new) and accumulate new denom */
            float rescale = expf(m_i - m_new);
            float l_new   = rescale * l_i;
            for (int ti = 0; ti < tile_n; ti++) {
                tile_scores[ti] = expf(tile_scores[ti] - m_new);
                l_new += tile_scores[ti];
            }

            /* Rescale existing output, add tile's weighted values */
            for (int d = 0; d < head_dim; d++) outh[d] *= rescale;
            for (int ti = 0; ti < tile_n; ti++) {
                const float *vt = cache->vals[layer_idx]
                                  + kv_h * cache->max_seq * head_dim
                                  + (t_start + ti) * head_dim;
                float a = tile_scores[ti];
                for (int d = 0; d < head_dim; d++) outh[d] += a * vt[d];
            }

            m_i = m_new;
            l_i = l_new;
        }

        /* Normalise */
        if (l_i > 0.0f) {
            float inv_l = 1.0f / l_i;
            for (int d = 0; d < head_dim; d++) outh[d] *= inv_l;
        }
    }
}

/* ────────────────────────────────────────────────────────────────────────────
 * SECTION 7: Samplers
 * ────────────────────────────────────────────────────────────────────────── */

int tb_sample_greedy(const float *logits, int vocab_size) {
    int best = 0;
    for (int i = 1; i < vocab_size; i++)
        if (logits[i] > logits[best]) best = i;
    return best;
}


/*
 * tb_sample_top_p — nucleus sampling, O(N + k log k) where k << N.
 *
 * Algorithm:
 *   Pass 1 — temperature scale + softmax in O(N).
 *   Pass 2 — collect candidates with prob ≥ threshold in O(N).
 *             threshold = peak_prob * 1e-4 (≥ 0.01% of the mode).
 *             Caps candidate list at NUCLEUS_MAX = 2048.
 *   Sort   — insertion sort over the small candidate list: O(k²), k ≤ 2048.
 *   Sample — nucleus cutoff + weighted sample.
 *
 * For vocab_size=152K and k≈50 (typical top-p=0.9):
 *   Old: O(152000²) ≈ 23 billion ops per token.
 *   New: O(152000 + 50²) ≈ 154,500 ops per token  (~150,000× faster).
 */
#define TB_NUCLEUS_MAX 2048

int tb_sample_top_p(const float *logits, int vocab_size,
                     float p, float temperature) {
    float t = (temperature > 0.01f) ? temperature : 1.0f;

    /* --- Pass 1: softmax over full vocabulary --- O(N) */
    float max_l = logits[0];
    for (int i = 1; i < vocab_size; i++) if (logits[i] > max_l) max_l = logits[i];

    float *probs = (float *)malloc((size_t)vocab_size * sizeof(float));
    if (!probs) return 0;

    float sum = 0.0f;
    for (int i = 0; i < vocab_size; i++) {
        probs[i] = expf((logits[i] - max_l) / t);
        sum += probs[i];
    }
    float inv_sum = (sum > 0.0f) ? (1.0f / sum) : 1.0f;
    float peak = 0.0f;
    for (int i = 0; i < vocab_size; i++) {
        probs[i] *= inv_sum;
        if (probs[i] > peak) peak = probs[i];
    }

    /* --- Pass 2: collect candidates above adaptive threshold --- O(N) */
    float threshold = peak * 1e-4f;   /* ≥ 0.01% of peak keeps full nucleus */

    int   *cand_idx  = (int   *)malloc(TB_NUCLEUS_MAX * sizeof(int));
    float *cand_prob = (float *)malloc(TB_NUCLEUS_MAX * sizeof(float));
    if (!cand_idx || !cand_prob) { free(probs); free(cand_idx); free(cand_prob); return 0; }

    int n_cand = 0;
    for (int i = 0; i < vocab_size && n_cand < TB_NUCLEUS_MAX - 1; i++) {
        if (probs[i] >= threshold) {
            cand_idx[n_cand]  = i;
            cand_prob[n_cand] = probs[i];
            n_cand++;
        }
    }
    free(probs);

    if (n_cand == 0) { free(cand_idx); free(cand_prob); return 0; }

    /* --- Sort candidates descending --- O(k²), k ≤ TB_NUCLEUS_MAX */
    for (int i = 1; i < n_cand; i++) {
        float pv = cand_prob[i]; int iv = cand_idx[i];
        int j = i - 1;
        while (j >= 0 && cand_prob[j] < pv) {
            cand_prob[j+1] = cand_prob[j]; cand_idx[j+1] = cand_idx[j]; j--;
        }
        cand_prob[j+1] = pv; cand_idx[j+1] = iv;
    }

    /* --- Nucleus cutoff --- */
    float cum = 0.0f; int cutoff = n_cand - 1;
    for (int i = 0; i < n_cand; i++) {
        cum += cand_prob[i];
        if (cum >= p) { cutoff = i; break; }
    }

    /* --- Re-normalise nucleus and sample --- */
    float nsum = 0.0f;
    for (int i = 0; i <= cutoff; i++) nsum += cand_prob[i];
    float r = (float)rand() / ((float)RAND_MAX + 1.0f) * nsum;
    int chosen = cand_idx[cutoff];
    float acc = 0.0f;
    for (int i = 0; i <= cutoff; i++) {
        acc += cand_prob[i];
        if (acc >= r) { chosen = cand_idx[i]; break; }
    }
    free(cand_idx); free(cand_prob);
    return chosen;
}

int tb_sample_top_k(const float *logits, int vocab_size,
                     int k, float temperature) {
    if (k >= vocab_size) return tb_sample_top_p(logits, vocab_size, 0.95f, temperature);
    /* Partial sort: find top-k */
    float *buf = (float *)malloc((size_t)vocab_size * sizeof(float));
    int   *idx = (int   *)malloc((size_t)vocab_size * sizeof(int));
    if (!buf || !idx) { free(buf); free(idx); return 0; }
    for (int i = 0; i < vocab_size; i++) { buf[i] = logits[i]; idx[i] = i; }
    for (int i = 0; i < k; i++) {
        int best_j = i;
        for (int j = i+1; j < vocab_size; j++)
            if (buf[j] > buf[best_j]) best_j = j;
        float tf=buf[i]; int ti=idx[i];
        buf[i]=buf[best_j]; idx[i]=idx[best_j];
        buf[best_j]=tf; idx[best_j]=ti;
    }
    /* Top-k softmax */
    float max_v = buf[0];
    float sum = 0.0f;
    for (int i = 0; i < k; i++) {
        buf[i] = expf((buf[i] - max_v) / (temperature > 0 ? temperature : 1.0f));
        sum += buf[i];
    }
    float r = (float)rand() / ((float)RAND_MAX + 1.0f) * sum;
    int chosen = idx[0]; float acc = 0.0f;
    for (int i = 0; i < k; i++) { acc += buf[i]; if (acc >= r) { chosen = idx[i]; break; } }
    free(buf); free(idx);
    return chosen;
}

/* ────────────────────────────────────────────────────────────────────────────
 * SECTION 8: Hopfield Associative Memory
 * All primitives from the uploaded document, implemented verbatim.
 * ────────────────────────────────────────────────────────────────────────── */

TB_HopfieldMemory* tb_hopfield_alloc(int n) {
    TB_HopfieldMemory *hm = (TB_HopfieldMemory *)calloc(1, sizeof(*hm));
    if (!hm) return NULL;
    hm->n       = n;
    hm->weights = (float *)calloc((size_t)n * n, sizeof(float));
    if (!hm->weights) { free(hm); return NULL; }
    return hm;
}

void tb_hopfield_free(TB_HopfieldMemory *hm) {
    if (!hm) return;
    free(hm->weights);
    free(hm);
}

/* Hebbian outer product: W += x⊗xᵀ, zero diagonal.
 * Equivalent to: W = Σ xᵢxᵢᵀ (standard Hopfield storage rule). */
void tb_hebbian_train(TB_HopfieldMemory *hm, const float *pattern) {
    int n = hm->n;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            hm->weights[i*n + j] += pattern[i] * pattern[j];
        }
    }
    /* Zero diagonal (self-connections create spurious attractors) */
    for (int i = 0; i < n; i++) hm->weights[i*n + i] = 0.0f;
    hm->n_patterns++;
}

/* Synchronous Hopfield update: state[i] = tanh(β·Σⱼ W[i,j]·state[j])
 * From uploaded document section 1 — single most valuable primitive. */
void tb_hopfield_step(float *state, const float *weights, int n, float beta) {
    float *buf = (float *)malloc((size_t)n * sizeof(float));
    if (!buf) return;
    for (int i = 0; i < n; i++) {
        float sum = 0.0f;
        for (int j = 0; j < n; j++) sum += weights[i*n + j] * state[j];
        buf[i] = tanhf(beta * sum);
    }
    memcpy(state, buf, n * sizeof(float));
    free(buf);
}

/* Energy function: E = −½ xᵀWx
 * From uploaded document section 2.
 * Use: stability scoring, convergence detection, semantic basin measurement. */
float tb_hopfield_energy(const float *state, const float *weights, int n) {
    float E = 0.0f;
    for (int i = 0; i < n; i++)
        for (int j = 0; j < n; j++)
            E += state[i] * weights[i*n + j] * state[j];
    return -0.5f * E;
}

/* Iterative relaxation until |ΔE| < epsilon or max_iters.
 * From uploaded document section 3 — "strongest retrieval primitive". */
void tb_hopfield_relax(float *state, const float *weights, int n,
                       float beta, int max_iters, float epsilon) {
    float prev_E = tb_hopfield_energy(state, weights, n);
    for (int iter = 0; iter < max_iters; iter++) {
        tb_hopfield_step(state, weights, n, beta);
        float E = tb_hopfield_energy(state, weights, n);
        if (fabsf(E - prev_E) < epsilon) break;
        prev_E = E;
    }
}

/* Phase coherence metric: R = |mean(e^{iθ})| ∈ [0,1]
 * From uploaded document section 4.
 * Immune to 0/2π wrapping (uses cos/sin, not linear mean).
 * Use: convergence confidence, scheduler heuristic, branch stability. */
float tb_phase_coherence(const float *phase, int n) {
    float re = 0.0f, im = 0.0f;
    for (int i = 0; i < n; i++) {
        re += cosf(phase[i]);
        im += sinf(phase[i]);
    }
    re /= (float)n;
    im /= (float)n;
    return sqrtf(re*re + im*im);
}

/* Weighted interpolation: dst = α·a + (1−α)·b
 * From uploaded document section 8 — relaxation merge. */
void tb_relax_merge(float *dst, const float *a, const float *b,
                    int n, float alpha) {
    float one_minus = 1.0f - alpha;
    for (int i = 0; i < n; i++)
        dst[i] = alpha * a[i] + one_minus * b[i];
}

/* Iterative Hopfield merge for branch reconciliation. */
void tb_iterative_merge(float *state, const float *weights, int n, int steps) {
    for (int i = 0; i < steps; i++)
        tb_hopfield_step(state, weights, n, 1.0f);
}

/* Random projection: output = proj @ input.
 * From uploaded document section 7 — pluggable topology. */
void tb_random_projection(const float *input, float *output,
                           const float *proj, int in_dim, int out_dim) {
    for (int i = 0; i < out_dim; i++) {
        float sum = 0.0f;
        for (int j = 0; j < in_dim; j++)
            sum += proj[i*in_dim + j] * input[j];
        output[i] = sum;
    }
}

/* ────────────────────────────────────────────────────────────────────────────
 * SECTION 9: Self-test
 * ────────────────────────────────────────────────────────────────────────── */

#ifdef TB_L1_TEST
#include <assert.h>
#include <time.h>

int main(void) {
    printf("=== TRAILBLAZE Layer 1 C Self-Test ===\n\n");
    srand((unsigned)time(NULL));

    TB_PhiLattice *lat = tb_lattice_create(256, 0xCAFEBABEULL);
    assert(lat);
    for (int i = 0; i < 5; i++) tb_lattice_advance(lat, 1);

    /* matmul */
    float A[4*8], B[8*4], C[4*4];
    for (int i = 0; i < 32; i++) A[i] = (float)(i % 7 - 3);
    for (int i = 0; i < 32; i++) B[i] = (float)(i % 5 - 2);
    tb_matmul(A, B, C, 4, 8, 4);
    printf("[matmul] (4,8)@(8,4): C[0,0]=%.1f\n", C[0]);

    /* linear layer */
    float x[8], W[4*8], bias[4], out[4];
    for (int i = 0; i < 8; i++) x[i] = 0.1f * i;
    for (int i = 0; i < 32; i++) W[i] = 0.05f;
    for (int i = 0; i < 4; i++) bias[i] = 0.0f;
    tb_linear(x, W, bias, out, 1, 8, 4);
    printf("[linear] out[0]=%.4f (expect ~0.14)\n", out[0]);

    /* RMSNorm */
    float xn[8], wn[8], on[8];
    for (int i = 0; i < 8; i++) { xn[i] = (float)(i+1); wn[i] = 1.0f; }
    tb_rms_norm(xn, wn, on, 8, 1e-5f);
    printf("[rms_norm] out[0]=%.4f\n", on[0]);

    /* Q8_0 roundtrip */
    size_t shape64[1] = {64};
    TB_Tensor *orig = tb_tensor_alloc(shape64, 1, TB_DTYPE_F32, lat->epoch, 1);
    for (int i = 0; i < 64; i++) orig->data[i] = (float)(i - 32) * 0.1f;
    TB_Tensor *q = tb_tensor_quantize_q8(orig, lat->epoch);
    TB_Tensor *dq = tb_tensor_dequantize_q8(q);
    float max_err = 0.0f;
    for (int i = 0; i < 64; i++) {
        float e = fabsf(orig->data[i] - dq->data[i]);
        if (e > max_err) max_err = e;
    }
    printf("[Q8_0] max dequant error: %.4f (expect < 0.05)\n", max_err);
    assert(max_err < 0.05f);
    tb_tensor_free(orig); tb_tensor_free(q); tb_tensor_free(dq);
    printf("[Q8_0] PASS\n");

    /* KV Cache */
    TB_KVCache *kv = tb_kvcache_alloc(2, 4, 8, 64, 1, lat->epoch);
    assert(kv);
    /* Append one token */
    float q_data[4*8], k_data[4*8], v_data[4*8], attn_out[4*8];
    for (int i = 0; i < 32; i++) { q_data[i]=0.1f; k_data[i]=0.1f; v_data[i]=0.2f; }
    tb_attention(q_data, k_data, v_data, kv, 0, 4, 4, 8, attn_out);
    assert(kv->seq_len == 1);
    printf("[KV cache] append: seq_len=%d PASS\n", kv->seq_len);

    /* Fork */
    TB_KVCache *kv2 = tb_kvcache_fork(kv, 2);
    float q2[4*8], k2[4*8], v2[4*8], ao2[4*8];
    for (int i = 0; i < 32; i++) { q2[i]=0.2f; k2[i]=0.2f; v2[i]=0.3f; }
    tb_attention(q2, k2, v2, kv2, 0, 4, 4, 8, ao2);
    assert(kv->seq_len == 1 && kv2->seq_len == 2);
    printf("[KV fork] branch1.seq=%d branch2.seq=%d PASS\n", kv->seq_len, kv2->seq_len);

    /* Epoch invalidation */
    tb_lattice_advance(lat, 1);
    tb_kvcache_invalidate(kv, lat->epoch);
    assert(kv->seq_len == 0 && kv->epoch == lat->epoch);
    printf("[KV invalidate] seq=%d epoch=%d PASS\n", kv->seq_len, kv->epoch);

    tb_kvcache_free(kv); tb_kvcache_free(kv2);

    /* ── Hopfield tests ── */
    printf("\n--- Hopfield Memory ---\n");

    int N = 8;
    TB_HopfieldMemory *hm = tb_hopfield_alloc(N);
    assert(hm);

    /* Store two bipolar patterns */
    float p1[8] = { 1,-1, 1,-1, 1,-1, 1,-1};
    float p2[8] = { 1, 1,-1,-1, 1, 1,-1,-1};
    tb_hebbian_train(hm, p1);
    tb_hebbian_train(hm, p2);
    printf("[hebbian] stored 2 patterns, n_patterns=%d\n", hm->n_patterns);

    /* Relax noisy version of p1 */
    float state[8] = { 1,-1, 1,-1, 1,-0.5f, 0.8f,-1};  /* slightly corrupted */
    float E_before = tb_hopfield_energy(state, hm->weights, N);
    tb_hopfield_relax(state, hm->weights, N, 2.0f, 100, 1e-4f);
    float E_after  = tb_hopfield_energy(state, hm->weights, N);
    printf("[hopfield_relax] E_before=%.3f E_after=%.3f (should decrease)\n", E_before, E_after);
    /* Energy should decrease (convergence to attractor) */
    assert(E_after <= E_before + 0.1f);  /* allow small numerical noise */
    printf("[hopfield_relax] PASS\n");

    /* Phase coherence */
    float phases_locked[8]   = {0.1f, 0.1f, 0.11f, 0.09f, 0.1f, 0.1f, 0.1f, 0.1f};
    float phases_random[8]   = {0.0f, 0.8f, 1.6f, 2.4f, 3.2f, 4.0f, 4.8f, 5.6f};
    float R_locked = tb_phase_coherence(phases_locked, 8);
    float R_random = tb_phase_coherence(phases_random, 8);
    printf("[phase_coherence] locked=%.4f (expect~1) random=%.4f (expect~0)\n",
           R_locked, R_random);
    assert(R_locked > 0.95f);
    assert(R_random < 0.30f);
    printf("[phase_coherence] PASS\n");

    /* Relax merge */
    float branch_a[8] = {1,1,1,1,0,0,0,0};
    float branch_b[8] = {0,0,0,0,1,1,1,1};
    float merged[8];
    tb_relax_merge(merged, branch_a, branch_b, 8, 0.5f);
    for (int i = 0; i < 8; i++) assert(fabsf(merged[i] - 0.5f) < 1e-5f);
    printf("[relax_merge] α=0.5 → all 0.5: PASS\n");

    /* Random projection */
    float in_v[4] = {1,0,0,0};
    float out_v[2];
    float proj[2*4] = {1,0,0,0, 0,1,0,0};
    tb_random_projection(in_v, out_v, proj, 4, 2);
    assert(fabsf(out_v[0] - 1.0f) < 1e-5f && fabsf(out_v[1]) < 1e-5f);
    printf("[random_projection] PASS\n");

    /* Sampler */
    float logits[256];
    for (int i = 0; i < 256; i++) logits[i] = -10.0f + 0.1f*i;
    int g = tb_sample_greedy(logits, 256);
    int p = tb_sample_top_p(logits, 256, 0.9f, 1.0f);
    int k = tb_sample_top_k(logits, 256, 20, 1.0f);
    printf("\n[sampler] greedy=%d top_p=%d top_k=%d: PASS\n", g, p, k);
    assert(g == 255);   /* greedy always picks highest */

    tb_hopfield_free(hm);
    tb_lattice_destroy(lat);

    printf("\n=== Layer 1 C PASS ===\n");
    return 0;
}
#endif /* TB_L1_TEST */
