#!/usr/bin/env python3
# Independent reference for the BabyBear field F_p and its degree-4 extension F_{p^4}, the
# FOUNDATION layer (component (a)) of the PQ STARK-verify precompile 0x0AE8 (direct SP1/Plonky3
# inner FRI/STARK verify). See docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 2.
#
# HONEST SCOPE. This is field arithmetic, NOT a STARK verifier. It verifies nothing about any
# proof. It is the bottom of the six-component port stack; every higher component (Poseidon2, the
# Merkle/MMCS commitment, FRI folding, the AIR) is built on it and none of them is implemented in
# this build. The top-level 0x0AE8 precompile stays fail-closed (returns EMPTY for every input)
# regardless of this module.
#
# What is proven offline here (no external fetch needed):
#   - the field axioms for F_p and F_{p^4} (associativity, distributivity, inverses, ...),
#   - p = 2^31 - 2^27 + 1 = 15 * 2^27 + 1 (2-adicity 27, p-1 = 2^27 * 3 * 5),
#   - the multiplicative generator has order exactly p-1 (full-factorization order check),
#   - the two-adic generator of the order-2^27 subgroup has order exactly 2^27,
#   - the extension non-residue W makes x^4 - W irreducible (W is a QNR and p == 1 mod 4), so
#     F_{p^4} = F_p[x]/(x^4 - W) is a genuine field,
#   - Fermat inverse agrees with an independent extended-Euclid bignum inverse (Python pow(a,-1,p)).
#
# [VERIFY] (cannot be confirmed offline against Plonky3 in this pass, marked, not faked):
#   - GENERATOR = 31 is A generator (proven) but that it is Plonky3 p3-baby-bear's chosen
#     multiplicative generator is [VERIFY] against the pinned SP1 v6.1.0 revision.
#   - W = 11 yields a real field (proven) but that it is Plonky3's exact BinomialExtensionField
#     <BabyBear, 4> non-residue is [VERIFY]/[MEASURE] against p3_baby_bear.
#   - the exact stored two_adic_generator(27) value Plonky3 uses (mine is a valid one derived from
#     the generator; order proven 2^27) is [VERIFY].

P = 2013265921            # 2^31 - 2^27 + 1 = 15 * 2^27 + 1 = 0x78000001
GENERATOR = 31            # [VERIFY] Plonky3 p3-baby-bear multiplicative generator; order proven = P-1
TWO_ADICITY = 27          # v2(P-1); p-1 = 2^27 * 3 * 5
W = 11                    # [VERIFY] Plonky3 BabyBear quartic non-residue; x^4 - 11 irreducibility proven
PM1_ODD_PRIMES = (3, 5)   # p-1 = 2^27 * 3 * 5; used for exact multiplicative-order checks


# ============================ base field F_p ============================

def reduce(a: int) -> int:
    return a % P


def add(a: int, b: int) -> int:
    return (a + b) % P


def sub(a: int, b: int) -> int:
    return (a - b) % P


def neg(a: int) -> int:
    return (-a) % P


def mul(a: int, b: int) -> int:
    return (a * b) % P


def pow_(base: int, e: int) -> int:
    """Exponentiation by squaring (deterministic, integer-only). e >= 0."""
    if e < 0:
        raise ValueError("use inv for negative exponents")
    b = base % P
    acc = 1
    while e > 0:
        if e & 1:
            acc = (acc * b) % P
        b = (b * b) % P
        e >>= 1
    return acc


def inv(a: int) -> int:
    """Multiplicative inverse via Fermat a^(p-2). inv(0) := 0 (callers guard div-by-zero), matching
    the Java precompile BabyBear.inv convention."""
    r = a % P
    if r == 0:
        return 0
    return pow_(r, P - 2)


def two_adic_generator(bits: int) -> int:
    """A generator of the order-2^bits subgroup of F_p^*, derived from the multiplicative generator
    as g^((p-1)/2^bits). Its order is exactly 2^bits (checked in the harness). bits <= TWO_ADICITY.
    [VERIFY] the exact value Plonky3 stores for two_adic_generator(bits)."""
    if not (0 <= bits <= TWO_ADICITY):
        raise ValueError("bits out of range [0, 27]")
    return pow_(GENERATOR, (P - 1) >> bits)


