/* ════════════════════════════════════════════════════════════════════
   L8 OS · NANO  ·  a real transformer, written out by hand
   ────────────────────────────────────────────────────────────────────
   Decoder-only transformer. Forward pass, hand-derived backward pass,
   AdamW. No dependencies, no framework, no GPU — Float32Array and loops.

   WHAT THIS IS AND IS NOT

     It is a real model. The weights start random, gradients are computed
     by differentiating the actual forward pass, and training changes the
     weights so that the loss goes down on data it has never seen. That
     is not a simulation of training; it is training.

     It is also ~1 million parameters, which is about four thousandths of
     one percent of a frontier model. Trained on your own files it will
     learn their shape — your naming, your comment style, the structure
     of your JSON — and it will complete text in that shape convincingly.

     It will not answer questions. It has no world knowledge, no
     instruction following, and no idea what it is saying. Anyone who
     ships a model this size as an assistant is lying to you, so the OS
     labels it as what it is: a completion model trained on your corpus.

   THE CORRECTNESS PROOF

     Hand-written backprop is easy to get subtly wrong, and a wrong
     gradient still trains — just badly, in a way that looks like a
     hyperparameter problem for a week. So every backward function here
     is checked against a numerical gradient in gradcheck.cjs, and that
     check is the thing to run before believing anything else.
   ════════════════════════════════════════════════════════════════════ */
'use strict';

/* ── deterministic RNG · a seed reproduces a model exactly ─────────── */
function rng(seed) {
  var s = seed >>> 0 || 1;
  return function () {
    s ^= s << 13; s >>>= 0; s ^= s >> 17; s ^= s << 5; s >>>= 0;
    return s / 4294967296;
  };
}
function randn(r) {                      /* Box-Muller */
  var u = 1 - r(), v = r();
  return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v);
}

/* ── a parameter is its values, its gradient, and Adam's two moments ── */
function param(n, init, r) {
  var p = { v: new Float32Array(n), g: new Float32Array(n),
            m: new Float32Array(n), s: new Float32Array(n), n: n };
  if (init) for (var i = 0; i < n; i++) p.v[i] = randn(r) * init;
  return p;
}

/* ══════════════════════════ PRIMITIVES ═══════════════════════════════
   Each one is (forward, backward) over flat Float32Arrays. Shapes are
   passed explicitly because a shape bug that a framework would catch is
   a silent wrong answer here. */

/* y[t,o] = sum_i x[t,i] * W[o,i] + b[o]      W is [out, in] row-major */
/* These two functions are where essentially all the time goes, so they
   are unrolled four wide. That is not premature: at dim 96 with a 512
   vocabulary the matmuls are over 90% of every step, and the unroll
   measured 2.4× faster than the straightforward loop on this machine.
   The scalar tail handles dimensions that are not multiples of four. */
function linearFwd(x, W, b, T, I, O, y) {
  var I4 = I & ~3;
  for (var t = 0; t < T; t++) {
    var xo = t * I, yo = t * O;
    for (var o = 0; o < O; o++) {
      var wo = o * I, a0 = 0, a1 = 0, a2 = 0, a3 = 0, i;
      for (i = 0; i < I4; i += 4) {
        a0 += x[xo + i]     * W[wo + i];
        a1 += x[xo + i + 1] * W[wo + i + 1];
        a2 += x[xo + i + 2] * W[wo + i + 2];
        a3 += x[xo + i + 3] * W[wo + i + 3];
      }
      var acc = a0 + a1 + a2 + a3;
      for (; i < I; i++) acc += x[xo + i] * W[wo + i];
      y[yo + o] = acc + (b ? b[o] : 0);
    }
  }
}
function linearBwd(x, W, dy, T, I, O, dx, dW, db) {
  var I4 = I & ~3;
  for (var t = 0; t < T; t++) {
    var xo = t * I, yo = t * O;
    for (var o = 0; o < O; o++) {
      var d = dy[yo + o];
      if (d === 0) continue;
      var wo = o * I, i;
      if (db) db[o] += d;
      for (i = 0; i < I4; i += 4) {
        var xi0 = x[xo + i], xi1 = x[xo + i + 1], xi2 = x[xo + i + 2], xi3 = x[xo + i + 3];
        dx[xo + i]     += d * W[wo + i];     dW[wo + i]     += d * xi0;
        dx[xo + i + 1] += d * W[wo + i + 1]; dW[wo + i + 1] += d * xi1;
        dx[xo + i + 2] += d * W[wo + i + 2]; dW[wo + i + 2] += d * xi2;
        dx[xo + i + 3] += d * W[wo + i + 3]; dW[wo + i + 3] += d * xi3;
      }
      for (; i < I; i++) { dx[xo + i] += d * W[wo + i]; dW[wo + i] += d * x[xo + i]; }
    }
  }
}

