// Standalone (no-Besu-classpath) reference + self-test for the GENERIC AIR constraint / quotient-
// consistency check, component (e) of the PQ STARK-verify precompile 0x0AE8 (direct SP1/Plonky3 inner
// FRI/STARK verify). Mirrors air_quotient_reference.py and air_quotient_reference.mjs. It re-checks the
// GENERIC quotient identity for real p3-uni-stark 0.4.3-succinct proofs of two KNOWN example AIRs (the
// Fibonacci AIR from tests/fib_air.rs; a degree-3 multiply AIR matching tests/mul_air.rs), emitted by the
// ground-truth extractor (pq-stark/airquotient-extractor) and accepted by the pinned library verifier.
//
// HONEST SCOPE: this implements the GENERIC quotient-consistency MECHANISM (p3-uni-stark verifier.rs
// lines 90-141) for a SUPPLIED AIR constraint evaluator: fold constraints with alpha (Horner),
// reconstruct quotient(zeta) from chunk openings + split-domain zps, compute Z_H(zeta) and the selectors,
// and check folded_constraints(zeta) == Z_H(zeta) * quotient(zeta), over F_{p^4}. The SP1 recursion AIR
// (its specific constraint set + vkey binding) is NOT ported, so the top-level 0x0AE8 STAYS FAIL-CLOSED
// for real SP1 proofs. The F_{p^4} arithmetic below matches the precompile's BabyBearExt4 (component (a)).
//
//   javac -d out AirQuotientSelfTest.java
//   java  -cp out AirQuotientSelfTest [ground_truth.json]           # run the conformance self-test
//   java  -cp out AirQuotientSelfTest --emit [ground_truth.json]    # emit the shared vector set (JSON)

