// Standalone (no-Besu-classpath) copy of the Poseidon2-over-BabyBear width-16 permutation used by the
// PQ STARK-verify precompile 0x0AE8, plus a self-test. It is a verbatim mirror of the
// Sp1StarkVerifierPrecompiledContract.Poseidon2Bb logic (same P, same M4, same internal diagonal, same
// CONFIRMED round constants), compilable and runnable without the Besu EVM jars so the permutation
// can be exercised offline and cross-checked against the Python and Node references.
//
// HONEST SCOPE (CONFIRMED 2026-07-19). This is component (b) of the six-component STARK port
// (docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 3). The 141 round constants, internal diagonal, M4,
// round counts (ROUNDS_F=8, ROUNDS_P=13), 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), and a real known-answer test PASSES (conformanceKats: zeros/iota/testvec emitted by
// executing the pinned crates). 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. It verifies NOTHING
// about a real proof and the top-level 0x0AE8 precompile stays fail-closed. See
// poseidon2_babybear_reference.py for the full source citation.
//
// Usage:
//   javac -d out Poseidon2BabyBearSelfTest.java
//   java  -cp out Poseidon2BabyBearSelfTest            # run the self-test (prints PASS/FAIL counts)
//   java  -cp out Poseidon2BabyBearSelfTest --emit     # emit the shared cross-language vector set (JSON)

import java.math.BigInteger;
import java.util.Random;

public final class Poseidon2BabyBearSelfTest {

  static final long P = 2013265921L; // BabyBear 2^31 - 2^27 + 1
  // 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.
  static final long R_INV = 943718400L; // = (2^32)^{-1} mod p, verified against the pinned library's matrix
  static final int WIDTH = 16;
  static final int SBOX_DEGREE = 7;  // x^7: smallest d>1 with gcd(d, p-1)=1. CONFIRMED.
  static final int ROUNDS_F = 8;     // external (full) rounds. CONFIRMED (round_numbers).
  static final int ROUNDS_P = 13;    // internal (partial) rounds. CONFIRMED (round_numbers).

  static final long[][] M4 = {
    {2, 3, 1, 1},
    {1, 2, 3, 1},
    {1, 1, 2, 3},
    {3, 1, 1, 2},
  };

  // diag(M_I - I) = D; internal layer M_I = R^{-1} * (J + diag(D)). CONFIRMED against p3-baby-bear
  // 0.4.3-succinct POSEIDON2_INTERNAL_MATRIX_DIAG_16_BABYBEAR_MONTY (canonical form).
  static final long[] INTERNAL_DIAG_M1_16 = {
    P - 2, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 32768
  };

  // ---- base field ----
  static long reduce(final long a) {
    long m = a % P;
    if (m < 0) m += P;
    return m;
  }

  static long add(final long a, final long b) {
    long s = a + b;
    if (s >= P) s -= P;
    return s;
  }

  static long sub(final long a, final long b) {
    long s = a - b;
    if (s < 0) s += P;
    return s;
  }

  static long mul(final long a, final long b) {
    return (reduce(a) * reduce(b)) % P; // both < 2^31, product < 2^62
  }

  static long pow(final long base, long e) {
    long b = reduce(base);
    long acc = 1L;
    while (e > 0) {
      if ((e & 1L) == 1L) acc = mul(acc, b);
      b = mul(b, b);
      e >>= 1;
    }
    return acc;
  }

  static long sboxMono(final long x) {
    final long x2 = mul(x, x);
    final long x4 = mul(x2, x2);
    return mul(x4, mul(x2, x)); // x^7
  }

  // 7^{-1} mod (p-1): exists because gcd(7, p-1) = 1. Used only by the inverse permutation.
  static final long SBOX_INV_E =
      BigInteger.valueOf(SBOX_DEGREE).modInverse(BigInteger.valueOf(P - 1)).longValueExact();

  static long sboxMonoInv(final long y) {
    return pow(y, SBOX_INV_E);
  }

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

  static final long[] RC_INTERNAL = {
    402771160L, 320708227L, 1122772462L, 100431997L, 202594011L, 1226485372L,
    1088619034L, 64118538L, 109828860L, 724723599L, 1662837151L, 797753907L,
    1075635743L,
  };

  // ---- linear layers ----
  static long[] m4Apply(final long[] t) {
    final long[] out = new long[4];
    for (int row = 0; row < 4; row++) {
      out[row] = reduce(M4[row][0] * t[0] + M4[row][1] * t[1] + M4[row][2] * t[2] + M4[row][3] * t[3]);
    }
    return out;
  }