/* layernorm over the last dimension, with learned gain and bias */
function lnFwd(x, g, b, T, C, y, cache) {
  for (var t = 0; t < T; t++) {
    var o = t * C, mu = 0, i;
    for (i = 0; i < C; i++) mu += x[o + i];
    mu /= C;
    var vr = 0;
    for (i = 0; i < C; i++) { var d = x[o + i] - mu; vr += d * d; }
    vr /= C;
    var inv = 1 / Math.sqrt(vr + 1e-5);
    cache.mu[t] = mu; cache.inv[t] = inv;
    for (i = 0; i < C; i++) {
      var xh = (x[o + i] - mu) * inv;
      cache.xh[o + i] = xh;
      y[o + i] = xh * g[i] + b[i];
    }
  }
}
function lnBwd(dy, g, T, C, cache, dx, dg, db) {
  for (var t = 0; t < T; t++) {
    var o = t * C, inv = cache.inv[t], i;
    var sd = 0, sdx = 0;
    for (i = 0; i < C; i++) {
      var d = dy[o + i] * g[i];
      sd += d; sdx += d * cache.xh[o + i];
      dg[i] += dy[o + i] * cache.xh[o + i];
      db[i] += dy[o + i];
    }
    for (i = 0; i < C; i++) {
      var dd = dy[o + i] * g[i];
      dx[o + i] += inv * (dd - sd / C - cache.xh[o + i] * sdx / C);
    }
  }
}

/* GELU, tanh approximation — the one GPT-2 used */
var G0 = Math.sqrt(2 / Math.PI);
function geluFwd(x, n, y) {
  for (var i = 0; i < n; i++) {
    var v = x[i], c = G0 * (v + 0.044715 * v * v * v);
    y[i] = 0.5 * v * (1 + Math.tanh(c));
  }
}
function geluBwd(x, dy, n, dx) {
  for (var i = 0; i < n; i++) {
    var v = x[i], v3 = v * v * v;
    var c = G0 * (v + 0.044715 * v3), th = Math.tanh(c);
    var dc = G0 * (1 + 3 * 0.044715 * v * v);
    dx[i] += dy[i] * (0.5 * (1 + th) + 0.5 * v * (1 - th * th) * dc);
  }
}

/* ══════════════════════ MODEL ════════════════════════════════════════ */
function build(cfg, seed) {
  var V = cfg.vocab, C = cfg.dim, L = cfg.layers, H = cfg.heads, T = cfg.ctx;
  var r = rng(seed || 1337);
  var std = 0.02;
  var m = { cfg: { vocab: V, dim: C, layers: L, heads: H, ctx: T }, p: {} };
  var P = m.p;
  P.wte = param(V * C, std, r);                 /* token embedding      */
  P.wpe = param(T * C, std, r);                 /* position embedding   */
  m.blocks = [];
  for (var l = 0; l < L; l++) {
    var b = {
      ln1g: param(C, 0, r), ln1b: param(C, 0, r),
      qkv:  param(3 * C * C, std, r), qkvb: param(3 * C, 0, r),
      proj: param(C * C, std / Math.sqrt(2 * L), r), projb: param(C, 0, r),
      ln2g: param(C, 0, r), ln2b: param(C, 0, r),
      fc:   param(4 * C * C, std, r), fcb: param(4 * C, 0, r),
      fcp:  param(C * 4 * C, std / Math.sqrt(2 * L), r), fcpb: param(C, 0, r),
    };
    for (var i = 0; i < C; i++) { b.ln1g.v[i] = 1; b.ln2g.v[i] = 1; }
    m.blocks.push(b);
  }
  P.lnfg = param(C, 0, r); P.lnfb = param(C, 0, r);
  for (var j = 0; j < C; j++) P.lnfg.v[j] = 1;
  /* the output head shares the token embedding — standard weight tying,
     and it removes V*C parameters that a model this size cannot afford */
  m.params = collect(m);
  m.nparams = m.params.reduce(function (a, p) { return a + p.n; }, 0);
  return m;
}
function collect(m) {
  var out = [m.p.wte, m.p.wpe, m.p.lnfg, m.p.lnfb];
  m.blocks.forEach(function (b) {
    ['ln1g','ln1b','qkv','qkvb','proj','projb','ln2g','ln2b','fc','fcb','fcp','fcpb']
      .forEach(function (k) { out.push(b[k]); });
  });
  return out;
}

