// Independent (second-language) reference for the Mixed Matrix Commitment Scheme (MMCS) that
// Plonky3/SP1 use over BabyBear, component (c) of the PQ STARK-verify precompile 0x0AE8. Mirrors
// mmcs_babybear_reference.py and the standalone Java MmcsBabyBearSelfTest. It serves cross-language
// agreement (three independent implementations produce byte-identical output) AND conformance (it
// reproduces real known-answer vectors emitted by the pinned Plonky3 crates). It verifies NOTHING
// about a real proof and the top-level 0x0AE8 verifier stays fail-closed. See
// docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 4.
//
// HONEST SCOPE (CONFIRMED 2026-07-19): the construction is the exact SP1 inner config, verbatim from
// p3-merkle-tree's own mmcs.rs tests:
//   Perm       = Poseidon2<BabyBear, ..., 16, 7>
//   MyHash     = PaddingFreeSponge<Perm, 16, 8, 8>
//   MyCompress = TruncatedPermutation<Perm, 2, 8, 16>
//   MyMmcs     = FieldMerkleTreeMmcs<Packing, Packing, MyHash, MyCompress, 8>   (DIGEST_ELEMS = 8)
// and its root/opening/verify outputs reproduce real known-answer vectors from the pinned
// p3-merkle-tree / p3-symmetric / p3-commit 0.4.3-succinct crates (checksums matched; see
// spec-mmcs-babybear.md). Leaf hashing/compression use the CONFIRMED Poseidon2-BabyBear permutation.
//
// Uses BigInt throughout (the permutation exceeds 2^53). Run standalone to emit JSON:
//   node mmcs_babybear_reference.mjs

import { permute as p2permute, P as P_BIG } from "./poseidon2_babybear_reference.mjs";

export const P = P_BIG;
export const WIDTH = 16;
export const RATE = 8;
export const OUT = 8;
export const DIGEST_ELEMS = 8;
const DEFAULT_DIGEST = Array(DIGEST_ELEMS).fill(0n);

const mod = (a) => ((a % P) + P) % P;

// ---- primitives ----
export function permute(state) {
  return p2permute(state);
}

// PaddingFreeSponge<Perm, 16, 8, 8>: overwrite-mode, padding-free sponge (p3-symmetric sponge.rs).
export function hashIter(elems) {
  let state = Array(WIDTH).fill(0n);
  const e = elems.map((x) => mod(BigInt(x)));
  for (let i = 0; i < e.length; i += RATE) {
    const chunk = e.slice(i, i + RATE);
    for (let j = 0; j < chunk.length; j++) state[j] = chunk[j];
    state = permute(state);
  }
  return state.slice(0, OUT);
}

export const hashSlice = (elems) => hashIter([...elems]);
export const hashItem = (x) => hashIter([x]);

// TruncatedPermutation<Perm, 2, 8, 16>: permute(left||right) truncated to 8 lanes.
export function compress2to1(left8, right8) {
  if (left8.length !== DIGEST_ELEMS || right8.length !== DIGEST_ELEMS) {
    throw new Error("compress inputs must be length 8");
  }
  const pre = [...left8.map((x) => mod(BigInt(x))), ...right8.map((x) => mod(BigInt(x)))];
  return permute(pre).slice(0, DIGEST_ELEMS);
}

// ---- helpers ----
function log2Ceil(n) {
  if (n <= 1) return 0;
  return BigInt(n - 1).toString(2).length;
}
function nextPow2(n) {
  if (n <= 1) return 1;
  return 1 << (BigInt(n - 1).toString(2).length);
}

export class Matrix {
  constructor(rows) {
    this.rows = rows.map((r) => r.map((x) => mod(BigInt(x))));
    this.height = rows.length;
    this.width = rows.length ? rows[0].length : 0;
  }
  row(i) {
    return [...this.rows[i]];
  }
}

// ---- tree build (FieldMerkleTree::new) ----
export function buildTree(matrices) {
  if (!matrices.length) throw new Error("no matrices");
  const heights = matrices.map((m) => m.height).sort((a, b) => a - b);
  for (let k = 1; k < heights.length; k++) {
    if (heights[k - 1] !== heights[k] && nextPow2(heights[k - 1]) === nextPow2(heights[k])) {
      throw new Error("matrix heights that round up to the same power of two must be equal");
    }
  }
  // stable sort tallest-first
  const order = matrices.map((_, i) => i).sort((a, b) => matrices[b].height - matrices[a].height || a - b);
  const ordered = order.map((i) => matrices[i]);

  const maxHeight = ordered[0].height;
  const logMaxHeight = log2Ceil(maxHeight);
  const maxHeightPadded = nextPow2(maxHeight);

  let idx = 0;
  const tallest = [];
  while (idx < ordered.length && ordered[idx].height === maxHeight) tallest.push(ordered[idx++]);

  const layer0 = [];
  for (let i = 0; i < maxHeightPadded; i++) layer0.push([...DEFAULT_DIGEST]);
  for (let i = 0; i < maxHeight; i++) {
    const row = [];
    for (const m of tallest) row.push(...m.row(i));
    layer0[i] = hashIter(row);
  }
  const digestLayers = [layer0];

  while (digestLayers[digestLayers.length - 1].length > 1) {
    const prev = digestLayers[digestLayers.length - 1];
    const nextLenPadded = prev.length >> 1;
    const inject = [];
    while (idx < ordered.length && nextPow2(ordered[idx].height) === nextLenPadded) inject.push(ordered[idx++]);
    digestLayers.push(compressAndInject(prev, inject, nextLenPadded));
  }
  return { digestLayers, logMaxHeight };
}