import java.math.BigInteger;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public final class AirQuotientSelfTest {

  static final long P = 2013265921L; // BabyBear 2^31 - 2^27 + 1
  static final long GENERATOR = 31L; // BabyBear multiplicative generator (order p-1)
  static final long EXT_W = 11L; // F_{p^4} = F_p[x]/(x^4 - 11) non-residue

  // ================= 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 mul(final long a, final long b) { return (reduce(a) * reduce(b)) % P; }
  static long powBase(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 invBase(final long a) { return powBase(reduce(a), P - 2); }
  static long twoAdicGenerator(final int bits) { return powBase(GENERATOR, (P - 1) >> bits); }

  // ================= F_{p^4} = F_p[x]/(x^4 - 11), elements long[4] = [c0,c1,c2,c3] =================
  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[] 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(EXT_W, t[4])), add(t[1], mul(EXT_W, t[5])), add(t[2], mul(EXT_W, t[6])), t[3]};
  }
  static long[] extPow(final long[] a, final BigInteger e) {
    long[] base = a.clone(); 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;
  }
  static long[] extInv(final long[] a) {
    if (a[0] == 0 && a[1] == 0 && a[2] == 0 && a[3] == 0) return new long[]{0, 0, 0, 0};
    final BigInteger p = BigInteger.valueOf(P);
    return extPow(a, p.pow(4).subtract(BigInteger.TWO));
  }
  static long[] extDiv(final long[] a, final long[] b) { return extMul(a, extInv(b)); }
  static boolean extEq(final long[] a, final long[] b) {
    for (int i = 0; i < 4; i++) if (reduce(a[i]) != reduce(b[i])) return false;
    return true;
  }
  static final long[] ONE = new long[]{1, 0, 0, 0};

  static long[] extExpPow2(final long[] x, final int logN) {
    long[] r = x.clone();
    for (int k = 0; k < logN; k++) r = extMul(r, r);
    return r;
  }

  // ================= domain / selector math (p3-commit domain.rs) =================
  /** Z_D(point) for a coset of size 2^logN and base-field shift: (point*shift^-1)^(2^logN) - 1. */
  static long[] zpAtPoint(final int logN, final long shiftBase, final long[] pointExt) {
    final long[] shiftInv = extFromBase(invBase(shiftBase));
    return extSub(extExpPow2(extMul(pointExt, shiftInv), logN), ONE);
  }

  static final class Selectors { long[] isFirst, isLast, isTransition, invZeroifier; }

  /** domain.rs selectors_at_point for the trace domain (shift=1, logN=degreeBits). */
  static Selectors selectorsAtPoint(final int degreeBits, final long[] zeta) {
    final long g = twoAdicGenerator(degreeBits);
    final long[] gInv = extFromBase(invBase(g));
    final long[] unshifted = zeta.clone();
    final long[] zH = extSub(extExpPow2(unshifted, degreeBits), ONE);
    final Selectors s = new Selectors();
    s.isFirst = extDiv(zH, extSub(unshifted, ONE));
    s.isLast = extDiv(zH, extSub(unshifted, gInv));
    s.isTransition = extSub(unshifted, gInv);
    s.invZeroifier = extInv(zH);
    return s;
  }

  static long[] monomial(final int e) {
    final long[] m = new long[4];
    m[e] = 1;
    return m;
  }

  /** Reconstruct quotient(zeta) from the chunk openings (verifier.rs lines 90-116). */
  static long[] reconstructQuotient(final int degreeBits, final long[][][] chunks, final long[] zeta) {
    final int quotientDegree = chunks.length;
    final int logQuotientDegree = Integer.numberOfTrailingZeros(quotientDegree);
    final int logNq = degreeBits + logQuotientDegree;
    final long genQ = twoAdicGenerator(logNq);
    final long[] shifts = new long[quotientDegree];
    for (int i = 0; i < quotientDegree; i++) shifts[i] = mul(GENERATOR, powBase(genQ, i));

    final long[][] zps = new long[quotientDegree][];
    for (int i = 0; i < quotientDegree; i++) {
      long[] acc = ONE.clone();
      final long[] firstPointI = extFromBase(shifts[i]);
      for (int j = 0; j < quotientDegree; j++) {
        if (j == i) continue;
        final long[] num = zpAtPoint(degreeBits, shifts[j], zeta);
        final long[] den = zpAtPoint(degreeBits, shifts[j], firstPointI);
        acc = extMul(acc, extMul(num, extInv(den)));
      }
      zps[i] = acc;
    }

    long[] quotient = new long[]{0, 0, 0, 0};
    for (int chI = 0; chI < chunks.length; chI++) {
      for (int eI = 0; eI < chunks[chI].length; eI++) {
        quotient = extAdd(quotient, extMul(extMul(zps[chI], monomial(eI)), chunks[chI][eI]));
      }
    }
    return quotient;
  }

  // ================= example AIR constraint evaluators (transcribed from their eval()) =================
  static long[] horner(final long[][] constraints, final long[] alpha) {
    long[] acc = new long[]{0, 0, 0, 0};
    for (final long[] c : constraints) acc = extAdd(extMul(acc, alpha), c);
    return acc;
  }

  static long[] foldFibonacci(final long[][] tl, final long[][] tn, final long[] pis, final Selectors s, final long[] alpha) {
    final long[] a = extFromBase(pis[0]);
    final long[] b = extFromBase(pis[1]);
    final long[] x = extFromBase(pis[2]);
    final long[] left = tl[0], right = tl[1];
    final long[] nleft = tn[0], nright = tn[1];
    final long[] c1 = extMul(s.isFirst, extSub(left, a));
    final long[] c2 = extMul(s.isFirst, extSub(right, b));
    final long[] c3 = extMul(s.isTransition, extSub(right, nleft));
    final long[] c4 = extMul(s.isTransition, extSub(extAdd(left, right), nright));
    final long[] c5 = extMul(s.isLast, extSub(right, x));
    return horner(new long[][]{c1, c2, c3, c4, c5}, alpha);
  }

  static long[] foldMulDeg3(final long[][] tl, final long[][] tn, final long[] pis, final Selectors s, final long[] alpha) {
    final long[] a = tl[0], b = tl[1], c = tl[2];
    final long[] nextA = tn[0];
    final long[] c1 = extSub(extMul(extMul(a, a), b), c);
    final long[] c2 = extMul(s.isFirst, extSub(extAdd(extMul(a, a), ONE), b));
    final long[] c3 = extMul(s.isTransition, extSub(extAdd(a, ONE), nextA));
    return horner(new long[][]{c1, c2, c3}, alpha);
  }

  static long[] foldAir(final String air, final long[][] tl, final long[][] tn, final long[] pis, final Selectors s, final long[] alpha) {
    if ("fibonacci".equals(air)) return foldFibonacci(tl, tn, pis, s, alpha);
    if ("mul_deg3".equals(air)) return foldMulDeg3(tl, tn, pis, s, alpha);
    throw new IllegalArgumentException("unknown AIR: " + air);
  }

  // ================= case model + the generic identity check =================
  static final class AirCase {
    String name, air;
    int degreeBits, quotientDegree;
    boolean libraryAccept;
    long[] pis, alpha, zeta;
    long[][] traceLocal, traceNext;
    long[][][] chunks;
  }

  static final class Derived { long[] quotient, folded, invZeroifier; }

  static Derived derive(final AirCase c) {
    final Selectors s = selectorsAtPoint(c.degreeBits, c.zeta);
    final Derived d = new Derived();
    d.quotient = reconstructQuotient(c.degreeBits, c.chunks, c.zeta);
    d.folded = foldAir(c.air, c.traceLocal, c.traceNext, c.pis, s, c.alpha);
    d.invZeroifier = s.invZeroifier;
    return d;
  }

  static boolean checkIdentity(final AirCase c) {
    final Derived d = derive(c);
    return extEq(extMul(d.folded, d.invZeroifier), d.quotient);
  }

  static AirCase copyCase(final AirCase c) {
    final AirCase o = new AirCase();
    o.name = c.name; o.air = c.air; o.degreeBits = c.degreeBits; o.quotientDegree = c.quotientDegree;
    o.libraryAccept = c.libraryAccept; o.pis = c.pis.clone();
    o.alpha = c.alpha.clone(); o.zeta = c.zeta.clone();
    o.traceLocal = new long[c.traceLocal.length][]; for (int i = 0; i < c.traceLocal.length; i++) o.traceLocal[i] = c.traceLocal[i].clone();
    o.traceNext = new long[c.traceNext.length][]; for (int i = 0; i < c.traceNext.length; i++) o.traceNext[i] = c.traceNext[i].clone();
    o.chunks = new long[c.chunks.length][][];
    for (int i = 0; i < c.chunks.length; i++) { o.chunks[i] = new long[c.chunks[i].length][]; for (int j = 0; j < c.chunks[i].length; j++) o.chunks[i][j] = c.chunks[i][j].clone(); }
    return o;
  }

  static boolean checkCase(final AirCase c, final String tamper) {
    final AirCase cc = copyCase(c);
    if ("trace".equals(tamper)) cc.traceLocal[0][0] = (cc.traceLocal[0][0] + 1) % P;
    else if ("quotient".equals(tamper)) cc.chunks[0][0][0] = (cc.chunks[0][0][0] + 1) % P;
    else if ("alpha".equals(tamper)) cc.alpha[0] = (cc.alpha[0] + 1) % P;
    return checkIdentity(cc);
  }

  // ================= minimal JSON parser =================
  static final class JP {
    final String s; int i;
    JP(final String s) { this.s = s; }
    void ws() { while (i < s.length() && Character.isWhitespace(s.charAt(i))) i++; }
    Object val() {
      ws(); final char c = s.charAt(i);
      if (c == '{') return obj();
      if (c == '[') return arr();
      if (c == '"') return str();
      if (c == 't') { i += 4; return Boolean.TRUE; }
      if (c == 'f') { i += 5; return Boolean.FALSE; }
      if (c == 'n') { i += 4; return null; }
      return num();
    }
    Map<String, Object> obj() {
      final Map<String, Object> m = new HashMap<>(); i++; ws();
      if (s.charAt(i) == '}') { i++; return m; }
      while (true) { ws(); final String k = str(); ws(); i++; m.put(k, val()); ws();
        if (s.charAt(i) == ',') { i++; continue; } i++; break; }
      return m;
    }
    List<Object> arr() {
      final List<Object> a = new ArrayList<>(); i++; ws();
      if (s.charAt(i) == ']') { i++; return a; }
      while (true) { a.add(val()); ws();
        if (s.charAt(i) == ',') { i++; continue; } i++; break; }
      return a;
    }
    String str() {
      final StringBuilder b = new StringBuilder(); i++;
      while (s.charAt(i) != '"') { if (s.charAt(i) == '\\') i++; b.append(s.charAt(i++)); }
      i++; return b.toString();
    }
    Long num() {
      final int start = i;
      while (i < s.length() && "+-0123456789.eE".indexOf(s.charAt(i)) >= 0) i++;
      return Long.parseLong(s.substring(start, i));
    }
  }

  @SuppressWarnings("unchecked")
  static long[] la(final Object o) {
    final List<Object> l = (List<Object>) o;
    final long[] out = new long[l.size()];
    for (int i = 0; i < out.length; i++) out[i] = (Long) l.get(i);
    return out;
  }
  @SuppressWarnings("unchecked")
  static long[][] laa(final Object o) {
    final List<Object> l = (List<Object>) o;
    final long[][] out = new long[l.size()][];
    for (int i = 0; i < out.length; i++) out[i] = la(l.get(i));
    return out;
  }
  @SuppressWarnings("unchecked")
  static long[][][] laaa(final Object o) {
    final List<Object> l = (List<Object>) o;
    final long[][][] out = new long[l.size()][][];
    for (int i = 0; i < out.length; i++) out[i] = laa(l.get(i));
    return out;
  }

  @SuppressWarnings("unchecked")
  static AirCase[] parseGroundTruth(final String json, final long[] valGenOut) {
    final JP jp = new JP(json);
    final Map<String, Object> root = (Map<String, Object>) jp.val();
    valGenOut[0] = (Long) root.get("val_generator");
    final List<Object> cases = (List<Object>) root.get("cases");
    final AirCase[] out = new AirCase[cases.size()];
    for (int ci = 0; ci < out.length; ci++) {
      final Map<String, Object> cm = (Map<String, Object>) cases.get(ci);
      final AirCase c = new AirCase();
      c.name = (String) cm.get("name");
      c.air = (String) cm.get("air");
      c.degreeBits = (int) (long) (Long) cm.get("degree_bits");
      c.quotientDegree = (int) (long) (Long) cm.get("quotient_degree");
      c.libraryAccept = (Boolean) cm.get("library_accept");
      c.pis = la(cm.get("public_values"));
      c.alpha = la(cm.get("alpha"));
      c.zeta = la(cm.get("zeta"));
      c.traceLocal = laa(cm.get("trace_local"));
      c.traceNext = laa(cm.get("trace_next"));
      c.chunks = laaa(cm.get("quotient_chunks"));
      out[ci] = c;
    }
    return out;
  }

  // ================= JSON emit (shared cross-language vector set) =================
  static String ja(final long[] v) {
    final StringBuilder b = new StringBuilder("[");
    for (int i = 0; i < v.length; i++) { if (i > 0) b.append(','); b.append(reduce(v[i])); }
    return b.append(']').toString();
  }

  static void emit(final String path) throws Exception {
    final long[] valGenOut = new long[1];
    final AirCase[] cases = parseGroundTruth(new String(Files.readAllBytes(Paths.get(path))), valGenOut);
    final StringBuilder sb = new StringBuilder();
    sb.append("{\"valGenerator\":").append(GENERATOR).append(",\"cases\":[");
    for (int ci = 0; ci < cases.length; ci++) {
      final AirCase c = cases[ci];
      if (ci > 0) sb.append(',');
      final Derived d = derive(c);
      sb.append("{\"name\":\"").append(c.name).append("\"");
      sb.append(",\"air\":\"").append(c.air).append("\"");
      sb.append(",\"degreeBits\":").append(c.degreeBits);
      sb.append(",\"quotientDegree\":").append(c.quotientDegree);
      sb.append(",\"quotient\":").append(ja(d.quotient));
      sb.append(",\"folded\":").append(ja(d.folded));
      sb.append(",\"invZeroifier\":").append(ja(d.invZeroifier));
      sb.append(",\"accept\":").append(checkCase(c, null));
      sb.append(",\"rejectTrace\":").append(!checkCase(c, "trace"));
      sb.append(",\"rejectQuotient\":").append(!checkCase(c, "quotient"));
      sb.append(",\"rejectAlpha\":").append(!checkCase(c, "alpha"));
      sb.append("}");
    }
    sb.append("]}");
    System.out.print(sb);
  }

  // ================= self-test (conformance) =================
  static int pass = 0, fail = 0;
  static void check(final String name, final boolean cond) { if (cond) pass++; else { fail++; System.out.println("FAIL " + name); } }

  static void selfTest(final String path) throws Exception {
    final long[] valGenOut = new long[1];
    final AirCase[] cases = parseGroundTruth(new String(Files.readAllBytes(Paths.get(path))), valGenOut);
    check("sanity.val_generator", valGenOut[0] == GENERATOR);
    for (final AirCase c : cases) {
      check("library_accept." + c.name, c.libraryAccept);
      check("accept." + c.name, checkCase(c, null));
      check("reject.trace." + c.name, !checkCase(c, "trace"));
      check("reject.quotient." + c.name, !checkCase(c, "quotient"));
      check("reject.alpha." + c.name, !checkCase(c, "alpha"));
    }
    System.out.println("AirQuotientSelfTest PASS=" + pass + " FAIL=" + fail);
    if (fail != 0) { System.out.println("AIR_QUOTIENT_SELFTEST_FAILED"); System.exit(1); }
    System.out.println("AIR_QUOTIENT_SELFTEST_OK");
  }

  public static void main(final String[] args) throws Exception {
    boolean emit = false;
    String path = "air_quotient_ground_truth.json";
    for (final String a : args) {
      if (a.equals("--emit")) emit = true;
      else if (!a.startsWith("--")) path = a;
    }
    if (emit) emit(path);
    else selfTest(path);
  }
}