import json, math, random
import numpy as np
import torch
from torch import nn
from symplectic_block import SymplecticBlock

SEED=646
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def gradV(q): return q+0.4*q**3
def H(q,p): return .5*(q*q+p*p)+.1*q**4

def leap(q,p,h):
    p=p-.5*h*gradV(q); q=q+h*p; p=p-.5*h*gradV(q); return q,p
def euler(q,p,h): return q+h*p,p-h*gradV(q)
def metrics(fn,q0=1.,p0=.3,h=.1,steps=10000):
    q,p=q0,p0; e0=float(H(q,p)); mx=me=0.; div=None; samples=[]
    for k in range(steps):
        try: q,p=fn(q,p,h); norm=math.hypot(float(q),float(p)); en=float(H(q,p))
        except (OverflowError,ValueError,ZeroDivisionError): div=k+1; norm=en=float('inf'); break
        if not math.isfinite(norm) or not math.isfinite(en): div=k+1; break
        mx=max(mx,norm); me=max(me,abs(en-e0))
        if k in (9,99,999,4999,9999): samples.append([k+1,norm,en-e0])
    return {'final_norm':norm,'max_norm':mx,'max_abs_energy_error':me,'diverged_at_step':div,'samples':samples}

def learned_regression():
    torch.manual_seed(SEED); n=256; x=torch.randn(n,1); q=torch.randn(n,2); p=torch.zeros_like(q)
    target_q=q+0.1*p; target_p=p-0.1*(q+0.4*q**3)
    sym=SymplecticBlock(2,1,substeps=2,hidden=16,step=.1)
    opt=torch.optim.Adam(sym.parameters(),lr=3e-3)
    for _ in range(80):
        opt.zero_grad(); qo,po=sym(q,p,x); loss=((qo-target_q)**2+(po-target_p)**2).mean(); loss.backward(); opt.step()
    qo,po=sym(q,p,x); symloss=float(((qo-target_q)**2+(po-target_p)**2).mean().detach())
    base=nn.Sequential(nn.Linear(5,16),nn.Tanh(),nn.Linear(16,4)); opt=torch.optim.Adam(base.parameters(),lr=3e-3)
    inp=torch.cat([q,p,x.expand(n,1)],1); target=torch.cat([target_q,target_p],1)
    for _ in range(80):
        opt.zero_grad(); loss=((base(inp)-target)**2).mean(); loss.backward(); opt.step()
    baseloss=float(((base(inp)-target)**2).mean().detach())
    return {'one_step_mse':{'euler_analogue':baseloss,'learned_symplectic':symloss},'parameters':{'euler':sum(p.numel() for p in base.parameters()),'symplectic':sum(p.numel() for p in sym.parameters())}}

def main():
    z=torch.tensor([.7,-.4],dtype=torch.float64,requires_grad=True)
    def lf(v): return torch.stack(leap(v[0],v[1],.2))
    def eu(v): return torch.stack(euler(v[0],v[1],.2))
    dl=float(torch.linalg.det(torch.autograd.functional.jacobian(lf,z))); de=float(torch.linalg.det(torch.autograd.functional.jacobian(eu,z)))
    qr,pr=1.,.3
    for _ in range(100000): qr,pr=leap(qr,pr,1e-3)
    errors={}
    for name,fn in [('leapfrog',leap),('euler',euler)]:
        try:
            q,p=1.,.3
            for _ in range(1000): q,p=fn(q,p,.1)
            errors[name]=float(math.hypot(q-qr,p-pr))
        except (OverflowError,ValueError): errors[name]=None
    out={'seed':SEED,'jacobian_determinant':{'leapfrog':dl,'euler':de},'rollouts':{str(h):{'leapfrog':metrics(leap,h=h),'euler':metrics(euler,h=h)} for h in (.1,.2)},'T100_error':errors,'learned_regression':learned_regression()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()