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

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

# Canonical BT unfolding:
# x' = y
# y' = mu*x + nu*y + alpha*x^2 + beta*x*y
# At mu=nu=0 and z*=0, J=[[0,1],[0,0]], q0=(1,0).
def field(z, p):
    x, y = z
    mu, nu, alpha, beta = p
    return torch.stack((y, mu*x + nu*y + alpha*x*x + beta*x*y))

def jacobian(z, p):
    # Explicit form keeps the z -> J -> invariant graph for higher derivatives.
    x, y = z
    mu, nu, alpha, beta = p
    return torch.stack((torch.stack((0*x, 1+0*x)),
                        torch.stack((mu + 2*alpha*x + beta*y, nu + beta*x))))

def invariants(J):
    e1 = J[0, 0] + J[1, 1]
    e2 = J[0, 0]*J[1, 1] - J[0, 1]*J[1, 0]
    return e1, e2  # e1, e2; polynomial form supports higher derivatives

def monitor(p, q_override=None, z_value=(0., 0.)):
    z = torch.tensor(z_value, dtype=p.dtype, requires_grad=True)
    J = jacobian(z, p)
    e1, e2 = invariants(J)
    if q_override is None:
        _, _, vh = torch.linalg.svd(J.detach())
        q = vh[-1].detach()
    else:
        q = torch.as_tensor(q_override, dtype=z.dtype)
    # For n=2: delta=e2/e0=e2 and tau=e1.
    gdelta = torch.autograd.grad(e2, z, retain_graph=True)[0]
    gtau = torch.autograd.grad(e1, z)[0]
    Ddelta = torch.dot(gdelta, q)
    Dtau = torch.dot(gtau, q)
    return {
        'J': J.detach().numpy().tolist(), 'q': q.numpy().tolist(),
        'delta': float(e2.detach()), 'tau': float(e1.detach()),
        'a': float((-0.5*Ddelta).detach()), 'b': float(Dtau.detach()),
        'e_transverse': 1.0,  # e_(n-2)=e0 for n=2
    }

def finite_difference(alpha, beta, h=1e-5):
    p = torch.tensor([0., 0., alpha, beta])
    def vals(x):
        r = monitor(p, [1., 0.], (x, 0.))
        return r['delta'], r['tau']
    dm = (vals(h)[0] - vals(-h)[0])/(2*h)
    dt = (vals(h)[1] - vals(-h)[1])/(2*h)
    return -0.5*dm, dt

def run_target(kind):
    torch.manual_seed(SEED)
    p = torch.nn.Parameter(torch.tensor([0.65, -0.55, 1.2, -0.8]))
    opt = torch.optim.Adam([p], lr=0.05)
    for _ in range(250):
        mu, nu = p[0], p[1]
        if kind == 'monitor':
            loss = mu*mu + nu*nu
        else:
            # Spectral-abscissa stability penalty: stable systems receive zero loss.
            disc = nu*nu + 4*mu
            rp = torch.where(disc >= 0, (nu + torch.sqrt(torch.clamp(disc, min=0)))/2, nu/2)
            loss = torch.relu(rp)**2
        opt.zero_grad(); loss.backward(); opt.step()
    r = monitor(p.detach(), [1., 0.])
    return {'final_p': p.detach().numpy().tolist(), 'delta': r['delta'], 'tau': r['tau'],
            'monitor_loss': r['delta']**2 + r['tau']**2}

def verify():
    # Prediction 1: exact unfolding identities delta=-mu, tau=nu over a grid.
    identity = []
    for mu in np.linspace(-1, 1,  nine := 9):
        for nu in np.linspace(-1, 1, nine):
            r = monitor(torch.tensor([mu, nu, 1., 1.]), [1., 0.])
            identity.append((mu, nu, r['delta'], r['tau']))
    err_delta = max(abs(d + mu) for mu, nu, d, t in identity)
    err_tau = max(abs(t - nu) for mu, nu, d, t in identity)

    # Prediction 2: a scales as alpha and b scales as beta at the BT point.
    coeff = []
    for alpha, beta in [(0.3,-0.7), (1.,2.), (-1.5,.25), (2.5,-3.)]:
        r = monitor(torch.tensor([0.,0.,alpha,beta]), [1.,0.])
        af, bf = finite_difference(alpha, beta)
        coeff.append({'alpha':alpha, 'beta':beta, 'a':r['a'], 'b':r['b'],
                      'a_fd':af, 'b_fd':bf, 'a_error':abs(r['a']-alpha),
                      'b_error':abs(r['b']-beta)})

    # Prediction 3: the candidate is simultaneous delta=tau=0 only at (mu,nu)=(0,0).
    candidates = []
    for mu in [-.2, 0., .2]:
        for nu in [-.2, 0., .2]:
            r = monitor(torch.tensor([mu,nu,1.,1.]), [1.,0.])
            candidates.append({'mu':mu, 'nu':nu, 'delta':r['delta'], 'tau':r['tau'],
                               'distance':abs(r['delta'])+abs(r['tau'])})
    out = {
        'prediction_1_unfolding': {'points':len(identity), 'max_abs_delta_minus_neg_mu':err_delta,
                                   'max_abs_tau_minus_nu':err_tau},
        'prediction_2_coefficients': coeff,
        'prediction_3_candidate_grid': candidates,
        'baseline': run_target('baseline'), 'idea_monitor': run_target('monitor')
    }
    return out

def main():
    result = verify()
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()