def multiplicative_order(a: int) -> int:
    """Exact order of a in F_p^* using the known factorization p-1 = 2^27 * 3 * 5."""
    r = a % P
    if r == 0:
        raise ValueError("0 has no multiplicative order")
    order = P - 1
    factors = [2] * TWO_ADICITY + list(PM1_ODD_PRIMES)
    for q in set(factors):
        while order % q == 0 and pow_(r, order // q) == 1:
            order //= q
    return order


def is_quadratic_non_residue(a: int) -> bool:
    """Euler's criterion: a is a QNR mod p iff a^((p-1)/2) == p-1 (== -1)."""
    return pow_(a % P, (P - 1) >> 1) == P - 1


# ==================== degree-4 extension F_{p^4} = F_p[x]/(x^4 - W) ====================
# Elements are 4-int lists [c0, c1, c2, c3] meaning c0 + c1*x + c2*x^2 + c3*x^3, x^4 = W.

def ext_from_base(a: int) -> list:
    return [a % P, 0, 0, 0]


def ext_add(a: list, b: list) -> list:
    return [(a[i] + b[i]) % P for i in range(4)]


def ext_sub(a: list, b: list) -> list:
    return [(a[i] - b[i]) % P for i in range(4)]


def ext_neg(a: list) -> list:
    return [(-a[i]) % P for i in range(4)]


def ext_mul(a: list, b: list) -> list:
    """Schoolbook convolution reduced by x^4 = W (matches the Java precompile BabyBearExt4.mul)."""
    t = [0] * 7
    for i in range(4):
        for j in range(4):
            t[i + j] = (t[i + j] + a[i] * b[j]) % P
    r0 = (t[0] + W * t[4]) % P
    r1 = (t[1] + W * t[5]) % P
    r2 = (t[2] + W * t[6]) % P
    r3 = t[3] % P
    return [r0, r1, r2, r3]


def ext_pow(a: list, e: int) -> list:
    """Exponentiation by squaring in F_{p^4}. e may be a large (multi-limb) integer."""
    if e < 0:
        raise ValueError("use ext_inv for negative exponents")
    base = [a[i] % P for i in range(4)]
    acc = [1, 0, 0, 0]
    while e > 0:
        if e & 1:
            acc = ext_mul(acc, base)
        base = ext_mul(base, base)
        e >>= 1
    return acc


def ext_frobenius(a: list) -> list:
    """The Frobenius endomorphism pi(a) = a^p. On F_{p^4} it generates the degree-4 Galois group,
    so Frobenius^4 == identity, and it fixes the base field. Computed generically via ext_pow(a, p)
    (a binomial extension also admits the closed form pi(a)_k = a_k * (x^p)_k, but the generic power
    is used here as the independent, obviously-correct reference)."""
    return ext_pow(a, P)


def ext_inv(a: list) -> list:
    """Inverse via Fermat in F_{p^4}: a^(p^4 - 2). Defined as [0,0,0,0] for the zero element (matching
    the base-field inv(0):=0 convention; callers guard div-by-zero)."""
    if a[0] % P == 0 and a[1] % P == 0 and a[2] % P == 0 and a[3] % P == 0:
        return [0, 0, 0, 0]
    return ext_pow(a, P ** 4 - 2)


def ext_is_one(a: list) -> bool:
    return a[0] % P == 1 and a[1] % P == 0 and a[2] % P == 0 and a[3] % P == 0


def ext_eq(a: list, b: list) -> bool:
    return all((a[i] - b[i]) % P == 0 for i in range(4))


# ==================== shared cross-language vector set ====================
# A fixed set of inputs all three languages (Python, Node, Java) evaluate; the harness parses each
# language's JSON and asserts byte-identical results. Three independent implementations agreeing
# validates the arithmetic against the public field definition.

_BASE_UNARY = [0, 1, 2, 31, 1000000, 123456789, P - 1]
_BASE_BINARY = [(2, 3), (P - 1, 1), (1000000, 999), (123456789, 987654321), (0, 5)]
_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],
]
_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]),
]
_POW_E = 12345


def shared_vectors() -> dict:
    return {
        "constants": {
            "P": P,
            "generator": GENERATOR,
            "twoAdicity": TWO_ADICITY,
            "twoAdicGen27": two_adic_generator(27),
            "W": W,
            "powE": _POW_E,
        },
        "base_unary": [
            {"a": a, "neg": neg(a), "inv": inv(a), "pow": pow_(a, _POW_E)}
            for a in _BASE_UNARY
        ],
        "base_binary": [
            {"a": a, "b": b, "add": add(a, b), "sub": sub(a, b), "mul": mul(a, b)}
            for a, b in _BASE_BINARY
        ],
        "ext_unary": [
            {"a": a, "neg": ext_neg(a), "inv": ext_inv(a), "frob": ext_frobenius(a)}
            for a in _EXT_UNARY
        ],
        "ext_binary": [
            {"a": a, "b": b, "add": ext_add(a, b), "sub": ext_sub(a, b), "mul": ext_mul(a, b)}
            for a, b in _EXT_BINARY
        ],
    }


if __name__ == "__main__":
    import json
    import sys
    json.dump(shared_vectors(), sys.stdout)