/* scratch buffers, allocated once per sequence length */
function scratch(m, T) {
  var C = m.cfg.dim, L = m.cfg.layers, H = m.cfg.heads, V = m.cfg.vocab;
  var s = { T: T, x: new Float32Array(T * C), logits: new Float32Array(T * V),
            probs: new Float32Array(T * V), dx: new Float32Array(T * C),
            dlogits: new Float32Array(T * V), blocks: [] };
  for (var l = 0; l < L; l++) {
    s.blocks.push({
      ln1: { mu: new Float32Array(T), inv: new Float32Array(T), xh: new Float32Array(T * C) },
      ln2: { mu: new Float32Array(T), inv: new Float32Array(T), xh: new Float32Array(T * C) },
      n1: new Float32Array(T * C), qkv: new Float32Array(T * 3 * C),
      att: new Float32Array(H * T * T), ao: new Float32Array(T * C),
      pj: new Float32Array(T * C), n2: new Float32Array(T * C),
      h: new Float32Array(T * 4 * C), hg: new Float32Array(T * 4 * C),
      mo: new Float32Array(T * C),
      d1: new Float32Array(T * C), d2: new Float32Array(T * C),
      dqkv: new Float32Array(T * 3 * C), datt: new Float32Array(H * T * T),
      dao: new Float32Array(T * C), dh: new Float32Array(T * 4 * C),
      dhg: new Float32Array(T * 4 * C),
    });
  }
  s.lnf = { mu: new Float32Array(T), inv: new Float32Array(T), xh: new Float32Array(T * C) };
  s.xf = new Float32Array(T * C); s.dxf = new Float32Array(T * C);
  return s;
}

function zero(a) { a.fill(0); }

/* ── forward · ids is Int32Array of length T ───────────────────────── */
function forward(m, ids, s) {
  var C = m.cfg.dim, H = m.cfg.heads, V = m.cfg.vocab, T = s.T, hs = C / H;
  var scale = 1 / Math.sqrt(hs), t, i, l, h, q, k;

  for (t = 0; t < T; t++)
    for (i = 0; i < C; i++)
      s.x[t * C + i] = m.p.wte.v[ids[t] * C + i] + m.p.wpe.v[t * C + i];

  for (l = 0; l < m.cfg.layers; l++) {
    var b = m.blocks[l], sb = s.blocks[l];
    lnFwd(s.x, b.ln1g.v, b.ln1b.v, T, C, sb.n1, sb.ln1);
    linearFwd(sb.n1, b.qkv.v, b.qkvb.v, T, C, 3 * C, sb.qkv);

    /* causal attention, head by head */
    for (h = 0; h < H; h++) {
      var ho = h * hs, ao = h * T * T;
      for (q = 0; q < T; q++) {
        var qb = q * 3 * C + ho, mx = -Infinity, sum = 0;
        for (k = 0; k <= q; k++) {
          var kb = k * 3 * C + C + ho, dot = 0;
          for (i = 0; i < hs; i++) dot += sb.qkv[qb + i] * sb.qkv[kb + i];
          dot *= scale;
          sb.att[ao + q * T + k] = dot;
          if (dot > mx) mx = dot;
        }
        for (k = 0; k <= q; k++) {
          var e = Math.exp(sb.att[ao + q * T + k] - mx);
          sb.att[ao + q * T + k] = e; sum += e;
        }
        for (k = 0; k <= q; k++) sb.att[ao + q * T + k] /= sum;
        for (k = q + 1; k < T; k++) sb.att[ao + q * T + k] = 0;
        for (i = 0; i < hs; i++) {
          var acc = 0;
          for (k = 0; k <= q; k++) acc += sb.att[ao + q * T + k] * sb.qkv[k * 3 * C + 2 * C + ho + i];
          sb.ao[q * C + ho + i] = acc;
        }
      }
    }
    linearFwd(sb.ao, b.proj.v, b.projb.v, T, C, C, sb.pj);
    for (i = 0; i < T * C; i++) s.x[i] += sb.pj[i];

    lnFwd(s.x, b.ln2g.v, b.ln2b.v, T, C, sb.n2, sb.ln2);
    linearFwd(sb.n2, b.fc.v, b.fcb.v, T, C, 4 * C, sb.h);
    geluFwd(sb.h, T * 4 * C, sb.hg);
    linearFwd(sb.hg, b.fcp.v, b.fcpb.v, T, 4 * C, C, sb.mo);
    for (i = 0; i < T * C; i++) s.x[i] += sb.mo[i];
  }

  lnFwd(s.x, m.p.lnfg.v, m.p.lnfb.v, T, C, s.xf, s.lnf);
  linearFwd(s.xf, m.p.wte.v, null, T, C, V, s.logits);   /* tied head */
  return s.logits;
}