function compressAndInject(prev, inject, nextLenPadded) {
  const out = [];
  for (let i = 0; i < nextLenPadded; i++) out.push([...DEFAULT_DIGEST]);
  if (!inject.length) {
    for (let i = 0; i < nextLenPadded; i++) out[i] = compress2to1(prev[2 * i], prev[2 * i + 1]);
    return out;
  }
  const injectHeight = inject[0].height;
  for (let i = 0; i < injectHeight; i++) {
    const digest = compress2to1(prev[2 * i], prev[2 * i + 1]);
    const row = [];
    for (const m of inject) row.push(...m.row(i));
    out[i] = compress2to1(digest, hashIter(row));
  }
  for (let i = injectHeight; i < nextLenPadded; i++) {
    const digest = compress2to1(prev[2 * i], prev[2 * i + 1]);
    out[i] = compress2to1(digest, [...DEFAULT_DIGEST]);
  }
  return out;
}

export function commit(matrices) {
  const { digestLayers, logMaxHeight } = buildTree(matrices);
  const root = digestLayers[digestLayers.length - 1][0];
  return { root, proverData: { matrices, digestLayers, logMaxHeight } };
}

export function openBatch(index, proverData) {
  const { matrices, digestLayers, logMaxHeight } = proverData;
  const openings = matrices.map((m) => {
    const bitsReduced = logMaxHeight - log2Ceil(m.height);
    return m.row(index >> bitsReduced);
  });
  const proof = [];
  for (let i = 0; i < logMaxHeight; i++) proof.push(digestLayers[i][(index >> i) ^ 1]);
  return { openings, proof };
}

// ---- verify_batch (the on-chain-shaped check) ----
export function verifyBatch(commitRoot, dimensions, index, openedValues, proof) {
  const n = dimensions.length;
  if (n === 0) return false;
  const order = dimensions.map((_, i) => i).sort((a, b) => dimensions[b][1] - dimensions[a][1] || a - b);
  const heights = order.map((i) => dimensions[i][1]);
  let ptr = 0;
  const takeGroup = (padded) => {
    const idxs = [];
    while (ptr < n && nextPow2(heights[ptr]) === padded) idxs.push(order[ptr++]);
    return idxs;
  };
  let currHeightPadded = nextPow2(heights[0]);
  let seed = [];
  for (const oi of takeGroup(currHeightPadded)) seed.push(...openedValues[oi]);
  let root = hashIter(seed);
  let idx = index;
  for (const sibling of proof) {
    const [left, right] = (idx & 1) === 0 ? [root, sibling] : [sibling, root];
    root = compress2to1(left, right);
    idx >>= 1;
    currHeightPadded >>= 1;
    if (ptr < n && nextPow2(heights[ptr]) === currHeightPadded) {
      const row = [];
      for (const oi of takeGroup(currHeightPadded)) row.push(...openedValues[oi]);
      root = compress2to1(root, hashIter(row));
    }
  }
  const want = commitRoot.map((x) => mod(BigInt(x)));
  return root.length === want.length && root.every((v, i) => v === want[i]);
}

// ---- matrix construction (shared with the extractor) ----
export function makeMatrix(height, width, base) {
  const rows = [];
  for (let i = 0; i < height; i++) {
    const r = [];
    for (let j = 0; j < width; j++) r.push(base + i * width + j);
    rows.push(r);
  }
  return new Matrix(rows);
}

// ---- CONFIRMED conformance vectors (from the pinned library) ----
export const PERM_ZEROS = [1787823396, 953829438, 89382455, 347481625, 1754527224, 916217775, 1056029082, 410644796, 1169123478, 1854704276, 1195829987, 1485264906, 1824644035, 1948268315, 847945433, 190591038];