  static long[] externalLayer(final long[] state) {
    final long[][] blocks = new long[4][];
    for (int b = 0; b < 4; b++) {
      blocks[b] = m4Apply(new long[] {state[4 * b], state[4 * b + 1], state[4 * b + 2], state[4 * b + 3]});
    }
    final long[] colSum = new long[4];
    for (int j = 0; j < 4; j++) {
      colSum[j] = reduce(blocks[0][j] + blocks[1][j] + blocks[2][j] + blocks[3][j]);
    }
    final long[] out = new long[WIDTH];
    for (int b = 0; b < 4; b++) for (int j = 0; j < 4; j++) out[4 * b + j] = add(blocks[b][j], colSum[j]);
    return out;
  }

  static long[] internalLayer(final long[] 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.
    long s = 0;
    for (final long v : state) s = add(s, v);
    final long[] out = new long[WIDTH];
    for (int i = 0; i < WIDTH; i++) out[i] = mul(R_INV, add(s, mul(state[i], INTERNAL_DIAG_M1_16[i])));
    return out;
  }

  // ---- explicit matrices (for invertibility / inverse round-trip only) ----
  static long[][] externalMatrix() {
    final long[][] m = new long[WIDTH][WIDTH];
    for (int bi = 0; bi < 4; bi++) {
      for (int bj = 0; bj < 4; bj++) {
        final long coef = bi == bj ? 2 : 1;
        for (int r = 0; r < 4; r++) for (int c = 0; c < 4; c++) m[4 * bi + r][4 * bj + c] = mul(coef, M4[r][c]);
      }
    }
    return m;
  }

  static long[][] internalMatrix() {
    // M_I = R^{-1} * (J + diag(D)): off-diagonal = R_INV, diagonal = R_INV*(1 + D[i]).
    final long[][] m = new long[WIDTH][WIDTH];
    for (int i = 0; i < WIDTH; i++) for (int j = 0; j < WIDTH; j++) m[i][j] = R_INV;
    for (int i = 0; i < WIDTH; i++) m[i][i] = mul(R_INV, add(1, INTERNAL_DIAG_M1_16[i]));
    return m;
  }

  static long[] matVec(final long[][] m, final long[] v) {
    final long[] out = new long[WIDTH];
    for (int i = 0; i < WIDTH; i++) {
      long acc = 0;
      for (int j = 0; j < WIDTH; j++) acc = add(acc, mul(m[i][j], v[j]));
      out[i] = acc;
    }
    return out;
  }