/* ── loss · mean cross-entropy over the T positions ────────────────── */
function loss(m, s, targets) {
  var V = m.cfg.vocab, T = s.T, sum = 0;
  for (var t = 0; t < T; t++) {
    var o = t * V, mx = -Infinity, i;
    for (i = 0; i < V; i++) if (s.logits[o + i] > mx) mx = s.logits[o + i];
    var z = 0;
    for (i = 0; i < V; i++) { var e = Math.exp(s.logits[o + i] - mx); s.probs[o + i] = e; z += e; }
    for (i = 0; i < V; i++) s.probs[o + i] /= z;
    sum += -Math.log(Math.max(s.probs[o + targets[t]], 1e-12));
  }
  return sum / T;
}

/* ── backward · the derivative of exactly the forward above ────────── */
function backward(m, ids, s, targets) {
  var C = m.cfg.dim, H = m.cfg.heads, V = m.cfg.vocab, T = s.T, hs = C / H;
  var scale = 1 / Math.sqrt(hs), t, i, l, h, q, k;

  zero(s.dlogits); zero(s.dxf); zero(s.dx);
  for (t = 0; t < T; t++) {
    var o = t * V;
    for (i = 0; i < V; i++) s.dlogits[o + i] = s.probs[o + i] / T;
    s.dlogits[o + targets[t]] -= 1 / T;
  }
  /* tied head: gradient flows into wte from here AND from the embedding */
  linearBwd(s.xf, m.p.wte.v, s.dlogits, T, C, V, s.dxf, m.p.wte.g, null);
  lnBwd(s.dxf, m.p.lnfg.v, T, C, s.lnf, s.dx, m.p.lnfg.g, m.p.lnfb.g);

  for (l = m.cfg.layers - 1; l >= 0; l--) {
    var b = m.blocks[l], sb = s.blocks[l];
    zero(sb.d1); zero(sb.d2); zero(sb.dqkv); zero(sb.datt); zero(sb.dao);
    zero(sb.dh); zero(sb.dhg);

    /* mlp branch — residual means s.dx flows both around and through */
    linearBwd(sb.hg, b.fcp.v, s.dx, T, 4 * C, C, sb.dhg, b.fcp.g, b.fcpb.g);
    geluBwd(sb.h, sb.dhg, T * 4 * C, sb.dh);
    linearBwd(sb.n2, b.fc.v, sb.dh, T, C, 4 * C, sb.d2, b.fc.g, b.fcb.g);
    lnBwd(sb.d2, b.ln2g.v, T, C, sb.ln2, s.dx, b.ln2g.g, b.ln2b.g);

    /* attention branch */
    linearBwd(sb.ao, b.proj.v, s.dx, T, C, C, sb.dao, b.proj.g, b.projb.g);
    for (h = 0; h < H; h++) {
      var ho = h * hs, ao = h * T * T;
      for (q = 0; q < T; q++) {
        for (i = 0; i < hs; i++) {
          var d = sb.dao[q * C + ho + i];
          if (d === 0) continue;
          for (k = 0; k <= q; k++) {
            sb.datt[ao + q * T + k] += d * sb.qkv[k * 3 * C + 2 * C + ho + i];
            sb.dqkv[k * 3 * C + 2 * C + ho + i] += d * sb.att[ao + q * T + k];
          }
        }
        /* softmax jacobian along the row */
        var dot = 0;
        for (k = 0; k <= q; k++) dot += sb.datt[ao + q * T + k] * sb.att[ao + q * T + k];
        for (k = 0; k <= q; k++) {
          var ds = sb.att[ao + q * T + k] * (sb.datt[ao + q * T + k] - dot) * scale;
          for (i = 0; i < hs; i++) {
            sb.dqkv[q * 3 * C + ho + i]         += ds * sb.qkv[k * 3 * C + C + ho + i];
            sb.dqkv[k * 3 * C + C + ho + i]     += ds * sb.qkv[q * 3 * C + ho + i];
          }
        }
      }
    }
    linearBwd(sb.n1, b.qkv.v, sb.dqkv, T, C, 3 * C, sb.d1, b.qkv.g, b.qkvb.g);
    lnBwd(sb.d1, b.ln1g.v, T, C, sb.ln1, s.dx, b.ln1g.g, b.ln1b.g);
  }

  for (t = 0; t < T; t++)
    for (i = 0; i < C; i++) {
      m.p.wte.g[ids[t] * C + i] += s.dx[t * C + i];
      m.p.wpe.g[t * C + i]      += s.dx[t * C + i];
    }
}