export const CONFORMANCE_CASES = [
  { name: "single_8x2", matspec: [[8, 2, 10]], index: 3,
    root: [35761595, 1133593632, 733114748, 517674920, 1449914397, 74512820, 381712048, 469819303],
    openings: [[16, 17]],
    proof: [[1513856679, 1890540197, 1949407553, 586338782, 197781917, 573541314, 1955108120, 661229895], [57846071, 657849236, 262378030, 1091294242, 669362455, 676261940, 190131986, 664894526], [1143223726, 1793027126, 1353127284, 491971160, 1959272008, 1195367436, 417822678, 1372117039]] },
  { name: "single_6x2", matspec: [[6, 2, 100]], index: 5,
    root: [1043564500, 1812685593, 1597353357, 1392680266, 1355213241, 1159759876, 1586103175, 799081422],
    openings: [[110, 111]],
    proof: [[1089158652, 427625262, 1157475631, 1692145862, 1671348918, 1986544368, 497315490, 1960470114], [1787823396, 953829438, 89382455, 347481625, 1754527224, 916217775, 1056029082, 410644796], [747170349, 1608481817, 1569266992, 1154307224, 561605764, 1907985030, 1454474361, 357054770]] },
  { name: "single_8x1", matspec: [[8, 1, 200]], index: 6,
    root: [447902873, 1264273264, 1781377203, 392525377, 451230220, 285479002, 124104889, 737840045],
    openings: [[206]],
    proof: [[854476755, 1853946983, 1409001429, 1897109439, 1880752521, 1036730533, 1817549359, 907866104], [757904997, 571589503, 673816220, 1066413580, 473597103, 1271204909, 363424036, 1450652756], [402779725, 764084359, 1406275333, 1158469515, 32391761, 236523006, 1039588073, 264072612]] },
  { name: "mixed_8x2_4x3", matspec: [[8, 2, 10], [4, 3, 300]], index: 5,
    root: [902263792, 353615665, 93922356, 1251424848, 837594048, 2005066510, 431363100, 287722769],
    openings: [[20, 21], [306, 307, 308]],
    proof: [[1732750591, 1876733121, 1364724824, 1847015379, 1792368333, 483325461, 1174939365, 1939380322], [932176417, 320077372, 1748871293, 625438914, 179521004, 655947905, 828588297, 31952751], [930401557, 471785873, 904189917, 1626475238, 550043796, 1590748862, 306781755, 496390863]] },
  { name: "mixed_8x1_4x2_2x2", matspec: [[8, 1, 1], [4, 2, 50], [2, 2, 400]], index: 6,
    root: [439675894, 37405611, 255560432, 531897590, 349476430, 1463998586, 540993751, 1307259928],
    openings: [[7], [56, 57], [402, 403]],
    proof: [[723290459, 412642417, 484634914, 785252854, 182396116, 1183763863, 630430401, 175545021], [1257356383, 1204259532, 1915921675, 53621155, 1015367194, 742423503, 1998129904, 1554700909], [13474675, 1966185163, 1255599565, 471961853, 1823737756, 1607265064, 1964440250, 1219296485]] },
  { name: "mixed_5x2_3x1", matspec: [[5, 2, 20], [3, 1, 70]], index: 4,
    root: [700351145, 394944774, 166238365, 1683786139, 1659356855, 1460803726, 573770743, 1196572690],
    openings: [[28, 29], [72]],
    proof: [[0, 0, 0, 0, 0, 0, 0, 0], [1714214715, 1781831455, 644202942, 667527548, 1395787195, 671219385, 1433934871, 1883643874], [1521468259, 1627990232, 1145489794, 839124709, 507242897, 777414579, 429233007, 1147797058]] },
];

const caseMatrices = (c) => c.matspec.map(([h, w, b]) => makeMatrix(h, w, b));
const caseDims = (c) => c.matspec.map(([h, w]) => [w, h]);
const toNums = (arr) => arr.map((x) => Number(x));

// ---- shared cross-language vector set ----
export function sharedVectors() {
  const prim = {
    hash_item_5: toNums(hashItem(5)),
    hash_slice_2: toNums(hashSlice([1, 2])),
    hash_slice_3: toNums(hashSlice([1, 2, 3])),
    hash_slice_8: toNums(hashSlice([1, 2, 3, 4, 5, 6, 7, 8])),
    hash_slice_9: toNums(hashSlice([1, 2, 3, 4, 5, 6, 7, 8, 9])),
    compress_h2_h3: toNums(compress2to1(hashSlice([1, 2]), hashSlice([1, 2, 3]))),
  };
  const cases = [];
  for (const c of CONFORMANCE_CASES) {
    const mats = caseMatrices(c);
    const { root, proverData } = commit(mats);
    const { openings, proof } = openBatch(c.index, proverData);
    const dims = caseDims(c);
    cases.push({
      name: c.name,
      index: c.index,
      root: toNums(root),
      openings: openings.map(toNums),
      proof: proof.map(toNums),
      rootMatch: toNums(root).every((v, i) => v === c.root[i]),
      openingsMatch: JSON.stringify(openings.map(toNums)) === JSON.stringify(c.openings),
      proofMatch: JSON.stringify(proof.map(toNums)) === JSON.stringify(c.proof),
      verifyOk: verifyBatch(c.root, dims, c.index, c.openings, c.proof),
    });
  }
  return {
    constants: { P: Number(P), width: WIDTH, rate: RATE, out: OUT, digestElems: DIGEST_ELEMS },
    permZeros: toNums(permute(Array(16).fill(0n))),
    prim,
    cases,
  };
}

if (import.meta.url === `file://${process.argv[1]}` || process.argv[1]?.endsWith("mmcs_babybear_reference.mjs")) {
  process.stdout.write(JSON.stringify(sharedVectors()));
}