# -*- coding: utf-8 -*-
"""
SOL AUDIT 5 — full graded gauge test.

Equivalence group for weight-graded maps (source weights (1,-1,-2), target
(-2,-1,1)): torus scalings AND equivariant shears:
  source S:  x -> a1*x
             y -> b1*y + q1*x*z            (xz has weight -1)
             z -> g1*z + (p0 + p1*x*y)*y^2 (y^2, xy^3 have weight -2)
  target T:  Y1 -> d1*Y1 + w1*Y2^2         (Y2^2 has weight -2)
             Y2 -> e1*Y2 + v1*Y1*Y3        (Y1Y3 has weight -1)
             Y3 -> z1*Y3
S and T are REQUIRED to be automorphisms: impose det J_S == const != 0 and
det J_T == const != 0 as coefficient equations.

Question: is each audit-3 specimen G expressible as T o F_alp o S?
YES for all -> Annie's rigidity is TRUE up to full graded gauge (statement
needs a gauge caveat, substance intact).
NO for some -> genuinely new counterexamples; her Rigidity Theorem is false.
"""
import sympy as sp
import json, time

t0 = time.time()
x, y, z = sp.symbols('x y z')
a1, b1, g1, d1, e1, z1 = sp.symbols('a1 b1 g1 d1 e1 z1')
q1, p0, p1, w1, v1 = sp.symbols('q1 p0 p1 w1 v1')
UNK = [a1, b1, g1, d1, e1, z1, q1, p0, p1, w1, v1]

F_alp = [(1+x*y)**3*z + y**2*(1+x*y)*(4+3*x*y),
         y + 3*x*(1+x*y)**2*z + 3*x*y**2*(4+3*x*y),
         2*x - 3*x**2*y - x**3*z]

# source substitution
Sx, Sy, Sz = a1*x, b1*y + q1*x*z, g1*z + (p0 + p1*x*y)*y**2
JS = sp.Matrix([Sx, Sy, Sz]).jacobian([x, y, z])
detS = sp.expand(JS.det())
pS = sp.Poly(detS, x, y, z)
eqs_S = [co for mo, co in zip(pS.monoms(), pS.coeffs()) if mo != (0, 0, 0)]
constS = pS.coeff_monomial(1)

FS = [sp.expand(f.subs({x: Sx, y: Sy, z: Sz}, simultaneous=True)) for f in F_alp]
# target composition
G1c = sp.expand(d1*FS[0] + w1*FS[1]**2)
G2c = sp.expand(e1*FS[1] + v1*FS[0]*FS[2])
G3c = sp.expand(z1*FS[2])
Y1, Y2, Y3 = sp.symbols('Y1 Y2 Y3')
JT = sp.Matrix([d1*Y1 + w1*Y2**2, e1*Y2 + v1*Y1*Y3, z1*Y3]).jacobian([Y1, Y2, Y3])
detT = sp.expand(JT.det())
pT = sp.Poly(detT, Y1, Y2, Y3)
eqs_T = [co for mo, co in zip(pT.monoms(), pT.coeffs()) if mo != (0, 0, 0)]
constT = pT.coeff_monomial(1)

u_ = 1 + x*y
aa = x*y

def build(co, slope):
    P  = sp.sympify(co['P0']) + sp.sympify(co['P1'])*aa + sp.sympify(co['P2'])*aa**2
    Qt = sp.sympify(co['Q0']) + sp.sympify(co['Q1'])*aa + sp.sympify(co['Q2'])*aa**2
    Rr = sp.sympify(co['R0']) + sp.sympify(co['R1'])*aa
    return [sp.expand(y**2*P + z*slope[0]),
            sp.expand(y*Qt + z*slope[1]),
            sp.expand(x*Rr + z*slope[2])]

def full_gauge_equiv(G, tag):
    eqs = list(eqs_S) + list(eqs_T)
    for comp, target in ((G1c, G[0]), (G2c, G[1]), (G3c, G[2])):
        d = sp.expand(comp - target)
        eqs += sp.Poly(d, x, y, z).coeffs()
    sols = sp.solve(eqs, UNK, dict=True)
    for so in sols:
        cS = sp.simplify(constS.subs(so)); cT = sp.simplify(constT.subs(so))
        if cS != 0 and cT != 0:
            print(f"[GAUGE-EQUIV] {tag}: S,T found (detS={cS}, detT={cT})  "
                  f"{ {str(k): str(v) for k, v in so.items()} }", flush=True)
            return True
    print(f"[*** STILL NEW ***] {tag}: no automorphism pair (S,T) exists in shear class", flush=True)
    return False

R = json.load(open('post-jc-program/audit-sol/audit3_results.json'))
vA = (u_**3, 3*x*u_**2, -x**3)
tested = equiv = 0
reps = []
# representative from probe A (R0=1/2)
for entry in R.get("probeA", []):
    for br in entry["result"]:
        for smp in br.get("samples", []):
            if isinstance(smp, dict) and 'SPECIMEN' in str(smp.get("verdict", "")) and smp.get("coeffs"):
                reps.append((f"probeA m=3 R0={smp['coeffs'].get('R0')}", smp["coeffs"], vA))
# representatives from probe B k=3 (both slope families)
seen_slopes = set()
for entry in R.get("probeB", []):
    if entry["k"] != 3: continue
    for br in entry.get("integrability_branches", []):
        for kk in br.get("keller", []):
            for rr in kk.get("result", []):
                for smp in rr.get("samples", []):
                    if isinstance(smp, dict) and 'SPECIMEN' in str(smp.get("verdict", "")) and smp.get("coeffs"):
                        key = (br["psi"], br["chi"], kk["slope_vals"])
                        if key in seen_slopes: continue
                        seen_slopes.add(key)
                        vals = {sp.Symbol(k2.strip()): sp.sympify(v2) for k2, v2 in
                                (pair.split(':') for pair in kk["slope_vals"].strip('{}').split(',') if ':' in pair)}
                        psi_i = sp.expand(sp.sympify(br["psi"]).subs(vals))
                        chi_i = sp.expand(sp.sympify(br["chi"]).subs(vals))
                        slope = (u_**3, sp.expand(x*psi_i), sp.expand(x**3*chi_i))
                        reps.append((f"probeB k=3 psi={psi_i} chi={chi_i}", smp["coeffs"], slope))

print(f"testing {len(reps)} representative specimens against the FULL graded gauge group...", flush=True)
out = []
for tag, co, slope in reps:
    G = build(co, slope)
    ok = full_gauge_equiv(G, tag)
    tested += 1; equiv += int(ok)
    out.append({"tag": tag, "gauge_equiv": ok})

print(f"\nRESULT: {tested} tested, {equiv} equivalent under full graded gauge, {tested-equiv} new", flush=True)
json.dump(out, open('post-jc-program/audit-sol/audit5_verdicts.json', 'w'), indent=2)
print(f"({round(time.time()-t0,1)}s)", flush=True)