# -*- coding: utf-8 -*-
"""
SOL AUDIT 3 — adversarial probe of the RIGIDITY THEOREM.

Annie's claim: in class A=(y^2 P(a), y+xy^2 Q(a), xR(a)), slope v_m, only m=3
works and the solution is unique. Two structural concerns:

(S1) Her A2-shape pins the coefficient of y to exactly 1. The general
     torus-equivariant weight-(-1) z-free component is y*Qt(a) with Qt(0) FREE.
     Was uniqueness an artifact of that normalization?

(S2) Her slope v_m = (u^m, m x u^{m-1}, -x^m) is torus-equivariant ONLY at m=3
     (weight of -x^m must be 3). So her m != 3 exclusions test a mixed
     (non-equivariant-slope + equivariant-A) ansatz. The honest equivariant
     rigidity question: general slope v = (phi(a), x psi(a), x^3 chi(a)).
     Does integrability + Keller force Alpoge's slope up to gauge?

Probe (a): rerun her class with Qt(0)=q0 free, m=1..5.
Probe (b): pin phi=(1+a)^k for k=1..4 (equivariant now, chi general!), solve
           integrability for (psi,chi), then staged Keller for general A.
Certify every solution: exact det, fiber degree by elimination (deg1 = automorphism).
Incremental JSON. Compact prints.
"""
import sympy as sp
import json, time, random

t0 = time.time()
random.seed(4242)
x, y, z = sp.symbols('x y z')
OUT = 'post-jc-program/audit-sol/audit3_results.json'
RES = {"probeA": [], "probeB": []}
def save(): json.dump(RES, open(OUT, 'w'), indent=2, default=str)

def fiber_degs(F, tries=2):
    out = []
    for _ in range(tries):
        t1, t2, t3 = [sp.Rational(random.randint(5, 70), random.randint(1, 7)) for _ in range(3)]
        zs = sp.solve(F[2] - t3, z)
        if len(zs) != 1: return None
        G1 = sp.numer(sp.together(F[0].subs(z, zs[0]) - t1))
        G2 = sp.numer(sp.together(F[1].subs(z, zs[0]) - t2))
        Rr = sp.expand(sp.resultant(G1, G2, y))
        if Rr == 0: return None
        p = sp.Poly(Rr, x)
        k = min(mm[0] for mm in p.monoms())
        e = sp.expand(Rr / x**k)
        g = sp.gcd(e, sp.diff(e, x))
        out.append(int(sp.degree(sp.cancel(e / g), x)))
    return out

def staged_solve(v, dP, dQ, dR, q0_free, tag):
    """A = (y^2 P, y*Qt, x*R). Returns branch info list."""
    aa = x*y
    ps = sp.symbols(f'P0:{dP+1}'); qs = sp.symbols(f'Q0:{dQ+1}'); rs = sp.symbols(f'R0:{dR+1}')
    U = list(ps) + list(qs) + list(rs)
    P = sum(c*aa**i for i, c in enumerate(ps))
    Qt = (sum(c*aa**i for i, c in enumerate(qs)) if q0_free
          else 1 + aa*sum(c*aa**i for i, c in enumerate(qs)))
    R = sum(c*aa**i for i, c in enumerate(rs))
    A = sp.Matrix([y**2*P, y*Qt, x*R])
    Ax, Ay = A.diff(x), A.diff(y)
    vx, vy = v.diff(x), v.diff(y)
    M2 = sp.expand(sp.Matrix.hstack(vx, vy, v).det())
    if M2 != 0:
        return [{"tag": tag, "status": "slope fails integrability (z^2 layer != 0)"}]
    M1 = sp.expand(sp.Matrix.hstack(vx, Ay, v).det() + sp.Matrix.hstack(Ax, vy, v).det())
    M0 = sp.expand(sp.Matrix.hstack(Ax, Ay, v).det())
    eqs1 = sp.Poly(M1, x, y).coeffs()
    lin = sp.linsolve(eqs1, U)
    if lin == sp.EmptySet:
        return [{"tag": tag, "status": "stage1 EMPTY (excluded)"}]
    sol = list(lin)[0]
    sub1 = dict(zip(U, sol))
    free1 = sorted({s_ for e in sol for s_ in e.free_symbols if s_ in U}, key=str)
    M0s = sp.expand(M0.subs(sub1, simultaneous=True))
    p0 = sp.Poly(M0s, x, y)
    c_expr, eqs2 = sp.Integer(0), []
    for mono, co in zip(p0.monoms(), p0.coeffs()):
        (eqs2.append(co) if mono != (0, 0) else None)
        if mono == (0, 0): c_expr = co
    brs = (sp.solve(eqs2, free1, dict=True) if eqs2 else [dict()]) if free1 \
          else ([dict()] if all(sp.expand(e) == 0 for e in eqs2) else [])
    out = []
    for bi, br in enumerate(brs):
        cb = sp.expand(c_expr.subs(br))
        full = {k: sp.expand(vv.subs(br)) for k, vv in sub1.items()}
        rem = sorted({s_ for e in list(full.values()) + [cb] for s_ in e.free_symbols if s_ in U}, key=str)
        info = {"tag": tag, "branch": bi, "c": str(cb), "free": [str(r) for r in rem], "samples": []}
        vals_list = [dict()] if not rem else \
            [dict(zip(rem, tuple(random.choice([1, -1, 2, sp.Rational(1,2), 3]) for _ in rem))) for _ in range(3)]
        for vals in vals_list:
            cv = sp.expand(cb.subs(vals))
            if cv == 0:
                info["samples"].append({"vals": str(vals), "verdict": "c=0 not Keller"}); continue
            inst = {k: sp.expand(vv.subs(vals)) for k, vv in full.items()}
            Fi = [sp.expand((A[i] + z*v[i]).subs(inst, simultaneous=True)) for i in range(3)]
            dJ = sp.expand(sp.Matrix(Fi).jacobian([x, y, z]).det())
            if not (dJ == cv and dJ.free_symbols == set()):
                info["samples"].append({"vals": str(vals), "verdict": "det check FAILED"}); continue
            fd = fiber_degs(Fi)
            vd = ("AUTOMORPHISM (deg1)" if fd and set(fd) == {1}
                  else f"SPECIMEN deg{set(fd)}" if fd and min(fd) >= 2 else f"inconclusive {fd}")
            info["samples"].append({"vals": {str(k2): str(v2) for k2, v2 in vals.items()},
                                    "coeffs": {str(k2): str(v2) for k2, v2 in inst.items()},
                                    "verdict": vd})
        out.append(info)
    return out

