# -*- coding: utf-8 -*-
"""
PHASE 2 — Normalize Alpoge's map to BCW input form  G = X + H,  det J_G = 1.

This is the entry gate for the Bass--Connell--Wright / Druzkowski reduction (1b):
every published reduction assumes a Keller map of the form identity + higher-order
terms with Jacobian determinant exactly 1. We produce it, verify it, and profile
the monomial structure H to size the lift.
Also: push the collision witnesses through the normalization (they must survive).
"""
import sympy as sp
import json, time

t0 = time.time()
x, y, z = sp.symbols('x y z')
X = sp.Matrix([x, y, z])

f1 = (1 + x*y)**3 * z + y**2 * (1 + x*y) * (4 + 3*x*y)
f2 = y + 3*x*(1 + x*y)**2 * z + 3*x*y**2 * (4 + 3*x*y)
f3 = 2*x - 3*x**2*y - x**3*z
F = sp.Matrix([f1, f2, f3])
J = F.jacobian([x, y, z])

# Linear part L = J(0)
L = J.subs({x: 0, y: 0, z: 0})
print("linear part L =", L.tolist(), " det L =", L.det())

# Constant part must be 0 for X + H form
F0 = F.subs({x: 0, y: 0, z: 0})
ok_const = all(c == 0 for c in F0)
print(f"[{'PASS' if ok_const else 'FAIL'}] F(0) = 0 (no translation needed)")

# Normalized map G = L^{-1} F
Gm = (L.inv() * F).applyfunc(sp.expand)
H = (Gm - X).applyfunc(sp.expand)

# check 1: linear part of G is identity <=> H has only degree >= 2 terms
ok_lin = True
min_deg = 99
max_deg = 0
mono_count = 0
profile = []
for i, h in enumerate(H):
    p = sp.Poly(h, x, y, z)
    degs = sorted({sum(m) for m in p.monoms()})
    if degs and degs[0] < 2:
        ok_lin = False
    min_deg = min(min_deg, degs[0] if degs else 99)
    max_deg = max(max_deg, degs[-1] if degs else 0)
    mono_count += len(p.monoms())
    profile.append({
        "component": f"H{i+1}",
        "n_monomials": len(p.monoms()),
        "degrees_present": degs,
        "H": str(h)
    })
print(f"[{'PASS' if ok_lin else 'FAIL'}] G = X + H with deg(H terms) in [{min_deg},{max_deg}], "
      f"total monomials in H: {mono_count}")

# check 2: det J_G == 1 identically (full expansion)
detG = sp.expand(Gm.jacobian([x, y, z]).det())
ok_det = (detG == 1)
print(f"[{'PASS' if ok_det else 'FAIL'}] det J_G == 1 identically  (got: {detG})")

# check 3: collisions survive normalization (L^{-1} is a bijection)
pts = [(sp.Integer(0), sp.Integer(0), sp.Rational(-1, 4)),
       (sp.Integer(1), sp.Rational(-3, 2), sp.Rational(13, 2)),
       (sp.Integer(-1), sp.Rational(3, 2), sp.Rational(13, 2))]
imgs = [tuple(sp.simplify(v) for v in Gm.subs({x: p[0], y: p[1], z: p[2]})) for p in pts]
ok_coll = (imgs[0] == imgs[1] == imgs[2])
print(f"[{'PASS' if ok_coll else 'FAIL'}] 3-point collision survives: G(p*) = {imgs[0]}")

out = {
    "normalized_map_G": [str(g) for g in Gm],
    "H_profile": profile,
    "L": [[str(L[i, j]) for j in range(3)] for i in range(3)],
    "checks": {"F0_zero": ok_const, "linear_part_identity": ok_lin,
               "detJG_equals_1": ok_det, "collision_survives": ok_coll},
    "collision_image_under_G": [str(c) for c in imgs[0]],
    "elapsed_sec": round(time.time() - t0, 2)
}
with open('post-jc-program/normalized_map.json', 'w', encoding='utf-8') as fh:
    json.dump(out, fh, indent=2)

all_ok = ok_const and ok_lin and ok_det and ok_coll
print()
print(f"{'='*60}")
print(f"PHASE 2: ALL_CHECKS_PASS = {all_ok}  ({out['elapsed_sec']}s)")
print("artifact: normalized_map.json  (input to the BCW degree-reduction pipeline)")