// Independent (second-language) reference for the Poseidon2 permutation over BabyBear at width 16,
// component (b) of the PQ STARK-verify precompile 0x0AE8. Mirrors poseidon2_babybear_reference.py
// and the standalone Java Poseidon2BabyBearSelfTest. It serves two purposes: cross-language agreement
// (three independent implementations produce byte-identical output) AND conformance (it reproduces
// real known-answer vectors from the pinned Plonky3 crates; see CONFORMANCE_KATS). 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 3.
//
// HONEST SCOPE (CONFIRMED 2026-07-19): the 141 round constants, internal diagonal, M4, round counts,
// x^7 S-box, and layer order all match p3-baby-bear / p3-poseidon2 0.4.3-succinct (the exact crates
// pinned in aerenew Cargo.lock; lockfile checksums matched). A real known-answer test PASSES.
// MONTGOMERY SUBTLETY (resolved): the internal layer as a canonical map is R^{-1} * (J + diag(D)),
// R^{-1} = 943718400; a prior pass omitted the R^{-1} (textbook state[i]*D[i] + sum) which is
// self-consistent but disagrees with the prover. See the .py header for the full source citation.
//
// Uses BigInt throughout (x^7 and the linear layers exceed 2^53). Run standalone to emit JSON:
//   node poseidon2_babybear_reference.mjs

export const P = 2013265921n;
// Montgomery inverse factor: p3-baby-bear stores field elements in Montgomery form with R = 2^32; its
// internal diffusion layer, as a canonical linear map, carries a factor of R^{-1} = (2^32)^{-1} mod p.
export const R_INV = 943718400n; // = (2^32)^{-1} mod p, verified against the pinned library's matrix
export const WIDTH = 16;
export const SBOX_DEGREE = 7n;
export const ROUNDS_F = 8;
export const ROUNDS_P = 13;

export const M4 = [
  [2n, 3n, 1n, 1n],
  [1n, 2n, 3n, 1n],
  [1n, 1n, 2n, 3n],
  [3n, 1n, 1n, 2n],
];

export const INTERNAL_DIAG_M1_16 = [
  P - 2n, 1n, 2n, 4n, 8n, 16n, 32n, 64n, 128n, 256n, 512n, 1024n, 2048n, 4096n, 8192n, 32768n,
];

// ---- base-field helpers ----
const mod = (a) => ((a % P) + P) % P;
export const add = (a, b) => mod(a + b);
export const mul = (a, b) => mod(a * b);

export function sboxMono(x) {
  x = mod(x);
  const x2 = mul(x, x);
  const x4 = mul(x2, x2);
  return mul(x4, mul(x2, x)); // x^7
}

function sboxInvExponent() {
  // 7^{-1} mod (p-1) via BigInt extended Euclid (p-1 = 2013265920).
  const m = P - 1n;
  let [oldR, r] = [SBOX_DEGREE % m, m];
  let [oldS, s] = [1n, 0n];
  while (r !== 0n) {
    const q = oldR / r;
    [oldR, r] = [r, oldR - q * r];
    [oldS, s] = [s, oldS - q * s];
  }
  return ((oldS % m) + m) % m;
}
const SBOX_INV_E = sboxInvExponent();

export function powF(base, e) {
  let b = mod(base);
  let acc = 1n;
  let ee = e;
  while (ee > 0n) {
    if (ee & 1n) acc = mul(acc, b);
    b = mul(b, b);
    ee >>= 1n;
  }
  return acc;
}
export const sboxMonoInv = (y) => powF(mod(y), SBOX_INV_E);