  static long[][] matInverse(final long[][] m) {
    final int n = WIDTH;
    final long[][] a = new long[n][2 * n];
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < n; j++) a[i][j] = reduce(m[i][j]);
      a[i][n + i] = 1;
    }
    for (int col = 0; col < n; col++) {
      int piv = -1;
      for (int r = col; r < n; r++) if (a[r][col] != 0) { piv = r; break; }
      if (piv < 0) throw new IllegalStateException("singular matrix");
      final long[] tmp = a[col]; a[col] = a[piv]; a[piv] = tmp;
      final long invP = pow(a[col][col], P - 2);
      for (int k = 0; k < 2 * n; k++) a[col][k] = mul(a[col][k], invP);
      for (int r = 0; r < n; r++) {
        if (r != col && a[r][col] != 0) {
          final long f = a[r][col];
          for (int k = 0; k < 2 * n; k++) a[r][k] = sub(a[r][k], mul(f, a[col][k]));
        }
      }
    }
    final long[][] inv = new long[n][n];
    for (int i = 0; i < n; i++) for (int j = 0; j < n; j++) inv[i][j] = a[i][n + j];
    return inv;
  }

  // ---- the permutation ----
  static long[] permute(final long[] state) {
    if (state.length != WIDTH) throw new IllegalArgumentException("state must have WIDTH elements");
    long[] s = new long[WIDTH];
    for (int i = 0; i < WIDTH; i++) s[i] = reduce(state[i]);

    s = externalLayer(s);
    final int half = ROUNDS_F / 2;
    for (int r = 0; r < half; r++) {
      for (int i = 0; i < WIDTH; i++) s[i] = add(s[i], RC_EXTERNAL[r][i]);
      for (int i = 0; i < WIDTH; i++) s[i] = sboxMono(s[i]);
      s = externalLayer(s);
    }
    for (int r = 0; r < ROUNDS_P; r++) {
      s[0] = add(s[0], RC_INTERNAL[r]);
      s[0] = sboxMono(s[0]);
      s = internalLayer(s);
    }
    for (int r = half; r < ROUNDS_F; r++) {
      for (int i = 0; i < WIDTH; i++) s[i] = add(s[i], RC_EXTERNAL[r][i]);
      for (int i = 0; i < WIDTH; i++) s[i] = sboxMono(s[i]);
      s = externalLayer(s);
    }
    return s;
  }

  static long[] permuteInverse(final long[] state) {
    long[] s = new long[WIDTH];
    for (int i = 0; i < WIDTH; i++) s[i] = reduce(state[i]);
    final long[][] extInv = matInverse(externalMatrix());
    final long[][] intInv = matInverse(internalMatrix());
    final int half = ROUNDS_F / 2;

    for (int r = ROUNDS_F - 1; r >= half; r--) {
      s = matVec(extInv, s);
      for (int i = 0; i < WIDTH; i++) s[i] = sboxMonoInv(s[i]);
      for (int i = 0; i < WIDTH; i++) s[i] = sub(s[i], RC_EXTERNAL[r][i]);
    }
    for (int r = ROUNDS_P - 1; r >= 0; r--) {
      s = matVec(intInv, s);
      s[0] = sboxMonoInv(s[0]);
      s[0] = sub(s[0], RC_INTERNAL[r]);
    }
    for (int r = half - 1; r >= 0; r--) {
      s = matVec(extInv, s);
      for (int i = 0; i < WIDTH; i++) s[i] = sboxMonoInv(s[i]);
      for (int i = 0; i < WIDTH; i++) s[i] = sub(s[i], 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) ----
  static final long[][] KAT_IN = {
    {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
    {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
    {894848333L, 1437655012L, 1200606629L, 1690012884L, 71131202L, 1749206695L, 1717947831L,
      120589055L, 19776022L, 42382981L, 1831865506L, 724844064L, 171220207L, 1299207443L,
      227047920L, 1783754913L},
  };
  static final long[][] KAT_OUT = {
    {1787823396L, 953829438L, 89382455L, 347481625L, 1754527224L, 916217775L, 1056029082L,
      410644796L, 1169123478L, 1854704276L, 1195829987L, 1485264906L, 1824644035L, 1948268315L,
      847945433L, 190591038L},
    {157639285L, 1851003038L, 1852457045L, 1920360618L, 779990819L, 1080011039L, 585017685L,
      1093051731L, 249426030L, 967262243L, 623744062L, 280332881L, 1995600430L, 1751988435L,
      317724737L, 1895035071L},
    {512585766L, 975869435L, 1921378527L, 1238606951L, 899635794L, 132650430L, 1426417547L,
      1734425242L, 57415409L, 67173027L, 1535042492L, 1318033394L, 1070659233L, 17258943L,
      856719028L, 1500534995L},
  };

  // ================================= self-test =================================
  static int pass = 0;
  static int fail = 0;

  static void check(final String name, final boolean cond) {
    if (cond) pass++;
    else {
      fail++;
      System.out.println("FAIL " + name);
    }
  }

  static long gcd(long a, long b) {
    while (b != 0) { final long t = a % b; a = b; b = t; }
    return a;
  }

  static boolean eq(final long[] a, final long[] b) {
    if (a.length != b.length) return false;
    for (int i = 0; i < a.length; i++) if (reduce(a[i]) != reduce(b[i])) return false;
    return true;
  }

  static void selfTest() {
    // S-box degree: 7 is the smallest d>1 with gcd(d, p-1)=1 (p-1 = 2^27*3*5).
    check("sbox.deg.min", SBOX_DEGREE == 7);
    for (int d = 2; d <= 6; d++) check("sbox.deg.notcoprime.d" + d, gcd(d, P - 1) != 1);
    check("sbox.deg7.coprime", gcd(7, P - 1) == 1);

    final Random rnd = new Random(0x0AE8);
    // x^7 is a bijection: inverse round-trips
    for (int it = 0; it < 5000; it++) {
      final long x = Math.floorMod(rnd.nextLong(), P);
      check("sbox.bijection", sboxMonoInv(sboxMono(x)) == x);
    }

    // linear layers equal their explicit matrices; matrices are invertible
    final long[][] ext = externalMatrix();
    final long[][] inte = internalMatrix();
    final long[][] extInv = matInverse(ext);
    final long[][] intInv = matInverse(inte);
    for (int it = 0; it < 2000; it++) {
      final long[] v = new long[WIDTH];
      for (int i = 0; i < WIDTH; i++) v[i] = Math.floorMod(rnd.nextLong(), P);
      check("ext.layer.eq.matrix", eq(externalLayer(v), matVec(ext, v)));
      check("int.layer.eq.matrix", eq(internalLayer(v), matVec(inte, v)));
      check("ext.inv.roundtrip", eq(matVec(extInv, matVec(ext, v)), v));
      check("int.inv.roundtrip", eq(matVec(intInv, matVec(inte, v)), v));
    }

    // permutation: determinism + inverse round-trip (genuine bijection)
    for (final long[] st : inputStates()) {
      check("permute.deterministic", eq(permute(st), permute(st)));
      check("permute.bijection", eq(permuteInverse(permute(st)), st));
    }
    for (int it = 0; it < 500; it++) {
      final long[] st = new long[WIDTH];
      for (int i = 0; i < WIDTH; i++) st[i] = Math.floorMod(rnd.nextLong(), P);
      check("permute.rand.bijection", eq(permuteInverse(permute(st)), st));
    }

    // CONFORMANCE: reproduce the pinned-library known-answer vectors exactly.
    for (int k = 0; k < KAT_IN.length; k++) {
      check("conformance.kat" + k, eq(permute(KAT_IN[k]), KAT_OUT[k]));
    }

    System.out.println("Poseidon2BabyBearSelfTest PASS=" + pass + " FAIL=" + fail);
    if (fail != 0) {
      System.out.println("POSEIDON2_BABYBEAR_SELFTEST_FAILED");
      System.exit(1);
    }
    System.out.println("POSEIDON2_BABYBEAR_SELFTEST_OK");
  }

  // ================================= shared-vector emit (JSON) =================================
  static long[][] inputStates() {
    final long[][] fixed = new long[5][WIDTH];
    for (int i = 0; i < WIDTH; i++) {
      fixed[0][i] = 0;
      fixed[1][i] = i;
      fixed[2][i] = P - 1;
      fixed[3][i] = ((long) i * 2654435761L) % P;
      fixed[4][i] = 0;
    }
    fixed[4][0] = 11;
    fixed[4][15] = 1;
    final long[][] states = new long[5 + 8][WIDTH];
    for (int i = 0; i < 5; i++) states[i] = fixed[i];
    long x = 123456789L;
    for (int k = 0; k < 8; k++) {
      for (int i = 0; i < WIDTH; i++) {
        x = (1103515245L * x + 12345L) % 2147483648L;
        states[5 + k][i] = x % P;
      }
    }
    return states;
  }

  static String jArr(final long[] a) {
    final StringBuilder sb = new StringBuilder("[");
    for (int i = 0; i < a.length; i++) {
      if (i > 0) sb.append(',');
      sb.append(a[i]);
    }
    return sb.append(']').toString();
  }

  static void emit() {
    final StringBuilder sb = new StringBuilder();
    sb.append('{');
    // constants
    final long[] rcExtFlat = new long[ROUNDS_F * WIDTH];
    for (int r = 0; r < ROUNDS_F; r++) for (int i = 0; i < WIDTH; i++) rcExtFlat[r * WIDTH + i] = RC_EXTERNAL[r][i];
    final long[] m4flat = new long[16];
    for (int r = 0; r < 4; r++) for (int c = 0; c < 4; c++) m4flat[r * 4 + c] = M4[r][c];
    sb.append("\"constants\":{")
        .append("\"P\":").append(P)
        .append(",\"rInv\":").append(R_INV)
        .append(",\"width\":").append(WIDTH)
        .append(",\"sboxDegree\":").append(SBOX_DEGREE)
        .append(",\"roundsF\":").append(ROUNDS_F)
        .append(",\"roundsP\":").append(ROUNDS_P)
        .append(",\"internalDiagM1\":").append(jArr(INTERNAL_DIAG_M1_16))
        .append(",\"m4\":").append(jArr(m4flat))
        .append(",\"rcExternalFlat\":").append(jArr(rcExtFlat))
        .append(",\"rcInternal\":").append(jArr(RC_INTERNAL))
        .append('}');
    sb.append(",\"permute\":[");
    final long[][] states = inputStates();
    for (int k = 0; k < states.length; k++) {
      if (k > 0) sb.append(',');
      final long[] out = permute(states[k]);
      final boolean rt = eq(permuteInverse(out), states[k]);
      sb.append("{\"in\":").append(jArr(states[k]))
          .append(",\"out\":").append(jArr(out))
          .append(",\"roundtrip\":").append(rt ? "true" : "false")
          .append('}');
    }
    sb.append(']');
    sb.append(",\"conformanceKats\":[");
    final String[] katNames = {"zeros", "iota", "testvec"};
    for (int k = 0; k < KAT_IN.length; k++) {
      if (k > 0) sb.append(',');
      final long[] out = permute(KAT_IN[k]);
      sb.append("{\"name\":\"").append(katNames[k]).append('"')
          .append(",\"in\":").append(jArr(KAT_IN[k]))
          .append(",\"out\":").append(jArr(out))
          .append(",\"match\":").append(eq(out, KAT_OUT[k]) ? "true" : "false")
          .append('}');
    }
    sb.append(']');
    sb.append('}');
    System.out.print(sb);
  }

  public static void main(final String[] args) {
    if (args.length > 0 && args[0].equals("--emit")) emit();
    else selfTest();
  }
}