# -*- coding: utf-8 -*-
"""
AUDIT 3 — THE HARD ADVERSARIAL TEST of the Rigidity Theorem.

Objection: the forge fixed the slope to v_m = (u^m, m x u^{m-1}, -x^m). Maybe
rigidity is an artifact of that choice. Here we take the FULL torus-equivariant
z-affine class under the (1,-1,-2) weights — slope completely free:

  F1 = y^2 P(a) + z S1(a)      P<=3, S1<=3
  F2 = y + x y^2 Q(a) + x z S2(a)   Q<=2, S2<=2
  F3 = x R(a) + x^3 z S3(a)         R<=2, S3<=1     (a = xy)

19 coefficient unknowns + c. Alpoge = one point theta*. We compute:
  (1) K(theta*) = 0 exactly (Keller equations vanish; c = -2),
  (2) M = dK/dtheta at theta*, exact nullspace  = infinitesimal deformations,
  (3) the span of KNOWN TRIVIAL directions (pre/post diagonal scalings and
      source z-shears z -> z + eps y^2 a^j, all staying in-class),
  (4) verdict: ker M == trivial span  => infinitesimally rigid MODULO
      automorphisms in the FULL equivariant z-affine class (slope not assumed);
      ker M > trivial span => mystery deformation directions exist -> list them.
"""
import sympy as sp
import json, time

t0 = time.time()
x, y, z = sp.symbols('x y z')
a = x*y

ps  = sp.symbols('p0:4');  s1s = sp.symbols('s10:14')
qs  = sp.symbols('q0:3');  s2s = sp.symbols('s20:23')
rs  = sp.symbols('r0:3');  s3s = sp.symbols('s30:32')
c   = sp.Symbol('c')
TH  = list(ps)+list(s1s)+list(qs)+list(s2s)+list(rs)+list(s3s)+[c]

def poly(cs): return sum(co*a**i for i, co in enumerate(cs))
F1 = y**2*poly(ps) + z*poly(s1s)
F2 = y + x*y**2*poly(qs) + x*z*poly(s2s)
F3 = x*poly(rs) + x**3*z*poly(s3s)
D  = sp.expand(sp.Matrix([F1,F2,F3]).jacobian([x,y,z]).det() - c)
K  = sp.Poly(D, x, y, z).coeffs()
print(f"Keller system: {len(K)} equations, {len(TH)} unknowns", flush=True)

theta_star = dict(zip(ps,  [4,7,3,0])); theta_star.update(zip(s1s, [1,3,3,1]))
theta_star.update(zip(qs,  [12,9,0]));  theta_star.update(zip(s2s, [3,6,3]))
theta_star.update(zip(rs,  [2,-3,0]));  theta_star.update(zip(s3s, [-1,0]))
theta_star[c] = -2

ok1 = all(sp.expand(k.subs(theta_star)) == 0 for k in K)
print(f"[{'PASS' if ok1 else 'FAIL'}] K(theta*) = 0: Alpoge's map satisfies the full-class "
      f"Keller system with c = -2", flush=True)

M = sp.Matrix([[sp.diff(k, v).subs(theta_star) for v in TH] for k in K])
ns = M.nullspace()
print(f"dim ker M (infinitesimal in-class Keller deformations at Alpoge) = {len(ns)}", flush=True)

# ---- trivial directions, computed by exact reparametrization -----------------
def extract(Fs):
    """project a map back to theta coordinates; None if it leaves the class."""
    out = {}
    specs = [ (Fs[0], [( (j, j+2, 0), ps[j]) for j in range(4)] + [((j, j, 1), s1s[j]) for j in range(4)]),
              (sp.expand(Fs[1]-y), [((j+1, j+2, 0), qs[j]) for j in range(3)] + [((j+1, j, 1), s2s[j]) for j in range(3)]),
              (Fs[2], [((j+1, j, 0), rs[j]) for j in range(3)] + [((j+3, j, 1), s3s[j]) for j in range(2)]) ]
    for comp, slots in specs:
        p = sp.Poly(sp.expand(comp), x, y, z)
        table = {mono: co for mono, co in zip(p.monoms(), p.coeffs())}
        for mono, sym in slots:
            out[sym] = table.pop(mono, 0)
        if table:  # residual monomials outside the class
            return None
    return out

