"""Stage 3 — symbolic verification (SymPy, exact) of the algebraic identities used in the paper.

Checks (each returns True or raises):
  S1  The affine addition and doubling formulas map curve points to curve points (identity modulo the curve equations).
  S2  φ(x, y) = (βx, y) with β³ = 1 is a group endomorphism: φ(P + Q) = φ(P) + φ(Q) and φ(2P) = 2φ(P) (a = 0).
  S3  The x-only differential addition formula for y² = x³ + b (Brier–Joye 2002 specialised to a = 0):
        x(P+Q) · x(P−Q) · (x_P − x_Q)² = (x_P x_Q − a)² − 4b(x_P + x_Q)    [a = 0 here]
  S4  Semaev's third summation polynomial S3 for y² = x³ + ax + b, derived from the group law, vanishes on x-coordinates
      of P, Q, P+Q and P, Q, P−Q (symmetric form), and its total degree is 2 in each variable.
  S5  ECDSA malleability: with s' = −s mod n the verification point is −X, whose x-coordinate equals that of X.
  S6  Schnorr correctness: s = k + e·d ⇒ sG = R + eP (as scalar identities).
Writes results/theory/symbolic_checks.json.
"""
import json, sys, time, pathlib
import sympy as sp
ROOT = pathlib.Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
from wp34.provenance import record

x1, y1, x2, y2, a, b, beta = sp.symbols("x1 y1 x2 y2 a b beta")


def reduce_mod_curve(expr, pts):
    """Reduce a rational expression modulo y_i² = x_i³ + a x_i + b for the given points, return canonical numerator."""
    num, den = sp.fraction(sp.together(expr))
    G = [yi**2 - (xi**3 + a * xi + b) for xi, yi in pts]
    gens = [yi for _, yi in pts] + [xi for xi, _ in pts] + [a, b]
    _, r = sp.reduced(sp.expand(num), G, *gens)
    return sp.simplify(r)


def add_formula(P, Q):
    (xa, ya), (xb, yb) = P, Q
    lam = (yb - ya) / (xb - xa)
    x3 = lam**2 - xa - xb
    y3 = lam * (xa - x3) - ya
    return x3, y3


def dbl_formula(P):
    xa, ya = P
    lam = (3 * xa**2 + a) / (2 * ya)
    x3 = lam**2 - 2 * xa
    y3 = lam * (xa - x3) - ya
    return x3, y3


def S1():
    x3, y3 = add_formula((x1, y1), (x2, y2))
    on = y3**2 - (x3**3 + a * x3 + b)
    r_add = reduce_mod_curve(on, [(x1, y1), (x2, y2)])
    x3d, y3d = dbl_formula((x1, y1))
    r_dbl = reduce_mod_curve(y3d**2 - (x3d**3 + a * x3d + b), [(x1, y1)])
    return r_add == 0 and r_dbl == 0


def S2():
    # a = 0 and beta^3 = 1
    x3, y3 = add_formula((x1, y1), (x2, y2))
    phi_sum = (beta * x3, y3)
    x3p, y3p = add_formula((beta * x1, y1), (beta * x2, y2))
    d1 = sp.together(phi_sum[0] - x3p); d2 = sp.together(phi_sum[1] - y3p)
    # reduce beta^3 → 1 : substitute beta^3 = 1 via polynomial remainder
    ok = True
    for d in (d1, d2):
        num, _ = sp.fraction(d)
        _, r = sp.reduced(sp.expand(num), [beta**3 - 1], beta, x1, y1, x2, y2)
        ok &= sp.simplify(r) == 0
    xd, yd = dbl_formula((x1, y1))
    xd = xd.subs(a, 0); yd = yd.subs(a, 0)
    xdp, ydp = dbl_formula((beta * x1, y1)); xdp = xdp.subs(a, 0); ydp = ydp.subs(a, 0)
    for d in (beta * xd - xdp, yd - ydp):
        num, _ = sp.fraction(sp.together(d))
        _, r = sp.reduced(sp.expand(num), [beta**3 - 1], beta, x1, y1)
        ok &= sp.simplify(r) == 0
    return bool(ok)


def S3():
    xs, _ = add_formula((x1, y1), (x2, y2))
    xd, _ = add_formula((x1, y1), (x2, -y2))
    lhs = xs * xd * (x1 - x2)**2
    rhs = (x1 * x2 - a)**2 - 4 * b * (x1 + x2)
    r = reduce_mod_curve(lhs - rhs, [(x1, y1), (x2, y2)])
    return r == 0


def S4():
    x3 = sp.symbols("x3")
    # Semaev (2004): S3(x1,x2,x3) = (x1−x2)²x3² − 2((x1+x2)(x1x2+a) + 2b)x3 + (x1x2−a)² − 4b(x1+x2)
    S3p = (x1 - x2)**2 * x3**2 - 2 * ((x1 + x2) * (x1 * x2 + a) + 2 * b) * x3 + (x1 * x2 - a)**2 - 4 * b * (x1 + x2)
    xs, _ = add_formula((x1, y1), (x2, y2))
    xd, _ = add_formula((x1, y1), (x2, -y2))
    ok = reduce_mod_curve(S3p.subs(x3, xs), [(x1, y1), (x2, y2)]) == 0
    ok &= reduce_mod_curve(S3p.subs(x3, xd), [(x1, y1), (x2, y2)]) == 0
    degs = [sp.Poly(S3p, v).degree() for v in (x1, x2, x3)]
    sym = sp.expand(S3p - S3p.subs({x1: x2, x2: x3, x3: x1}, simultaneous=True)) == 0
    return bool(ok and degs == [2, 2, 2] and sym)


def S5():
    n, z, r, s, d = sp.symbols("n z r s d")
    # verification point X = u1 G + u2 Q = (z/s) G + (r/s) Q. With s' = -s: X' = -(z/s) G - (r/s) Q = -X.
    w = 1 / s
    u1, u2 = z * w, r * w
    wp = 1 / (-s)
    u1p, u2p = z * wp, r * wp
    return sp.simplify(u1p + u1) == 0 and sp.simplify(u2p + u2) == 0  # scalars negate ⇒ point negates ⇒ same x


def S6():
    k, e, d = sp.symbols("k e d")
    s = k + e * d
    return sp.simplify(s - (k + e * d)) == 0  # sG = kG + e(dG) = R + eP by the Z-module structure


def main():
    t0 = time.time()
    res = {"S1_formulas_preserve_curve": S1(), "S2_phi_is_endomorphism": S2(), "S3_xonly_differential_addition": S3(),
           "S4_semaev_S3_polynomial": S4(), "S5_ecdsa_malleability_scalars": S5(), "S6_schnorr_correctness_scalars": S6()}
    record("symbolic_checks_sympy", {"sympy": sp.__version__}, res, out=str(ROOT / "results/theory/symbolic_checks.json"), started=t0)
    print(json.dumps(res, indent=1)); print("ALL:", all(res.values()))


if __name__ == "__main__":
    main()