// Standalone (no-Besu-classpath) copy of the BabyBear field + degree-4 extension arithmetic used by
// the PQ STARK-verify precompile 0x0AE8, plus a self-test. It is a verbatim mirror of the
// Sp1StarkVerifierPrecompiledContract.BabyBear / BabyBearExt4 logic (same P, same W, same reduction),
// compilable and runnable without the Besu EVM jars so the field layer can be exercised offline and
// cross-checked against the Python and Node references.
//
// HONEST SCOPE. This is field arithmetic, the FOUNDATION component (a) of the six-component STARK
// port (docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 2). It verifies NOTHING about a real proof;
// the top-level 0x0AE8 precompile stays fail-closed. GENERATOR = 31 and W = 11 are proven here to be
// a real generator and a real non-residue (so F_{p^4} is a field), but that they are Plonky3's exact
// chosen constants is [VERIFY] against p3-baby-bear at the pinned SP1 v6.1.0 revision.
//
// Usage:
//   javac -d out BabyBearFieldSelfTest.java
//   java  -cp out BabyBearFieldSelfTest            # run the self-test (prints PASS/FAIL counts)
//   java  -cp out BabyBearFieldSelfTest --emit     # emit the shared cross-language vector set (JSON)

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

public final class BabyBearFieldSelfTest {

  static final long P = 2013265921L; // 2^31 - 2^27 + 1 = 15 * 2^27 + 1 = 0x78000001
  static final long GENERATOR = 31L; // [VERIFY] Plonky3 generator; order proven = P-1
  static final int TWO_ADICITY = 27;
  static final long W = 11L;         // [VERIFY] Plonky3 BabyBear quartic non-residue; irreducibility proven
  static final BigInteger BP = BigInteger.valueOf(P);

  // ---- base field F_p ----
  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 neg(final long a) {
    final long r = reduce(a);
    return r == 0 ? 0 : P - r;
  }

  static long mul(final long a, final long b) {
    return (reduce(a) * reduce(b)) % P; // both < 2^31 so product < 2^62, safe in signed 64-bit
  }

  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 inv(final long a) {
    final long r = reduce(a);
    if (r == 0) return 0; // inv(0) := 0, callers guard div-by-zero
    return pow(r, P - 2);
  }

  static long twoAdicGenerator(final int bits) {
    if (bits < 0 || bits > TWO_ADICITY) throw new IllegalArgumentException("bits out of range");
    return pow(GENERATOR, (P - 1) >> bits);
  }

  /** Exact multiplicative order using the known factorization p-1 = 2^27 * 3 * 5. */
  static long multiplicativeOrder(final long a) {
    final long r = reduce(a);
    if (r == 0) throw new IllegalArgumentException("0 has no order");
    long order = P - 1;
    final int[] primes = {2, 3, 5};
    for (final int q : primes) {
      while (order % q == 0 && pow(r, order / q) == 1) order /= q;
    }
    return order;
  }

  static boolean isQuadraticNonResidue(final long a) {
    return pow(reduce(a), (P - 1) >> 1) == P - 1;
  }

  // ---- degree-4 extension F_{p^4} = F_p[x]/(x^4 - W); elements are long[4] = c0 + c1 x + c2 x^2 + c3 x^3 ----
  static long[] extFromBase(final long a) {
    return new long[] {reduce(a), 0, 0, 0};
  }

  static long[] extAdd(final long[] a, final long[] b) {
    return new long[] {add(a[0], b[0]), add(a[1], b[1]), add(a[2], b[2]), add(a[3], b[3])};
  }

  static long[] extSub(final long[] a, final long[] b) {
    return new long[] {sub(a[0], b[0]), sub(a[1], b[1]), sub(a[2], b[2]), sub(a[3], b[3])};
  }

  static long[] extNeg(final long[] a) {
    return new long[] {neg(a[0]), neg(a[1]), neg(a[2]), neg(a[3])};
  }

  /** Schoolbook convolution reduced by x^4 = W (identical to the precompile BabyBearExt4.mul). */
  static long[] extMul(final long[] a, final long[] b) {
    final long[] t = new long[7];
    for (int i = 0; i < 4; i++) {
      for (int j = 0; j < 4; j++) {
        t[i + j] = add(t[i + j], mul(a[i], b[j]));
      }
    }
    return new long[] {
      add(t[0], mul(W, t[4])), add(t[1], mul(W, t[5])), add(t[2], mul(W, t[6])), reduce(t[3])
    };
  }