e = sp.Symbol('epsilon')
F0 = [F1.subs(theta_star), sp.expand(F2.subs(theta_star)), F3.subs(theta_star)]
dirs, names = [], []
def add_dir(Fe, cval, name):
    ex = extract([sp.expand(f) for f in Fe])
    if ex is None:
        print(f"        [{name}] leaves the class -> not an in-class direction", flush=True)
        return
    vec = [sp.diff(ex.get(v_, 0), e).subs(e, 0) if v_ != c else sp.diff(cval, e).subs(e, 0)
           for v_ in TH]
    dirs.append(vec); names.append(name)

lam = 1 + e
# source scalings
add_dir([f.subs({x: lam*x}, simultaneous=True) for f in F0], -2*lam, "pre-scale x")
add_dir([f.subs({y: lam*y}, simultaneous=True) for f in F0], -2*lam, "pre-scale y")
add_dir([f.subs({z: lam*z}, simultaneous=True) for f in F0], -2*lam, "pre-scale z")
# target scalings
add_dir([lam*F0[0], F0[1], F0[2]], -2*lam, "post-scale F1")
add_dir([F0[0], lam*F0[1], F0[2]], -2*lam, "post-scale F2")
add_dir([F0[0], F0[1], lam*F0[2]], -2*lam, "post-scale F3")
# source z-shears  z -> z + e*y^2*a^j  (unipotent: c unchanged)
for j in (0, 1):
    add_dir([f.subs({z: z + e*y**2*a**j}, simultaneous=True) for f in F0],
            sp.Integer(-2)*sp.Integer(1), f"z-shear y^2 a^{j}")
# source y-shear candidate: y -> y*(1+e) is pre-scale y (already); x-shear x -> x(1+e) same.

Tspan = sp.Matrix(dirs).T if dirs else sp.zeros(len(TH), 0)
rankT = Tspan.rank()
print(f"trivial (automorphism) directions constructed: {len(dirs)}, rank of span = {rankT}", flush=True)

# every trivial direction must lie in ker M (self-test of M)
okT = all(all(sp.simplify(val) == 0 for val in M*sp.Matrix(d)) for d in dirs)
print(f"[{'PASS' if okT else 'FAIL'}] all trivial directions lie in ker M (consistency)", flush=True)

# final comparison
if ns:
    Kmat = sp.Matrix([[v for v in vec] for vec in [list(n) for n in ns]]).T
    joint = Kmat.row_join(Tspan)
    extra = len(ns) - rankT if joint.rank() == len(ns) else f"(span mismatch: joint rank {joint.rank()})"
else:
    extra = 0
print(f"dim ker M = {len(ns)}   vs   rank(trivial span) = {rankT}", flush=True)
verdict = (len(ns) == rankT)
print(f"[{'RIGID' if verdict else 'NOT-RIGID'}] "
      + ("ker M == automorphism directions: Alpoge's map is INFINITESIMALLY RIGID "
         "modulo automorphisms in the FULL equivariant z-affine class (slope free)."
         if verdict else
         f"MYSTERY DIRECTIONS EXIST: {extra} deformation direction(s) beyond automorphisms "
         "-> potential new specimen families; dumping kernel for inspection."), flush=True)

out = {"K_eqs": len(K), "unknowns": len(TH), "K_at_alpoge_zero": bool(ok1),
       "dim_ker": len(ns), "trivial_rank": rankT, "trivial_in_kernel": bool(okT),
       "verdict_rigid_mod_autos": bool(verdict),
       "kernel_vectors": [[str(v) for v in list(n)] for n in ns],
       "trivial_names": names, "elapsed": round(time.time()-t0, 2)}
json.dump(out, open('post-jc-program/audit/audit3.json','w'), indent=2)
print(f"artifact: audit3.json  ({out['elapsed']}s)", flush=True)
