import json, math, random
from pathlib import Path
import numpy as np
import torch

SEED = 2270
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)

# A two-dimensional neural-ODE-like polynomial vector field.
# At p=(0,0), J=[[0,1],[0,0]], hence J is rank-one nilpotent and ker(J)=span((1,0)).
# F2=c*x^2 + k*x^3 + d*x*y + e*x^2*y + r*y^2.
# Along q0=(1,0): a=c and b=d exactly.
def field(z, c, d, k=0.37, e=-0.21, r=0.13, m=0.8):
    x, y = z[0], z[1]
    return torch.stack((y, c*x*x + k*x**3 + m*x**4 + d*x*y + e*x*x*y + r*y*y))

def jacobian_at(z, c, d, k=0.37, e=-0.21, r=0.13, m=0.8):
    z = z.detach().clone().requires_grad_(True)
    return torch.autograd.functional.jacobian(lambda zz: field(zz, c, d, k, e, r, m), z)

def monitor(c, d, h):
    p = torch.zeros(2)
    J = jacobian_at(p, c, d)
    # right singular vector for the smallest singular value
    _, _, vh = torch.linalg.svd(J)
    q = vh[-1]
    q = q / torch.linalg.vector_norm(q)
    # orient deterministically to make reported values comparable
    if q[0] < 0: q = -q
    jp = jacobian_at(p + h*q, c, d)
    jm = jacobian_at(p - h*q, c, d)
    a = -0.5 * (torch.linalg.det(jp)-torch.linalg.det(jm))/(2*h)
    b = (torch.trace(jp)-torch.trace(jm))/(2*h)
    return float(a), float(b), float(torch.linalg.det(J)), float(torch.trace(J)), q.numpy().tolist()

def math_checks():
    # Prediction 1: a scales linearly with c; prediction 2: b scales linearly with d.
    cs = np.linspace(-1.2, 1.2, 9); ds = np.linspace(-1.2, 1.2, 9)
    av = np.array([monitor(float(c), .41, 1e-3)[0] for c in cs])
    bv = np.array([monitor(.41, float(d), 1e-3)[1] for d in ds])
    slope_a, intercept_a = np.polyfit(cs, av, 1)
    slope_b, intercept_b = np.polyfit(ds, bv, 1)
    # Prediction 3: symmetric finite differences are stable as h is reduced, including at zero crossings.
    hs = [1e-1, 5e-2, 2.5e-2, 1.25e-2, 6.25e-3]
    conv = [monitor(.73, -.62, h)[:2] for h in hs]
    zero_a = monitor(0., .4, 1e-3)[0]; zero_b = monitor(.4, 0., 1e-3)[1]
    return {
        'predictions': {
            'a(c) predicted slope/intercept': [1.0, 0.0],
            'a(c) observed slope/intercept': [float(slope_a), float(intercept_a)],
            'b(d) predicted slope/intercept': [1.0, 0.0],
            'b(d) observed slope/intercept': [float(slope_b), float(intercept_b)],
            'zero crossing a at c=0 observed': zero_a,
            'zero crossing b at d=0 observed': zero_b,
            'finite_difference predicted a(h),b(h)': [[0.73 + 1.6*h*h, -0.62] for h in hs],
            'finite_difference observed h=[0.1,...]': conv,
        },
        'nilpotent_check': monitor(.73, -.62, 1e-3)[2:]
    }

def train(use_bt, steps=180, lr=.035):
    # Same tiny task: recover coefficients c=d=1 from a weak signal. BT term avoids
    # the degenerate region but does not replace the task objective.
    c = torch.tensor(0.015, requires_grad=True)
    d = torch.tensor(-0.02, requires_grad=True)
    opt = torch.optim.Adam([c, d], lr=lr)
    eps = .35; lam = 0.8
    history=[]
    for step in range(steps):
        task = (c-1.)**2 + (d-1.)**2
        a,b,_,_,_ = monitor(float(c.detach()), float(d.detach()), 1e-3)
        # Rebuild differentiable exact monitor for this controlled family: a=c,b=d.
        bt = lam*(torch.relu(torch.tensor(eps)-torch.abs(c))**2 + torch.relu(torch.tensor(eps)-torch.abs(d))**2)
        loss = task + (bt if use_bt else 0.)
        opt.zero_grad(); loss.backward(); opt.step()
        if step in (0, 9, 29, 89, 179): history.append((step, float(task), float(a), float(b)))
    return {'c':float(c), 'd':float(d), 'task_loss':float((c-1.)**2+(d-1.)**2), 'history':history}

def main():
    out = {'math': math_checks(), 'training': {'baseline':train(False), 'bt_penalty':train(True)}}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()