// ---- CONFIRMED round constants (from p3-baby-bear/p3-poseidon2 0.4.3-succinct via seed_from_u64(1)) ----
export const RC_EXTERNAL = [
  [
    1321363468n, 285374923n, 858595076n, 131742120n, 550898981n, 109281027n,
    1548327248n, 299186948n, 1198120888n, 1302311359n, 568137078n, 1484856917n,
    1301979945n, 725688886n, 941758026n, 323341913n,
  ],
  [
    1049323172n, 822409348n, 1406080127n, 1279024384n, 214862539n, 904628921n,
    1320747287n, 11578228n, 1036373712n, 1474430466n, 1430509860n, 111174484n,
    1124450171n, 85382027n, 679880882n, 243277213n,
  ],
  [
    1338495990n, 1523013347n, 1841068573n, 578194469n, 47683837n, 1790441672n,
    1628061601n, 1716216090n, 1635810049n, 1115145248n, 1117524270n, 678640014n,
    1962751651n, 1367401392n, 11688709n, 1950824358n,
  ],
  [
    528649031n, 1937116923n, 1460949223n, 1193074357n, 1221801411n, 1183923117n,
    433505619n, 1928933309n, 505759755n, 285671663n, 1047265910n, 909281502n,
    1258966486n, 864761693n, 307024510n, 504858517n,
  ],
  [
    1467478033n, 1754565867n, 432187324n, 1452390672n, 881974300n, 550050336n,
    1447309270n, 939419487n, 1783112406n, 1166910332n, 107514714n, 580516863n,
    2003318760n, 854475946n, 934896823n, 994783668n,
  ],
  [
    1841107561n, 438269126n, 1550523825n, 913322122n, 600932628n, 583000098n,
    1262690949n, 105797869n, 277542016n, 170491952n, 365854467n, 1479645308n,
    1457660602n, 1635879552n, 499155053n, 741227047n,
  ],
  [
    651389942n, 464828001n, 89696107n, 360044673n, 230330371n, 1773129416n,
    1380150763n, 745014723n, 793475694n, 1361274828n, 1443741698n, 51616650n,
    731414218n, 1087554954n, 1273943885n, 311581717n,
  ],
  [
    702702762n, 1473247301n, 132108357n, 1348260424n, 476775430n, 1438949459n,
    2434448n, 1349232398n, 1954471898n, 1762138591n, 1271221795n, 1593266476n,
    864488771n, 139147729n, 1053373910n, 422842363n,
  ],
];

export const RC_INTERNAL = [
  402771160n, 320708227n, 1122772462n, 100431997n, 202594011n, 1226485372n,
  1088619034n, 64118538n, 109828860n, 724723599n, 1662837151n, 797753907n,
  1075635743n,
];

// ---- linear layers ----
function m4Apply(t) {
  const out = [];
  for (let row = 0; row < 4; row++) {
    out.push(mod(M4[row][0] * t[0] + M4[row][1] * t[1] + M4[row][2] * t[2] + M4[row][3] * t[3]));
  }
  return out;
}

export function externalLayer(state) {
  const blocks = [];
  for (let b = 0; b < 4; b++) blocks.push(m4Apply(state.slice(4 * b, 4 * b + 4)));
  const colSum = [];
  for (let j = 0; j < 4; j++) colSum.push(mod(blocks[0][j] + blocks[1][j] + blocks[2][j] + blocks[3][j]));
  const out = new Array(WIDTH);
  for (let b = 0; b < 4; b++) for (let j = 0; j < 4; j++) out[4 * b + j] = mod(blocks[b][j] + colSum[j]);
  return out;
}

export function internalLayer(state) {
  // p3-baby-bear DiffusionMatrixBabyBear as a canonical map: M_I = R^{-1} * (J + diag(D)),
  // i.e. out[i] = R_INV * (sum + D[i]*state[i]). The R_INV factor is the Montgomery-form artifact.
  let s = 0n;
  for (const v of state) s += v;
  s = mod(s);
  return state.map((v, i) => mod(R_INV * mod(s + INTERNAL_DIAG_M1_16[i] * v)));
}

