# -*- coding: utf-8 -*-
"""
SOL AUDIT 4 — the decisive question from audit 3's finding:

Audit 3 found MORE Keller specimens than Annie's "unique point" (a free
parameter R0, plus sign-flipped slopes in the equivariant class). Are these
  (i) genuinely NEW counterexamples  -> her Rigidity Theorem is FALSE, or
 (ii) the gauge orbit of Alpoge's map under diagonal linear changes of
      coordinates  -> her theorem is TRUE-but-misstated (unique up to gauge)?

Test: for each specimen G found by audit 3, solve for diagonal source/target
scalings with  G(x,y,z) == diag(d,e,f) * F_alp(alpha*x, beta*y, gamma*z).
Exact coefficient matching, sp.solve. Nonempty solution = gauge-equivalent.
"""
import sympy as sp
import json, time

t0 = time.time()
x, y, z = sp.symbols('x y z')
al, be, ga, de, ep, ze = sp.symbols('alpha beta gamma delta epsilon zeta')

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]
Fg = [sp.expand(de*F_alp[0].subs({x: al*x, y: be*y, z: ga*z}, simultaneous=True)),
      sp.expand(ep*F_alp[1].subs({x: al*x, y: be*y, z: ga*z}, simultaneous=True)),
      sp.expand(ze*F_alp[2].subs({x: al*x, y: be*y, z: ga*z}, simultaneous=True))]

def gauge_equiv(G):
    eqs = []
    for i in range(3):
        d = sp.expand(Fg[i] - sp.expand(G[i]))
        eqs += sp.Poly(d, x, y, z).coeffs()
    sols = sp.solve(eqs, [al, be, ga, de, ep, ze], dict=True)
    good = []
    for so in sols:
        if all(sp.simplify(v) != 0 for v in so.values()):
            good.append({str(k): str(v) for k, v in so.items()})
    return good

R = json.load(open('post-jc-program/audit-sol/audit3_results.json'))
u_ = 1 + x*y
aa = x*y
tested, equiv, newfound = 0, 0, []

def build_from_coeffs(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])]

# Probe A specimens: slope = v_3 = (u^3, 3x u^2, -x^3)
vA = (u_**3, 3*x*u_**2, -x**3)
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"):
                G = build_from_coeffs(smp["coeffs"], vA)
                dJ = sp.expand(sp.Matrix(G).jacobian([x, y, z]).det())
                assert dJ.free_symbols == set() and dJ != 0
                sols = gauge_equiv(G)
                tested += 1
                if sols: equiv += 1
                else: newfound.append(("probeA", smp["coeffs"]))
                print(f"[{'GAUGE-EQUIV' if sols else '*** NEW ***'}] probeA m=3 specimen "
                      f"R0={smp['coeffs'].get('R0')}  det={dJ}"
                      + (f"  gauge={sols[0]}" if sols else ""), flush=True)

# Probe B specimens: reconstruct slope from branch psi/chi (k=3 entries only)
for entry in R.get("probeB", []):
    k = entry["k"]
    for br in entry.get("integrability_branches", []):
        psi_t, chi_t = sp.sympify(br["psi"]), sp.sympify(br["chi"])
        for kk in br.get("keller", []):
            sv = kk["slope_vals"]
            try:
                vals = {sp.Symbol(k2.strip()): sp.sympify(v2) for k2, v2 in
                        (pair.split(':') for pair in sv.strip('{}').split(',') if ':' in pair)}
            except Exception:
                vals = {}
            psi_i = sp.expand(psi_t.subs(vals)); chi_i = sp.expand(chi_t.subs(vals))
            slope = (u_**k, sp.expand(x*psi_i), sp.expand(x**3*chi_i))
            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"):
                        G = build_from_coeffs(smp["coeffs"], slope)
                        dJ = sp.expand(sp.Matrix(G).jacobian([x, y, z]).det())
                        if not (dJ.free_symbols == set() and dJ != 0):
                            continue
                        sols = gauge_equiv(G)
                        tested += 1
                        if sols: equiv += 1
                        else: newfound.append((f"probeB k={k}", smp["coeffs"]))
                        print(f"[{'GAUGE-EQUIV' if sols else '*** NEW ***'}] probeB k={k} "
                              f"slope psi={psi_i} chi={chi_i}  det={dJ}"
                              + (f"  gauge={sols[0]}" if sols else ""), flush=True)

print(f"\nRESULT: {tested} specimens tested, {equiv} gauge-equivalent to Alpoge, "
      f"{len(newfound)} genuinely new", flush=True)
json.dump({"tested": tested, "gauge_equivalent": equiv, "new": newfound},
          open('post-jc-program/audit-sol/audit4_verdicts.json', 'w'), indent=2, default=str)
print(f"({round(time.time()-t0,1)}s) artifact: audit4_verdicts.json", flush=True)