  static long[] extPow(final long[] a, final BigInteger e) {
    if (e.signum() < 0) throw new IllegalArgumentException("use extInv for negative exponents");
    long[] base = new long[] {reduce(a[0]), reduce(a[1]), reduce(a[2]), reduce(a[3])};
    long[] acc = new long[] {1, 0, 0, 0};
    BigInteger ee = e;
    while (ee.signum() > 0) {
      if (ee.testBit(0)) acc = extMul(acc, base);
      base = extMul(base, base);
      ee = ee.shiftRight(1);
    }
    return acc;
  }

  /** Frobenius pi(a) = a^p; Frobenius^4 == id and fixes the base field. */
  static long[] extFrobenius(final long[] a) {
    return extPow(a, BP);
  }

  static long[] extInv(final long[] a) {
    if (reduce(a[0]) == 0 && reduce(a[1]) == 0 && reduce(a[2]) == 0 && reduce(a[3]) == 0) {
      return new long[] {0, 0, 0, 0};
    }
    return extPow(a, BP.pow(4).subtract(BigInteger.TWO)); // a^(p^4 - 2)
  }

  static boolean extEq(final long[] a, final long[] b) {
    return reduce(a[0]) == reduce(b[0]) && reduce(a[1]) == reduce(b[1])
        && reduce(a[2]) == reduce(b[2]) && reduce(a[3]) == reduce(b[3]);
  }

  // ================================= 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 void selfTest() {
    // constants
    check("P.value", P == (1L << 31) - (1L << 27) + 1 && P == 15L * (1L << 27) + 1 && P == 0x78000001L);
    // generator has full order p-1 (real proof, not a Plonky3 claim)
    check("generator.order", multiplicativeOrder(GENERATOR) == P - 1);
    // two-adic generator order is exactly 2^27
    final long tag = twoAdicGenerator(27);
    check("twoadic.order", pow(tag, 1L << 27) == 1 && pow(tag, 1L << 26) != 1);
    for (int bits = 1; bits <= 12; bits++) {
      final long g = twoAdicGenerator(bits);
      check("twoadic.order.b" + bits, pow(g, 1L << bits) == 1 && pow(g, 1L << (bits - 1)) != 1);
    }
    // W = 11 is a QNR and p == 1 mod 4 => x^4 - W irreducible => F_{p^4} is a field
    check("W.qnr", isQuadraticNonResidue(W));
    check("p.mod4", P % 4 == 1);
    // x^4 == W in the extension (defining relation)
    final long[] x = {0, 1, 0, 0};
    check("x4.eq.W", extEq(extPow(x, BigInteger.valueOf(4)), extFromBase(W)));

    final Random rnd = new Random(0x0AE8);
    for (int it = 0; it < 20000; it++) {
      final long a = Math.floorMod(rnd.nextLong(), P);
      final long b = Math.floorMod(rnd.nextLong(), P);
      final long c = Math.floorMod(rnd.nextLong(), P);
      // base axioms
      check("add.comm", add(a, b) == add(b, a));
      check("mul.comm", mul(a, b) == mul(b, a));
      check("distrib", mul(a, add(b, c)) == add(mul(a, b), mul(a, c)));
      check("add.assoc", add(add(a, b), c) == add(a, add(b, c)));
      check("sub.def", add(sub(a, b), b) == a);
      check("neg.def", add(a, neg(a)) == 0);
      if (a != 0) {
        check("inv.def", mul(a, inv(a)) == 1);
        check("fermat", pow(a, P - 1) == 1); // Fermat little theorem
      }
      check("frobenius.base", pow(a, P) == a); // a^p == a on the base field
      // (p-1) behaves as -1
      check("minus1.sq", mul(P - 1, P - 1) == 1);
      check("minus1.add", add(P - 1, 1) == 0);

      // extension axioms
      final long[] ea = {a, b, c, add(a, c)};
      final long[] eb = {b, c, a, sub(b, a)};
      check("ext.add.comm", extEq(extAdd(ea, eb), extAdd(eb, ea)));
      check("ext.mul.comm", extEq(extMul(ea, eb), extMul(eb, ea)));
      check("ext.sub.def", extEq(extAdd(extSub(ea, eb), eb), ea));
      check("ext.neg.def", extEq(extAdd(ea, extNeg(ea)), new long[] {0, 0, 0, 0}));
      final boolean nz = a != 0 || b != 0 || c != 0 || add(a, c) != 0;
      if (nz) {
        check("ext.inv.def", extEq(extMul(ea, extInv(ea)), new long[] {1, 0, 0, 0}));
      }
      // Frobenius^4 == id, and Frobenius fixes base-embedded elements
      long[] f = ea;
      for (int k = 0; k < 4; k++) f = extFrobenius(f);
      check("ext.frob4.id", extEq(f, ea));
      check("ext.frob.base.fix", extEq(extFrobenius(extFromBase(a)), extFromBase(a)));
    }