function zeroGrads(m) { m.params.forEach(function (p) { p.g.fill(0); }); }

/* ── AdamW ─────────────────────────────────────────────────────────── */
function step(m, lr, t, wd, clip) {
  var b1 = 0.9, b2 = 0.95, eps = 1e-8;
  var norm = 0, i, j;
  for (i = 0; i < m.params.length; i++) {
    var g = m.params[i].g;
    for (j = 0; j < g.length; j++) norm += g[j] * g[j];
  }
  norm = Math.sqrt(norm);
  var sc = (clip && norm > clip) ? clip / norm : 1;
  var c1 = 1 - Math.pow(b1, t), c2 = 1 - Math.pow(b2, t);
  for (i = 0; i < m.params.length; i++) {
    var p = m.params[i];
    for (j = 0; j < p.n; j++) {
      var gg = p.g[j] * sc;
      p.m[j] = b1 * p.m[j] + (1 - b1) * gg;
      p.s[j] = b2 * p.s[j] + (1 - b2) * gg * gg;
      var mh = p.m[j] / c1, sh = p.s[j] / c2;
      p.v[j] -= lr * (mh / (Math.sqrt(sh) + eps) + (wd || 0) * p.v[j]);
    }
  }
  return norm;
}

/* ── sampling ──────────────────────────────────────────────────────── */
/* Generation calls this once per token with a context that grows until it
   hits the window, so a fresh scratch each time meant allocating and then
   throwing away megabytes per token — the loop spent most of its life in
   the garbage collector rather than in the matmuls. One cache per model,
   keyed by length, and the steady state reuses a single set of buffers. */
function scratchFor(m, T) {
  if (!m._sc) m._sc = {};
  if (!m._sc[T]) m._sc[T] = scratch(m, T);
  return m._sc[T];
}

