/* quant.js — the DMCRA level-set quantiser: a codebook chosen per block, from an elastic family.
 *
 * WHAT THIS IS, AND WHAT IT IS NOT
 * The DMCRA document proposes replacing fixed ternary `{-1,0,+1}` with an elastic family of grids
 * that "unclamps" as pressure demands. Stripped of the hardware vocabulary that a static page
 * cannot implement, that is one concrete, decidable mechanism: a weight matrix is cut into blocks,
 * each block is assigned a LEVEL SET from a nested family, and the assignment is chosen by a
 * measured criterion under a bit budget. This file is that mechanism and nothing else. There is no
 * routing fabric here, and no zero-bus: `0/0/0` is three zeros and this module reports, by exact
 * arithmetic, that it reconstructs identically to `{0}`.
 *
 * Two scale rules, on purpose:
 *   rule "meanabs" — scale = mean(|w|) per block, the convention this repo's weights.js already
 *                    uses for ternary. Reproduced here so a test can assert that this quantiser's
 *                    ternary is IDENTICAL to the already-verified `toTernary`/`fromTernary`. That
 *                    assertion is what stops this file from drifting away from the project.
 *   rule "lsq"     — least-squares optimal step per block (the classic LSQ iteration, Esser et al.
 *                    2020): s <- sum(w*q)/sum(q*q) with q = nearest level of w/s. Default for the
 *                    level-set ladder, because it gives EVERY grid its own optimal scale, so the
 *                    ladder compares codebooks rather than comparing one grid's scale artefact
 *                    against another's.
 *
 * Bit accounting is explicit and separated, because the interesting honest wrinkle is that they
 * disagree: `indexBits` is what an encoder must carry per weight, `effBits` adds the per-block scale
 * overhead (which is what makes small blocks expensive), and `entropyBits` is the measured Shannon
 * entropy of the level histogram. A grid that claims five channels but decodes to three distinct
 * values is reported as exactly that: `distinctLevels` < `levels.length`, `wastedChannels` > 0.
 */