    System.out.println("BabyBearFieldSelfTest PASS=" + pass + " FAIL=" + fail);
    if (fail != 0) {
      System.out.println("BABYBEAR_FIELD_SELFTEST_FAILED");
      System.exit(1);
    }
    System.out.println("BABYBEAR_FIELD_SELFTEST_OK");
  }

  // ================================= shared-vector emit (JSON) =================================
  static final long[] BASE_UNARY = {0, 1, 2, 31, 1000000, 123456789, P - 1};
  static final long[][] BASE_BINARY = {{2, 3}, {P - 1, 1}, {1000000, 999}, {123456789, 987654321}, {0, 5}};
  static final long[][] EXT_UNARY = {
    {1, 0, 0, 0}, {0, 1, 0, 0}, {2, 3, 5, 7}, {P - 1, P - 1, P - 1, P - 1}, {11, 0, 0, 0}, {123, 456, 789, 1011}
  };
  static final long[][][] EXT_BINARY = {
    {{1, 2, 3, 4}, {5, 6, 7, 8}}, {{0, 1, 0, 0}, {0, 1, 0, 0}}, {{2, 3, 5, 7}, {11, 0, 0, 0}}
  };
  static final long POW_E = 12345L;

  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('{');
    sb.append("\"constants\":{")
        .append("\"P\":").append(P)
        .append(",\"generator\":").append(GENERATOR)
        .append(",\"twoAdicity\":").append(TWO_ADICITY)
        .append(",\"twoAdicGen27\":").append(twoAdicGenerator(27))
        .append(",\"W\":").append(W)
        .append(",\"powE\":").append(POW_E)
        .append('}');
    sb.append(",\"base_unary\":[");
    for (int i = 0; i < BASE_UNARY.length; i++) {
      final long a = BASE_UNARY[i];
      if (i > 0) sb.append(',');
      sb.append("{\"a\":").append(a)
          .append(",\"neg\":").append(neg(a))
          .append(",\"inv\":").append(inv(a))
          .append(",\"pow\":").append(pow(a, POW_E))
          .append('}');
    }
    sb.append(']');
    sb.append(",\"base_binary\":[");
    for (int i = 0; i < BASE_BINARY.length; i++) {
      final long a = BASE_BINARY[i][0], b = BASE_BINARY[i][1];
      if (i > 0) sb.append(',');
      sb.append("{\"a\":").append(a).append(",\"b\":").append(b)
          .append(",\"add\":").append(add(a, b))
          .append(",\"sub\":").append(sub(a, b))
          .append(",\"mul\":").append(mul(a, b))
          .append('}');
    }
    sb.append(']');
    sb.append(",\"ext_unary\":[");
    for (int i = 0; i < EXT_UNARY.length; i++) {
      final long[] a = EXT_UNARY[i];
      if (i > 0) sb.append(',');
      sb.append("{\"a\":").append(jArr(a))
          .append(",\"neg\":").append(jArr(extNeg(a)))
          .append(",\"inv\":").append(jArr(extInv(a)))
          .append(",\"frob\":").append(jArr(extFrobenius(a)))
          .append('}');
    }
    sb.append(']');
    sb.append(",\"ext_binary\":[");
    for (int i = 0; i < EXT_BINARY.length; i++) {
      final long[] a = EXT_BINARY[i][0], b = EXT_BINARY[i][1];
      if (i > 0) sb.append(',');
      sb.append("{\"a\":").append(jArr(a)).append(",\"b\":").append(jArr(b))
          .append(",\"add\":").append(jArr(extAdd(a, b)))
          .append(",\"sub\":").append(jArr(extSub(a, b)))
          .append(",\"mul\":").append(jArr(extMul(a, b)))
          .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();
    }
  }
}