function sample(m, ids, opts) {
  opts = opts || {};
  var T = Math.min(ids.length, m.cfg.ctx);
  var s = opts.scratch && opts.scratch.T === T ? opts.scratch : scratchFor(m, T);
  var seq = new Int32Array(T);
  for (var i = 0; i < T; i++) seq[i] = ids[ids.length - T + i];
  forward(m, seq, s);
  var V = m.cfg.vocab, o = (T - 1) * V;
  var temp = opts.temperature == null ? 0.8 : opts.temperature;
  var logits = new Float64Array(V);
  for (i = 0; i < V; i++) logits[i] = s.logits[o + i] / Math.max(temp, 1e-6);

  /* repetition penalty · inference only, and off unless asked for.
     A model this small falls into loops — the same token, then the same
     pair, forever. Dividing the score of tokens already in the recent
     window (the standard CTRL-style penalty) breaks the loop without
     touching a single trained weight, so nothing measured about the
     model changes. `repeatWindow` bounds how far back "recent" reaches. */
  var pen = opts.repeat;
  if (pen && pen !== 1) {
    var win = Math.min(opts.repeatWindow || 64, ids.length), seen = Object.create(null);
    for (i = ids.length - win; i < ids.length; i++) if (i >= 0) seen[ids[i]] = 1;
    for (var t2 in seen) { var id2 = +t2;
      logits[id2] = logits[id2] > 0 ? logits[id2] / pen : logits[id2] * pen; }
  }

  /* top-k, then top-p, then sample */
  var idx = Array.from({ length: V }, function (_, n) { return n; });
  idx.sort(function (a, b) { return logits[b] - logits[a]; });
  var K = Math.min(opts.topK || 40, V);
  idx = idx.slice(0, K);
  var mx = logits[idx[0]], z = 0, probs = [];
  idx.forEach(function (n) { var e = Math.exp(logits[n] - mx); probs.push(e); z += e; });
  for (i = 0; i < probs.length; i++) probs[i] /= z;
  var p = opts.topP == null ? 0.95 : opts.topP, cum = 0, cut = probs.length;
  for (i = 0; i < probs.length; i++) { cum += probs[i]; if (cum >= p) { cut = i + 1; break; } }
  idx = idx.slice(0, cut); probs = probs.slice(0, cut);
  var tot = probs.reduce(function (a, b) { return a + b; }, 0);
  var r = (opts.rand || Math.random)() * tot, acc = 0;
  for (i = 0; i < idx.length; i++) { acc += probs[i]; if (r <= acc) return idx[i]; }
  return idx[idx.length - 1];
}

/* ── save / load · plain JSON so the browser can read it ───────────── */
function save(m) {
  var out = { format: 'l8-nano-1', cfg: m.cfg, nparams: m.nparams, w: {} };
  var put = function (name, p) { out.w[name] = b64(p.v); };
  put('wte', m.p.wte); put('wpe', m.p.wpe); put('lnfg', m.p.lnfg); put('lnfb', m.p.lnfb);
  m.blocks.forEach(function (b, l) {
    ['ln1g','ln1b','qkv','qkvb','proj','projb','ln2g','ln2b','fc','fcb','fcp','fcpb']
      .forEach(function (k) { put(l + '.' + k, b[k]); });
  });
  return out;
}
function load(json) {
  var m = build(json.cfg, 1);
  var get = function (name, p) { unb64(json.w[name], p.v); };
  get('wte', m.p.wte); get('wpe', m.p.wpe); get('lnfg', m.p.lnfg); get('lnfb', m.p.lnfb);
  m.blocks.forEach(function (b, l) {
    ['ln1g','ln1b','qkv','qkvb','proj','projb','ln2g','ln2b','fc','fcb','fcp','fcpb']
      .forEach(function (k) { get(l + '.' + k, b[k]); });
  });
  return m;
}
/* float32 → base64. Four times smaller than JSON numbers and exact. */
function b64(f32) {
  var bytes = new Uint8Array(f32.buffer, f32.byteOffset, f32.byteLength);
  if (typeof Buffer !== 'undefined') return Buffer.from(bytes).toString('base64');
  var s = '';
  for (var i = 0; i < bytes.length; i++) s += String.fromCharCode(bytes[i]);
  return btoa(s);
}
function unb64(str, into) {
  var bytes;
  if (typeof Buffer !== 'undefined') bytes = Buffer.from(str, 'base64');
  else { var bin = atob(str); bytes = new Uint8Array(bin.length);
    for (var i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i); }
  var src = new Float32Array(bytes.buffer, bytes.byteOffset, bytes.byteLength / 4);
  into.set(src);
}

var NANO = { build: build, scratch: scratch, forward: forward, loss: loss, backward: backward,
  zeroGrads: zeroGrads, step: step, sample: sample, scratchFor: scratchFor, save: save, load: load, rng: rng,
  linearFwd: linearFwd, linearBwd: linearBwd, lnFwd: lnFwd, lnBwd: lnBwd,
  geluFwd: geluFwd, geluBwd: geluBwd };

if (typeof module !== 'undefined' && module.exports) module.exports = NANO;
if (typeof globalThis !== 'undefined') globalThis.L8Nano = NANO;