(function (root, factory) {
  var Q = factory();
  if (typeof module !== "undefined" && module.exports) module.exports = Q;
  root.DMCRAQUANT = Q;
})(typeof globalThis !== "undefined" ? globalThis : this, function () {
  "use strict";

  /* ---------------------------------------------------------------- the elastic family
   * The document's family is nested: {+-1} -> {0,+-1} -> {0,+-1,+-2} -> ... . Two extra grids are
   * included because they are what an equal-bits contest actually needs: `quad` spends the SAME two
   * index bits as ternary but on two magnitudes instead of a zero, and `oct` spends the same three
   * bits as penta on four magnitudes instead of a zero. Without those two, "unclamping ternary"
   * would only ever be compared against a cheaper code, which would decide nothing.
   * `ternaryDup` is not a proposal; it is the document's `0/0/0` written down so it can be measured.
   */
  var GRIDS = {
    binary:    { levels: [0, 1],                note: "sign only, 1 index bit" },
    binaryPN:  { levels: [-1, 1],               note: "bipolar sign only, 1 index bit" },
    ternary:   { levels: [-1, 0, 1],            note: "BitNet b1.58, 2 index bits" },
    quad:      { levels: [-2, -1, 1, 2],        note: "2 magnitudes, no zero, 2 index bits" },
    ternaryDup:{ levels: [-1, 0, 0, 0, 1],      note: "the document's {0/0/0}: 5 channels, 3 distinct values" },
    penta:     { levels: [-2, -1, 0, 1, 2],     note: "3 index bits" },
    hepta:     { levels: [-3, -2, -1, 0, 1, 2, 3], note: "3 index bits, same cost as penta" },
    oct:       { levels: [-4, -3, -2, -1, 1, 2, 3, 4], note: "4 magnitudes, no zero, 3 index bits" },
    nona:      { levels: [-4, -3, -2, -1, 0, 1, 2, 3, 4], note: "4 index bits" }
  };

  var SCALE_BYTES = 4;

  function gridOf(grid) {
    if (typeof grid === "string") {
      if (!GRIDS[grid]) throw new Error("unknown grid: " + grid);
      return GRIDS[grid].levels.slice();
    }
    if (Object.prototype.toString.call(grid) === "[object Array]") {
      if (!grid.length) throw new Error("empty level set");
      return grid.slice();
    }
    throw new Error("grid must be a known name or an array of levels");
  }

  function indexBitsFor(k) { return Math.max(1, Math.ceil(Math.log2(k))); }

  /* Distinct DECODED values. This is what makes `ternaryDup` measurable rather than rhetorical:
   * three zero channels decode to one value, so the encoding is wider than the codebook. */
  function distinctValues(levels) {
    var seen = Object.create(null), n = 0, i, v;
    for (i = 0; i < levels.length; i++) { v = levels[i]; if (seen[v] !== 1) { seen[v] = 1; n++; } }
    return n;
  }

  function gridInfo(grid) {
    var levels = gridOf(grid);
    var maxLevel = 0, i;
    for (i = 0; i < levels.length; i++) if (Math.abs(levels[i]) > maxLevel) maxLevel = Math.abs(levels[i]);
    var d = distinctValues(levels);
    var idx = indexBitsFor(levels.length);
    var dec = indexBitsFor(d);
    return {
      levels: levels, count: levels.length, maxLevel: maxLevel,
      distinctLevels: d, distinctBits: dec,
      indexBits: idx, decodeBits: dec,
      wastedChannels: levels.length - d,
      /* Channels paid for that the decoder cannot distinguish. Zero for every honest grid. */
      wastedBits: idx - dec
    };
  }

  /* ---------------------------------------------------------------- nearest level */

  function isContiguousIntegers(levels) {
    for (var i = 1; i < levels.length; i++) if (levels[i] !== levels[i - 1] + 1) return false;
    return true;
  }

  /* Nearest level to x, given that the levels were built by scaling integers. Ties resolve to the
   * first (lowest) candidate, which is why three zero channels can never disagree: every one of
   * them is 0, so the reconstructed value is 0 whichever way the tie falls. */
  function nearestIndex(levels, x, contiguous, lo, hi) {
    if (contiguous) {
      var r = Math.round(x);
      if (r < lo) r = lo; else if (r > hi) r = hi;
      return r - lo;
    }
    var best = 0, bestD = Infinity, i, d;
    for (i = 0; i < levels.length; i++) {
      d = Math.abs(x - levels[i]);
      if (d < bestD) { bestD = d; best = i; }
    }
    return best;
  }

  /* ---------------------------------------------------------------- block statistics */

  /* Per-block moments. `variance` is the document's stated pressure; it is computed here so the
   * allocation experiment can test whether it is the RIGHT pressure, and `mse` per grid is kept so
   * the benefit-driven allocator can allocate by measured error reduction instead. */
  function blockStats(values, blockSize, gridNames, rule) {
    var n = values.length;
    var bs = blockSize > 0 ? blockSize : n;
    var out = [], off;
    for (off = 0; off < n; off += bs) {
      var len = Math.min(bs, n - off);
      var sum = 0, absSum = 0, sq = 0, i, v;
      for (i = 0; i < len; i++) { v = values[off + i]; sum += v; absSum += Math.abs(v); sq += v * v; }
      var mean = sum / len;
      var variance = Math.max(0, sq / len - mean * mean);
      var rec = { offset: off, len: len, mean: mean, meanAbs: absSum / len,
                  variance: variance, sd: Math.sqrt(variance), rms: Math.sqrt(sq / len), mse: {} };
      if (gridNames) {
        for (i = 0; i < gridNames.length; i++) {
          var q = quantiseBlock(values, off, len, gridNames[i], rule);
          rec.mse[gridNames[i]] = q.mse;
        }
      }
      out.push(rec);
    }
    return out;
  }

  /* ---------------------------------------------------------------- the scale step */

  function meanAbsOf(values, off, len) {
    var s = 0, i;
    for (i = 0; i < len; i++) s += Math.abs(values[off + i]);
    return s / len;
  }

  /* LSQ step: iterate s <- sum(w*q)/sum(q*q) until it stops moving. This is the same update used in
   * learned step-size quantization, run here as a few fixed-point iterations per block. */
  function lsqScale(values, off, len, levels, contiguous, lo, hi, seed) {
    var s = seed > 0 ? seed : meanAbsOf(values, off, len);
    if (!(s > 0)) {
      /* An all-zero block: any scale reconstructs it exactly. */
      return 0;
    }
    var it, i, num, den, q, w, s2;
    for (it = 0; it < 6; it++) {
      num = 0; den = 0;
      for (i = 0; i < len; i++) {
        w = values[off + i];
        q = levels[nearestIndex(levels, w / s, contiguous, lo, hi)];
        num += w * q; den += q * q;
      }
      if (den <= 0) break;
      s2 = num / den;
      if (!(s2 > 0)) break;
      if (Math.abs(s2 - s) <= 1e-12 * s) { s = s2; break; }
      s = s2;
    }
    return s;
  }

  /* ---------------------------------------------------------------- one block */

  /* Quantise one block under one grid. Returns the reconstruction, the scale, and both error and
   * bit accounts. `rule` is "lsq" (default) or "meanabs" (the weights.js ternary convention). */
  function quantiseBlock(values, off, len, grid, rule) {
    var info = gridInfo(grid);
    var levels = info.levels;
    var contiguous = isContiguousIntegers(levels);
    var lo = contiguous ? levels[0] : 0, hi = contiguous ? levels[levels.length - 1] : 0;
    var s;
    if (rule === "meanabs") s = meanAbsOf(values, off, len);
    else s = lsqScale(values, off, len, levels, contiguous, lo, hi, meanAbsOf(values, off, len));
    var rec = new Float32Array(len);
    var counts = new Array(levels.length);
    for (var c = 0; c < counts.length; c++) counts[c] = 0;
    var sse = 0, sae = 0, maxAbs = 0, i, w, qi, r, d;
    for (i = 0; i < len; i++) {
      w = values[off + i];
      qi = s > 0 ? nearestIndex(levels, w / s, contiguous, lo, hi) : 0;
      counts[qi]++;
      r = s > 0 ? levels[qi] * s : 0;
      rec[i] = r;
      d = w - r; sse += d * d; sae += Math.abs(d);
      if (Math.abs(d) > maxAbs) maxAbs = Math.abs(d);
    }
    return {
      grid: typeof grid === "string" ? grid : "custom", levels: levels, scale: s,
      values: rec, counts: counts,
      mse: len ? sse / len : 0, mae: len ? sae / len : 0, maxAbs: maxAbs,
      bits: blockBits(len, s, info, counts)
    };
  }

  /* Bit account for one block. `indexBits` travels per weight; the scale is one 4-byte float per
   * block, so its share is 32/len bits per weight — the term that punishes small blocks. */
  function blockBits(len, scale, info, counts) {
    var indexBits = info.indexBits;
    var scaleBits = len ? (SCALE_BYTES * 8) / len : 0;
    var entropy = 0, total = len || 1, j, p;
    if (counts) {
      for (j = 0; j < counts.length; j++) {
        if (!counts[j]) continue;
        p = counts[j] / total;
        entropy -= p * Math.log2(p);
      }
    }
    return {
      indexBits: indexBits, scaleBits: scaleBits, effBits: indexBits + scaleBits,
      entropyBits: entropy + (scale ? scaleBits : 0),
      decodeBits: info.decodeBits, wastedBits: info.wastedBits,
      wastedChannels: info.wastedChannels, distinctLevels: info.distinctLevels,
      channels: info.count, scaleZero: !(scale > 0)
    };
  }

  /* ---------------------------------------------------------------- whole tensor */

  /* Quantise a whole array under one grid. Returns the reconstruction plus an aggregate bit account
   * that is the length-weighted mean of the per-block accounts, so a caller can compare two grids
   * at equal bits rather than at equal nominal precision. */
  function quantise(values, opts) {
    opts = opts || {};
    var grid = opts.grid === undefined ? "ternary" : opts.grid;
    var rule = opts.rule === "meanabs" ? "meanabs" : "lsq";
    var bs = opts.blockSize | 0;
    var n = values.length;
    var size = bs > 0 ? bs : n;
    var out = new Float32Array(n);
    var blocks = [], off, acc = null;
    for (off = 0; off < n; off += size) {
      var len = Math.min(size, n - off);
      var q = quantiseBlock(values, off, len, grid, rule);
      out.set(q.values, off);
      blocks.push({ offset: off, len: len, scale: q.scale, mse: q.mse, bits: q.bits,
                    counts: q.counts, levels: q.levels });
      acc = acc ? acc : {
        indexBits: 0, scaleBits: 0, effBits: 0, entropyBits: 0, decodeBits: 0,
        wastedBits: 0, wastedChannels: q.bits.wastedChannels, distinctLevels: q.bits.distinctLevels,
        channels: q.bits.channels, grid: q.grid, name: typeof grid === "string" ? grid : "custom"
      };
      acc.indexBits += q.bits.indexBits * len / n;
      acc.scaleBits += q.bits.scaleBits * len / n;
      acc.effBits += q.bits.effBits * len / n;
      acc.entropyBits += q.bits.entropyBits * len / n;
      acc.decodeBits += q.bits.decodeBits * len / n;
      acc.wastedBits += q.bits.wastedBits * len / n;
    }
    if (!acc) {
      acc = { indexBits: 0, scaleBits: 0, effBits: 0, entropyBits: 0, decodeBits: 0,
              wastedBits: 0, wastedChannels: 0, distinctLevels: 0, channels: 0,
              grid: String(grid), name: typeof grid === "string" ? grid : "custom" };
    }
    acc.blockSize = bs; acc.rule = rule; acc.count = n;
    acc.storedBytes = Math.ceil(n * acc.indexBits / 8) + blocks.length * SCALE_BYTES;
    acc.floatBytes = n * 4;
    acc.ratio = acc.storedBytes ? acc.floatBytes / acc.storedBytes : 0;
    var err = measure(values, out);
    acc.mse = err.mse; acc.rmse = err.rmse; acc.mae = err.mae; acc.maxAbs = err.maxAbs;
    return { values: out, blocks: blocks, bits: acc, error: err };
  }

  /* ---------------------------------------------------------------- error */

  function measure(orig, rec) {
    var n = Math.min(orig.length, rec.length), i, d, sse = 0, sae = 0, maxAbs = 0, e0 = 0;
    for (i = 0; i < n; i++) {
      d = orig[i] - rec[i];
      sse += d * d; sae += Math.abs(d);
      if (Math.abs(d) > maxAbs) maxAbs = Math.abs(d);
      e0 += orig[i] * orig[i];
    }
    return {
      count: n, mse: n ? sse / n : 0, rmse: n ? Math.sqrt(sse / n) : 0,
      mae: n ? sae / n : 0, maxAbs: maxAbs,
      /* Fraction of the original energy the reconstruction retains; 1 means nothing was lost. */
      energyRetained: e0 > 0 ? 1 - sse / e0 : 1
    };
  }

  /* Exact equality of two reconstructions, used by the {0/0/0} identity test. Compares bit patterns
   * rather than a tolerance, because the claim under test is identity, not closeness. */
  function identical(a, b) {
    if (a.length !== b.length) return false;
    for (var i = 0; i < a.length; i++) if (a[i] !== b[i]) return false;
    return true;
  }

  /* ---------------------------------------------------------------- elastic allocation */

  /* The allocator. This is the DMCRA mechanism made decidable: given a bit budget, assign each block
   * a grid so that measured reconstruction error is as low as possible.
   *
   *   key "benefit"  (default) — greedy by MEASURED error reduction per extra bit, using each
   *                              block's per-grid MSE. The "pressure" is benefit/cost, not variance.
   *   key "variance" — the document's stated rule: rank blocks by variance and unclamp the most
   *                    dispersed ones first, ignoring how much each block's error would actually
   *                    improve. Implemented faithfully so it can be compared, not assumed inferior.
   *
   * Budget is an average of `indexBits` per weight. The greedy walks upgrades from the cheapest
   * candidate upward and stops when the next cheapest upgrade no longer fits, so the realised
   * average is <= budget and, when the budget happens to be reachable exactly, equals it.
   */
  function allocate(values, opts) {
    opts = opts || {};
    var candidates = (opts.candidates || ["binary", "ternary", "penta", "hepta"]).slice();
    var rule = opts.rule === "meanabs" ? "meanabs" : "lsq";
    var budget = opts.budget === undefined ? 2 : opts.budget;
    var key = opts.key === "variance" ? "variance" : "benefit";
    var bs = opts.blockSize | 0;
    var n = values.length;
    var size = bs > 0 ? bs : n;

    candidates.sort(function (a, b) { return gridInfo(a).indexBits - gridInfo(b).indexBits; });
    var stats = blockStats(values, size, candidates, rule);
    var nBlocks = stats.length;
    if (!nBlocks) {
      return { grids: [], stats: [], avgBits: 0, bits: null, values: new Float32Array(0),
               error: measure(values, new Float32Array(0)), key: key, candidates: candidates };
    }

    /* Costs, and the exact error reduction each single-step upgrade would buy. */
    var cost = candidates.map(function (c) { return gridInfo(c).indexBits; });
    var level = stats.map(function () { return 0; });      // index into candidates
    var totalWeight = 0;
    stats.forEach(function (s) { totalWeight += s.len; });

    function currentBits() {
      var b = 0;
      for (var i = 0; i < nBlocks; i++) b += cost[level[i]] * stats[i].len / totalWeight;
      return b;
    }

    /* A plain greedy on single-step upgrades. Unclamping is monotone: a block only moves to a
     * costlier grid, and only when that grid actually reduces its measured error, so a grid that
     * costs more without helping is never selected. Each pass takes the single best-scoring upgrade
     * that still fits the budget, which is the standard greedy for this separable problem. The two
     * keys differ only in the score: benefit uses measured error reduction per extra bit, variance
     * uses the block's own dispersion, which is the rule the document names. */
    var guard = 0;
    while (guard++ < nBlocks * candidates.length + 10) {
      var best = -1, bestScore = -Infinity, i, next, gain, dCost, score;
      for (i = 0; i < nBlocks; i++) {
        next = level[i] + 1;
        if (next >= candidates.length) continue;
        gain = (stats[i].mse[candidates[level[i]]] - stats[i].mse[candidates[next]]) * stats[i].len;
        if (gain <= 0) continue;
        dCost = (cost[next] - cost[level[i]]) * stats[i].len / totalWeight;
        if (dCost > 0 && currentBits() + dCost > budget + 1e-9) continue;
        if (key === "variance") score = stats[i].variance + gain * 1e-18;
        else score = dCost > 0 ? gain / dCost : gain;
        if (score > bestScore) { bestScore = score; best = i; }
      }
      if (best < 0) break;
      level[best]++;
    }

    /* Realise the assignment. */
    var out = new Float32Array(n);
    var grids = [], off = 0, i2;
    var accBits = { indexBits: 0, scaleBits: 0, effBits: 0, entropyBits: 0, storedBytes: 0 };
    for (i2 = 0; i2 < nBlocks; i2++) {
      var st = stats[i2];
      var name = candidates[level[i2]];
      var q = quantiseBlock(values, st.offset, st.len, name, rule);
      out.set(q.values, st.offset);
      grids.push(name);
      accBits.indexBits += q.bits.indexBits * st.len / n;
      accBits.scaleBits += q.bits.scaleBits * st.len / n;
      accBits.effBits += q.bits.effBits * st.len / n;
      accBits.entropyBits += q.bits.entropyBits * st.len / n;
      accBits.storedBytes += Math.ceil(st.len * q.bits.indexBits / 8) + SCALE_BYTES;
    }
    accBits.count = n; accBits.blockSize = bs; accBits.rule = rule;
    accBits.mode = "elastic"; accBits.key = key; accBits.budget = budget;
    accBits.storedBytes = Math.ceil(accBits.storedBytes);
    accBits.floatBytes = n * 4;
    accBits.ratio = accBits.storedBytes ? accBits.floatBytes / accBits.storedBytes : 0;
    var err2 = measure(values, out);
    accBits.mse = err2.mse; accBits.rmse = err2.rmse; accBits.mae = err2.mae;
    accBits.maxAbs = err2.maxAbs;
    var hist = Object.create(null);
    grids.forEach(function (g) { hist[g] = (hist[g] || 0) + 1; });
    accBits.gridHistogram = hist;
    return { values: out, grids: grids, stats: stats, bits: accBits, error: err2,
             avgBits: accBits.indexBits, key: key, candidates: candidates,
             /* The ladder the allocator was allowed to climb, for reporting. */
             costs: candidates.map(function (c) { return { grid: c, indexBits: gridInfo(c).indexBits }; }) };
  }

  /* ---------------------------------------------------------------- the {0/0/0} experiment */

  /* Everything needed to settle the document's multi-channel-centroid claim numerically: the same
   * weights quantised with one zero channel and with three, the reconstruction compared for exact
   * inequality, and the channel accounting of each. A three-channel zero that decodes identically
   * is not an information channel; the caller can print both rows and see that. */
  function zeroChannelTest(values, opts) {
    opts = opts || {};
    var bs = opts.blockSize | 0;
    var one = quantise(values, { grid: "ternary", blockSize: bs, rule: opts.rule });
    var three = quantise(values, { grid: "ternaryDup", blockSize: bs, rule: opts.rule });
    return {
      one: one, three: three,
      identical: identical(one.values, three.values),
      info: { ternary: gridInfo("ternary"), ternaryDup: gridInfo("ternaryDup") }
    };
  }

  /* ---------------------------------------------------------------- HDC-side helpers */

  /* Shannon entropy in bits of a categorical histogram, used by the bundling experiment to compare
   * multi-channel codes at equal information content rather than equal channel count. */
  function entropy(counts) {
    var total = 0, i;
    for (i = 0; i < counts.length; i++) total += counts[i];
    if (!total) return 0;
    var h = 0, p;
    for (i = 0; i < counts.length; i++) {
      if (!counts[i]) continue;
      p = counts[i] / total;
      h -= p * Math.log2(p);
    }
    return h;
  }

  return {
    GRIDS: GRIDS, gridInfo: gridInfo, gridOf: gridOf,
    quantise: quantise, quantiseBlock: quantiseBlock, blockStats: blockStats,
    allocate: allocate, zeroChannelTest: zeroChannelTest,
    measure: measure, identical: identical, entropy: entropy,
    indexBitsFor: indexBitsFor, SCALE_BYTES: SCALE_BYTES
  };
});