u_ = 1 + x*y
print("=== PROBE A: Annie's class but with Qt(0)=q0 FREE, m=1..5 ===", flush=True)
for m_ in (1, 2, 3, 4, 5):
    v = sp.Matrix([u_**m_, m_*x*u_**(m_-1), -x**m_])
    res = staged_solve(v, 2, 2, 1, q0_free=True, tag=f"m={m_},q0free")
    for r in res:
        print(f"  m={m_}: {r.get('status', 'branch ' + str(r.get('branch')) + ' c=' + r.get('c','?') + ' free=' + str(r.get('free')))}", flush=True)
        for smp in r.get("samples", []):
            print(f"      sample -> {smp['verdict']}", flush=True)
    RES["probeA"].append({"m": m_, "result": res}); save()

print("\n=== PROBE B: TRULY equivariant class: phi=(1+a)^k, psi,chi GENERAL ===", flush=True)
aa = x*y
for k_ in (1, 2, 3, 4):
    gs = sp.symbols(f'g0:3'); hs = sp.symbols(f'h0:2')
    psi = sum(c*aa**i for i, c in enumerate(gs))
    chi = sum(c*aa**i for i, c in enumerate(hs))
    v = sp.Matrix([u_**k_, x*psi, x**3*chi])
    M2 = sp.expand(sp.Matrix.hstack(v.diff(x), v.diff(y), v).det())
    eqs = sp.Poly(M2, x, y).coeffs()
    sols = sp.solve(eqs, list(gs) + list(hs), dict=True)
    nontriv = []
    for so in sols:
        psi_s = sp.expand(psi.subs(so)); chi_s = sp.expand(chi.subs(so))
        if psi_s == 0 and chi_s == 0: continue
        nontriv.append((psi_s, chi_s, so))
    print(f"  k={k_}: integrability solutions (nontrivial): {len(nontriv)}", flush=True)
    entry = {"k": k_, "integrability_branches": []}
    for psi_s, chi_s, so in nontriv:
        freeg = sorted({s_ for e in so.values() for s_ in e.free_symbols} |
                       {g for g in list(gs)+list(hs) if g not in so}, key=str)
        print(f"      psi={psi_s}  chi={chi_s}  free={freeg}", flush=True)
        binfo = {"psi": str(psi_s), "chi": str(chi_s), "free": [str(f) for f in freeg], "keller": []}
        vals_list = [dict()] if not freeg else \
            [dict(zip(freeg, tuple(random.choice([1, -1, 2]) for _ in freeg))) for _ in range(2)]
        for vals in vals_list:
            pv = sp.expand(psi_s.subs(vals)); cv_ = sp.expand(chi_s.subs(vals))
            if pv == 0 and cv_ == 0: continue
            vinst = sp.Matrix([u_**k_, x*pv, x**3*cv_])
            res = staged_solve(vinst, 2, 2, 1, q0_free=True, tag=f"k={k_} slope({pv};{cv_})")
            for r in res:
                stat = r.get('status', f"branch{r.get('branch')} c={r.get('c')} free={r.get('free')}")
                print(f"        Keller[{vals}]: {stat}", flush=True)
                for smp in r.get("samples", []):
                    print(f"           -> {smp['verdict']}", flush=True)
                    if 'SPECIMEN' in str(smp.get('verdict','')):
                        print(f"              coeffs: {smp.get('coeffs')}", flush=True)
            binfo["keller"].append({"slope_vals": str(vals), "result": res})
        entry["integrability_branches"].append(binfo)
    RES["probeB"].append(entry); save()

print(f"\nAUDIT 3 complete ({round(time.time()-t0,1)}s). artifact: audit3_results.json", flush=True)