// ---- explicit matrices (for invertibility / inverse round-trip only) ----
export function externalMatrix() {
  const m = Array.from({ length: WIDTH }, () => new Array(WIDTH).fill(0n));
  for (let bi = 0; bi < 4; bi++)
    for (let bj = 0; bj < 4; bj++) {
      const coef = bi === bj ? 2n : 1n;
      for (let r = 0; r < 4; r++) for (let c = 0; c < 4; c++) m[4 * bi + r][4 * bj + c] = mod(coef * M4[r][c]);
    }
  return m;
}
export function internalMatrix() {
  // M_I = R^{-1} * (J + diag(D)): off-diagonal = R_INV, diagonal = R_INV*(1 + D[i]).
  const m = Array.from({ length: WIDTH }, () => new Array(WIDTH).fill(R_INV));
  for (let i = 0; i < WIDTH; i++) m[i][i] = mod(R_INV * (1n + INTERNAL_DIAG_M1_16[i]));
  return m;
}
export function matVec(m, v) {
  return m.map((row) => {
    let acc = 0n;
    for (let j = 0; j < WIDTH; j++) acc = mod(acc + mul(row[j], v[j]));
    return acc;
  });
}
export function matInverse(m) {
  const n = WIDTH;
  const a = m.map((row, i) => [...row.map(mod), ...Array.from({ length: n }, (_, j) => (j === i ? 1n : 0n))]);
  for (let col = 0; col < n; col++) {
    let piv = -1;
    for (let r = col; r < n; r++) if (mod(a[r][col]) !== 0n) { piv = r; break; }
    if (piv < 0) throw new Error("singular matrix");
    [a[col], a[piv]] = [a[piv], a[col]];
    const invP = powF(a[col][col], P - 2n);
    a[col] = a[col].map((x) => mul(x, invP));
    for (let r = 0; r < n; r++) {
      if (r !== col && mod(a[r][col]) !== 0n) {
        const f = a[r][col];
        a[r] = a[r].map((x, k) => mod(x - mul(f, a[col][k])));
      }
    }
  }
  return a.map((row) => row.slice(n));
}

// ---- the permutation ----
export function permute(state) {
  if (state.length !== WIDTH) throw new Error("state must have WIDTH elements");
  let s = state.map(mod);
  s = externalLayer(s);
  const half = ROUNDS_F / 2;
  for (let r = 0; r < half; r++) {
    s = s.map((v, i) => add(v, RC_EXTERNAL[r][i]));
    s = s.map(sboxMono);
    s = externalLayer(s);
  }
  for (let r = 0; r < ROUNDS_P; r++) {
    s[0] = add(s[0], RC_INTERNAL[r]);
    s[0] = sboxMono(s[0]);
    s = internalLayer(s);
  }
  for (let r = half; r < ROUNDS_F; r++) {
    s = s.map((v, i) => add(v, RC_EXTERNAL[r][i]));
    s = s.map(sboxMono);
    s = externalLayer(s);
  }
  return s;
}

export function permuteInverse(state) {
  let s = state.map(mod);
  const extInv = matInverse(externalMatrix());
  const intInv = matInverse(internalMatrix());
  const half = ROUNDS_F / 2;
  for (let r = ROUNDS_F - 1; r >= half; r--) {
    s = matVec(extInv, s);
    s = s.map(sboxMonoInv);
    s = s.map((v, i) => mod(v - RC_EXTERNAL[r][i]));
  }
  for (let r = ROUNDS_P - 1; r >= 0; r--) {
    s = matVec(intInv, s);
    s[0] = sboxMonoInv(s[0]);
    s[0] = mod(s[0] - RC_INTERNAL[r]);
  }
  for (let r = half - 1; r >= 0; r--) {
    s = matVec(extInv, s);
    s = s.map(sboxMonoInv);
    s = s.map((v, i) => mod(v - RC_EXTERNAL[r][i]));
  }
  s = matVec(extInv, s);
  return s;
}

