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


def seed_all(s=7):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)


def get_device():
    if torch.cuda.is_available():
        try:
            torch.zeros(1, device='cuda')
            return torch.device('cuda')
        except Exception:
            pass
    return torch.device('cpu')


def leapfrog_matrix(h):
    # H=(q^2+p^2)/2; the displayed split leapfrog is a linear map.
    return np.array([[1-h*h/2, h], [-h*(1-h*h/4), 1-h*h/2]], dtype=float)


def symplectic_check():
    J=np.array([[0.,1.],[-1.,0.]])
    vals=[]
    for h in (0.05, 0.2, 0.8):
        A=leapfrog_matrix(h)
        vals.append({'h':h, 'det':float(np.linalg.det(A)),
                     'symplectic_error':float(np.linalg.norm(A.T@J@A-J)),
                     'energy_after_1000':float(np.linalg.norm(np.linalg.matrix_power(A,1000)@np.array([1.,0.]))**2/2)})
    # Euler control, which does not preserve the form.
    E=np.array([[1.,.2],[-.2,1.]])
    vals.append({'euler_h':.2, 'det':float(np.linalg.det(E)),
                 'symplectic_error':float(np.linalg.norm(E.T@J@E-J)),
                 'energy_after_1000':float(np.linalg.norm(np.linalg.matrix_power(E,1000)@np.array([1.,0.]))**2/2)})
    return vals


class HamNet(nn.Module):
    def __init__(self, width=32):
        super().__init__()
        self.net=nn.Sequential(nn.Linear(2,width), nn.Tanh(), nn.Linear(width,1))
    def forward(self,z): return self.net(z).squeeze(-1)


def ham_step(model,z,h,create_graph=True):
    # Two fixed-point iterations implement the implicit p half-step.
    q,p=z[:,0],z[:,1]
    ph=p
    for _ in range(2):
        zz=torch.stack((q,ph),1).requires_grad_(True)
        gq=torch.autograd.grad(model(zz).sum(),zz,create_graph=create_graph)[0][:,0]
        ph=p-.5*h*gq
    zz=torch.stack((q,ph),1).requires_grad_(True)
    gp=torch.autograd.grad(model(zz).sum(),zz,create_graph=create_graph)[0][:,1]
    qn=q+h*gp
    zz2=torch.stack((qn,ph),1).requires_grad_(True)
    gq2=torch.autograd.grad(model(zz2).sum(),zz2,create_graph=create_graph)[0][:,0]
    pn=ph-.5*h*gq2
    return torch.stack((qn,pn),1)


class Residual(nn.Module):
    def __init__(self, width=32, h=.15):
        super().__init__(); self.h=h
        self.net=nn.Sequential(nn.Linear(2,width),nn.Tanh(),nn.Linear(width,2))
    def forward(self,z): return z+self.h*self.net(z)


def make_data(n=256, length=31, h=.15, device='cpu'):
    # Exact oscillator rotations; each sequence has a random phase/amplitude.
    x=[]
    for _ in range(n):
        z=torch.randn(2)*.7
        seq=[z]
        for _ in range(length-1):
            q,p=seq[-1]; c=math.cos(h); s=math.sin(h)
            seq.append(torch.stack((c*q+s*p, -s*q+c*p)))
        x.append(torch.stack(seq))
    return torch.stack(x).to(device)


def train_model(model, data, steps=500, h=.15, device='cpu'):
    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    model.train()
    for it in range(steps):
        idx=torch.randint(0,data.shape[0],(64,),device=device)
        z=data[idx,0]; target=data[idx,1]
        pred=model(z)
        loss=((pred-target)**2).mean()
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10.); opt.step()
    return float(loss.detach().cpu())


def rollout(model, initial, horizon=200):
    z=initial.clone(); out=[z]
    # Hamiltonian evaluation needs autograd for input gradients, but no
    # parameter graph is needed.
    for _ in range(horizon):
        # enable_grad is required by Hamiltonian input derivatives; detaching
        # prevents graph growth and is harmless for rollout evaluation.
        with torch.enable_grad():
            z=model(z)
        z=z.detach()
        out.append(z)
    return torch.stack(out,1)


def evaluate(model, test, horizons=(10,50,200)):
    pred=rollout(model,test[:,0],max(horizons))
    result={}
    for k in horizons:
        result[str(k)]=float(((pred[:,k]-test[:,k])**2).mean().sqrt().cpu())
    energy=(pred.pow(2).sum(-1)/2)
    result['energy_abs_drift_200']=float((energy[:,200]-energy[:,0]).abs().mean().cpu())
    return result


def main():
    seed_all(7); dev=get_device(); h=.15
    check=symplectic_check()
    train=make_data(256,31,h,dev); test=make_data(128,201,h,dev)
    # Same architecture width, optimizer, steps and data for the learned transition.
    seed_all(11); base=Residual(32,h); train_model(base,train,500,h,dev); b=evaluate(base,test)
    seed_all(11); idea=HamNet(32); train_model(lambda z: ham_step(idea,z,h), train, 0, h, dev) if False else None
    idea.to(dev); opt=torch.optim.Adam(idea.parameters(),lr=3e-3); idea.train()
    for _ in range(500):
        idx=torch.randint(0,train.shape[0],(64,),device=dev); z=train[idx,0]; target=train[idx,1]
        pred=ham_step(idea,z,h); loss=((pred-target)**2).mean(); opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(idea.parameters(),10.); opt.step()
    idea.eval(); a=evaluate(lambda z: ham_step(idea,z,h),test)
    print(json.dumps({'device':str(dev),'symplectic_check':check,'baseline':b,'idea':a},indent=2))

if __name__=='__main__': main()