// AksharaLP v3 tiny sentence LM (JS port of bhc/nlm.py): Viterbi unigram-piece tokenizer + 2-layer LSTMP + tied embedding.
// File: u32 header length | JSON {cfg, bits, G, tensors:[[name, shape, kind, off, len]], vocab:[[piece, score]]} | data (4-byte aligned).
const SP = "\u2581";
function f16(h) {
  const s = h & 0x8000 ? -1 : 1, e = (h >> 10) & 0x1f, f = h & 0x3ff;
  if (e === 0) return s * f * 2 ** -24; if (e === 31) return f ? NaN : s * Infinity; return s * (1 + f / 1024) * 2 ** (e - 15);
}
function dq(data, off, [r, c], kind, G) {
  const out = new Float32Array(r * c), ng = Math.ceil(c / G), dv = new DataView(data.buffer, data.byteOffset);
  let q;
  if (kind === "q8") { q = new Int8Array(data.buffer, data.byteOffset + off, r * c); off += r * c; }
  else { const nb = (r * c + 1) >> 1; q = new Int8Array(nb * 2); for (let i = 0; i < nb; i++) { const u = data[off + i]; q[2 * i] = (u >> 4) - 8; q[2 * i + 1] = (u & 15) - 8; } off += nb; }
  off += (4 - (off % 4)) % 4;
  for (let i = 0; i < r; i++) for (let g = 0; g < ng; g++) {
    const sc = f16(dv.getUint16(off + 2 * (i * ng + g), true)), e = Math.min(c, (g + 1) * G);
    for (let j = g * G; j < e; j++) out[i * c + j] = q[i * c + j] * sc;
  }
  return out;
}
export class Tok {
  constructor(vocab) {
    this.score = new Map(); this.id = new Map(); this.maxlen = 0; let mn = Infinity;
    vocab.forEach(([p, s], i) => { this.score.set(p, s); this.id.set(p, i); this.maxlen = Math.max(this.maxlen, [...p].length); mn = Math.min(mn, s); });
    this.unk = this.id.get("<unk>"); this.eos = this.id.get("</s>"); this.unkScore = mn - 10; this.cache = new Map();
  }
  encodeWord(w) {
    let r = this.cache.get(w); if (r) return r;
    if (this.id.has(w) && w.startsWith("<")) { r = [this.id.get(w)]; this.cache.set(w, r); return r; }
    const s = [...(SP + w)], n = s.length, best = new Float64Array(n + 1).fill(-1e30), back = new Array(n + 1).fill(null); best[0] = 0;
    for (let i = 0; i < n; i++) {
      if (best[i] <= -1e29) continue; let hit = false;
      for (let j = i + 1; j <= Math.min(n, i + this.maxlen); j++) {
        const p = s.slice(i, j).join(""), sc = this.score.get(p);
        if (sc !== undefined && p !== "<unk>" && p !== "<s>" && p !== "</s>" && p !== "<n>") {
          hit = true; if (best[i] + sc > best[j]) { best[j] = best[i] + sc; back[j] = [i, this.id.get(p)]; }
        }
      }
      if (!hit || back[i + 1] === null) if (best[i] + this.unkScore > best[i + 1]) { best[i + 1] = best[i] + this.unkScore; back[i + 1] = [i, this.unk]; }
    }
    const out = []; for (let j = n; j > 0;) { const [i, pid] = back[j]; out.push(pid); j = i; }
    r = out.reverse(); if (this.cache.size < 200000) this.cache.set(w, r); return r;
  }
}
export class NLM {
  constructor(buf, norm) {
    const b = Uint8Array.from(buf instanceof ArrayBuffer ? new Uint8Array(buf) : buf), n = new DataView(b.buffer, b.byteOffset).getUint32(0, true);
    const meta = JSON.parse(new TextDecoder().decode(b.subarray(4, 4 + n))), data = b.subarray(4 + n), W = {};
    for (const [k, shape, kind, off] of meta.tensors) {
      const cnt = shape.reduce((a, x) => a * x, 1);
      W[k] = kind === "f32" ? new Float32Array(Uint8Array.prototype.slice.call(data, off, off + 4 * cnt).buffer) : dq(data, off, shape, kind, meta.G);
    }
    this.cfg = meta.cfg; this.W = W; this.tok = new Tok(meta.vocab); this.L = meta.cfg.layers; this.norm = norm || meta.cfg.norm || "exact";
    this.E = meta.cfg.emb; this.H = meta.cfg.hid; this.V = meta.vocab.length; this.emb = W["emb.weight"]; this.ob = W["out_bias"]; this._init = null;
  }
  static async load(bytes, norm) { return new NLM(bytes, norm); }
  // state = {h: [Float32Array(E) x L], c: [Float32Array(H) x L]}
  _gates(l, x, h, out) {            // out[4H] = Wih x + Whh h + b_ih + b_hh  (x may be null -> skip input part)
    const W = this.W, E = this.E, H4 = 4 * this.H, wih = W[`lstm.weight_ih_l${l}`], whh = W[`lstm.weight_hh_l${l}`], bi = W[`lstm.bias_ih_l${l}`], bh = W[`lstm.bias_hh_l${l}`];
    for (let r = 0, o = 0; r < H4; r++, o += E) {
      let a0 = 0, a1 = 0, a2 = 0, a3 = 0, i = 0;
      if (x && h) for (; i + 3 < E; i += 4) {
        a0 += wih[o + i] * x[i] + whh[o + i] * h[i]; a1 += wih[o + i + 1] * x[i + 1] + whh[o + i + 1] * h[i + 1];
        a2 += wih[o + i + 2] * x[i + 2] + whh[o + i + 2] * h[i + 2]; a3 += wih[o + i + 3] * x[i + 3] + whh[o + i + 3] * h[i + 3];
      } else if (x) for (; i + 3 < E; i += 4) { a0 += wih[o + i] * x[i]; a1 += wih[o + i + 1] * x[i + 1]; a2 += wih[o + i + 2] * x[i + 2]; a3 += wih[o + i + 3] * x[i + 3]; }
      else for (; i + 3 < E; i += 4) { a0 += whh[o + i] * h[i]; a1 += whh[o + i + 1] * h[i + 1]; a2 += whh[o + i + 2] * h[i + 2]; a3 += whh[o + i + 3] * h[i + 3]; }
      let a = out[r] + ((a0 + a1) + (a2 + a3));
      for (; i < E; i++) { if (x) a += wih[o + i] * x[i]; if (h) a += whh[o + i] * h[i]; }
      if (x) a += bi[r] + bh[r];
      out[r] = a;
    }
    return out;
  }
  _cell(l, g, c) {                  // -> {h, c}
    const H = this.H, E = this.E, whr = this.W[`lstm.weight_hr_l${l}`], cn = new Float32Array(H), m = new Float32Array(H), h = new Float32Array(E);
    const sg = (v) => 1 / (1 + Math.exp(-v));
    for (let j = 0; j < H; j++) { cn[j] = sg(g[H + j]) * c[j] + sg(g[j]) * Math.tanh(g[2 * H + j]); m[j] = sg(g[3 * H + j]) * Math.tanh(cn[j]); }
    for (let r = 0, o = 0; r < E; r++, o += H) {
      let a0 = 0, a1 = 0, a2 = 0, a3 = 0, j = 0;
      for (; j + 3 < H; j += 4) { a0 += whr[o + j] * m[j]; a1 += whr[o + j + 1] * m[j + 1]; a2 += whr[o + j + 2] * m[j + 2]; a3 += whr[o + j + 3] * m[j + 3]; }
      let s = (a0 + a1) + (a2 + a3); for (; j < H; j++) s += whr[o + j] * m[j]; h[r] = s;
    }
    return { h, c: cn };
  }
  _xg0(id) {                        // cached layer-0 input projection per piece
    this._xc = this._xc || new Map(); let g = this._xc.get(id); if (g) return g;
    g = this._gates(0, this.emb.subarray(id * this.E, (id + 1) * this.E), null, new Float32Array(4 * this.H));
    if (this._xc.size > 4000) this._xc.clear(); this._xc.set(id, g); return g;
  }
  step(st, id, hh0) {              // hh0: optional precomputed Whh0 h0
    const hs = [], cs = []; let x = null;
    for (let l = 0; l < this.L; l++) {
      let g;
      if (l === 0) { g = Float32Array.from(this._xg0(id)); const r = hh0 || this._gates(0, null, st.h[0], new Float32Array(4 * this.H)); for (let i = 0; i < g.length; i++) g[i] += r[i]; }
      else g = this._gates(l, x, st.h[l], new Float32Array(4 * this.H));
      const o = this._cell(l, g, st.c[l]); hs.push(o.h); cs.push(o.c); x = o.h;
    }
    return { h: hs, c: cs };
  }
  logZ(st) {
    if (this.norm !== "exact") return 0;
    const top = st.h[this.L - 1], E = this.E, lg = new Float64Array(this.V); let m = -Infinity;
    for (let v = 0; v < this.V; v++) { let s = this.ob[v]; const o = v * E; for (let i = 0; i < E; i++) s += this.emb[o + i] * top[i]; lg[v] = s; if (s > m) m = s; }
    let z = 0; for (let v = 0; v < this.V; v++) z += Math.exp(lg[v] - m); return m + Math.log(z);
  }
  logit(st, id) { const top = st.h[this.L - 1], E = this.E, o = id * E; let s = this.ob[id]; for (let i = 0; i < E; i++) s += this.emb[o + i] * top[i]; return s; }
  initState() {
    if (!this._init) { const z = { h: Array.from({ length: this.L }, () => new Float32Array(this.E)), c: Array.from({ length: this.L }, () => new Float32Array(this.H)) }; this._init = this.step(z, this.tok.eos); }
    return this._init;
  }
  /** states[B], words[K] -> {lp[B][K], ns[B][K]} (same math as bhc/nlm.py score_words) */
  scoreWords(states, words) {
    const seqs = words.map((w) => this.tok.encodeWord(w)), lp = [], ns = [];
    for (const st of states) {
      const z = this.logZ(st), hh0 = this._gates(0, null, st.h[0], new Float32Array(4 * this.H)), row = [], nrow = [];
      for (const q of seqs) {
        let s = this.logit(st, q[0]) - z, cur = this.step(st, q[0], hh0);
        for (let t = 1; t < q.length; t++) { s += this.logit(cur, q[t]) - this.logZ(cur); cur = this.step(cur, q[t]); }
        row.push(s); nrow.push(cur);
      }
      lp.push(row); ns.push(nrow);
    }
    return { lp, ns };
  }
  eosLogp(st) { return this.logit(st, this.tok.eos) - this.logZ(st); }
  sentLogp(words) { let st = this.initState(), tot = 0; for (const w of words) { const { lp, ns } = this.scoreWords([st], [w]); tot += lp[0][0]; st = ns[0][0]; } return tot + this.eosLogp(st); }
}