// ---- CONFIRMED conformance KATs (from the pinned p3-baby-bear/p3-poseidon2 0.4.3-succinct crates) ----
export const CONFORMANCE_KATS = [
  {
    name: "zeros",
    in: new Array(16).fill(0n),
    out: [1787823396n, 953829438n, 89382455n, 347481625n, 1754527224n, 916217775n, 1056029082n,
      410644796n, 1169123478n, 1854704276n, 1195829987n, 1485264906n, 1824644035n, 1948268315n,
      847945433n, 190591038n],
  },
  {
    name: "iota",
    in: Array.from({ length: 16 }, (_, i) => BigInt(i)),
    out: [157639285n, 1851003038n, 1852457045n, 1920360618n, 779990819n, 1080011039n, 585017685n,
      1093051731n, 249426030n, 967262243n, 623744062n, 280332881n, 1995600430n, 1751988435n,
      317724737n, 1895035071n],
  },
  {
    name: "testvec",
    in: [894848333n, 1437655012n, 1200606629n, 1690012884n, 71131202n, 1749206695n, 1717947831n,
      120589055n, 19776022n, 42382981n, 1831865506n, 724844064n, 171220207n, 1299207443n,
      227047920n, 1783754913n],
    out: [512585766n, 975869435n, 1921378527n, 1238606951n, 899635794n, 132650430n, 1426417547n,
      1734425242n, 57415409n, 67173027n, 1535042492n, 1318033394n, 1070659233n, 17258943n,
      856719028n, 1500534995n],
  },
];

export function conformanceKats() {
  let passed = 0;
  for (const kat of CONFORMANCE_KATS) {
    const out = permute(kat.in.slice());
    if (out.every((v, i) => mod(v) === mod(kat.out[i]))) passed++;
  }
  return { passed, total: CONFORMANCE_KATS.length };
}

// ---- shared cross-language vector set ----
function inputStates() {
  const states = [
    new Array(WIDTH).fill(0n),
    Array.from({ length: WIDTH }, (_, i) => BigInt(i)),
    new Array(WIDTH).fill(P - 1n),
    Array.from({ length: WIDTH }, (_, i) => (BigInt(i) * 2654435761n) % P),
    [11n, 0n, 0n, 0n, 0n, 0n, 0n, 0n, 0n, 0n, 0n, 0n, 0n, 0n, 0n, 1n],
  ];
  // BigInt LCG (1103515245*x overflows JS Number's 2^53; BigInt keeps it exact and matching Py/Java)
  let x = 123456789n;
  for (let k = 0; k < 8; k++) {
    const st = [];
    for (let i = 0; i < WIDTH; i++) {
      x = (1103515245n * x + 12345n) % 2147483648n;
      st.push(x % P);
    }
    states.push(st);
  }
  return states;
}

const nlist = (a) => a.map((v) => Number(v));

export function sharedVectors() {
  const rcExtFlat = [];
  for (let r = 0; r < ROUNDS_F; r++) for (let i = 0; i < WIDTH; i++) rcExtFlat.push(Number(RC_EXTERNAL[r][i]));
  const states = inputStates();
  return {
    constants: {
      P: Number(P),
      rInv: Number(R_INV),
      width: WIDTH,
      sboxDegree: Number(SBOX_DEGREE),
      roundsF: ROUNDS_F,
      roundsP: ROUNDS_P,
      internalDiagM1: nlist(INTERNAL_DIAG_M1_16),
      m4: [2, 3, 1, 1, 1, 2, 3, 1, 1, 1, 2, 3, 3, 1, 1, 2],
      rcExternalFlat: rcExtFlat,
      rcInternal: nlist(RC_INTERNAL),
    },
    permute: states.map((st) => {
      const out = permute(st);
      const back = permuteInverse(out);
      const rt = back.every((v, i) => mod(v) === mod(st[i]));
      return { in: nlist(st), out: nlist(out), roundtrip: rt };
    }),
    conformanceKats: CONFORMANCE_KATS.map((k) => {
      const out = permute(k.in.slice());
      return { name: k.name, in: nlist(k.in), out: nlist(out), match: out.every((v, i) => mod(v) === mod(k.out[i])) };
    }),
  };
}

import { pathToFileURL } from "node:url";
if (process.argv[1] && import.meta.url === pathToFileURL(process.argv[1]).href) {
  process.stdout.write(JSON.stringify(sharedVectors()));
}