import json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

SEED = 3038
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(min(12, torch.get_num_threads()))
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if DEVICE == "cuda":
        torch.cuda.set_device(0)
        torch.zeros(1, device="cuda")
except Exception:
    DEVICE = "cpu"

# Discrete Duffing oscillator. The target is position, while the learned state
# contains position and velocity, so long-horizon state propagation is nontrivial.
def duffing(x, u=0.0, dt=0.08):
    q, v = x[..., 0], x[..., 1]
    a = -0.8*q - 0.25*v - 0.15*q**3 + u
    return torch.stack((q + dt*v, v + dt*a), -1)

def make_data(n, T=140, noise=0.015):
    xs = torch.zeros(n, T+1, 2)
    xs[:, 0] = torch.randn(n, 2) * torch.tensor([0.45, 0.20])
    for k in range(T):
        # Mild known excitation makes the identification problem better conditioned.
        u = 0.12 * torch.sin(torch.tensor(0.11*k))
        xs[:, k+1] = duffing(xs[:, k], u) + noise*torch.randn(n, 2)
    return xs

class Dynamics(nn.Module):
    def __init__(self, hidden=32):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 2))
    def forward(self, x):
        return x + 0.08*torch.tanh(self.net(x))

class Storage(nn.Module):
    def __init__(self, hidden=24, delta=0.01):
        super().__init__(); self.delta = delta
        self.g = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1))
    def forward(self, x):
        return self.g(x).square().squeeze(-1) + self.delta*x.square().sum(-1)

def rollout(model, x0, H):
    out=[x0]; x=x0
    for _ in range(H):
        x=model(x); out.append(x)
    return torch.stack(out, 1)

def train(train_x, regularized, steps=260, horizon=15):
    model=Dynamics().to(DEVICE)
    storage=Storage().to(DEVICE) if regularized else None
    params=list(model.parameters()) + ([] if storage is None else list(storage.parameters()))
    opt=torch.optim.Adam(params, lr=2e-3)
    ema=None; rho=0.08; lam=0.35
    n=train_x.shape[0]; rng=np.random.default_rng(SEED + (1 if regularized else 0))
    for step in range(steps):
        idx=torch.tensor(rng.integers(0,n,32), device=DEVICE)
        k=torch.tensor(rng.integers(0, train_x.shape[1]-horizon, 32), device=DEVICE)
        # Gather independent starting points and targets for a short rollout.
        x0=train_x[idx,k].to(DEVICE)
        target=torch.stack([train_x[idx,k+j+1,0].to(DEVICE) for j in range(horizon)], 1)
        pred=rollout(model,x0,horizon)[:,1:]
        losses=(pred[:,:,0]-target).square()
        ordinary=losses.mean()
        if ema is None: ema=ordinary.detach()
        else: ema=0.995*ema + 0.005*ordinary.detach()
        loss=ordinary
        if regularized:
            V=torch.stack([storage(pred[:,j]) for j in range(horizon+0)],1)
            V0=storage(x0)
            Vprev=torch.cat([V0[:,None], V[:,:-1]],1)
            residual=V-Vprev + rho*losses - ordinary.detach() + ema.detach()
            diss=F.relu(residual).square().mean()
            loss=ordinary+lam*diss
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(params, 5.0); opt.step()
    return model, storage

@torch.no_grad()
def evaluate(model, storage, test_x, H=100):
    x0=test_x[:,0].to(DEVICE)
    pred=rollout(model,x0,H)[:,1:]
    target=test_x[:,1:H+1].to(DEVICE)
    err=(pred[:,:,0]-target[:,:,0]).square()
    # Evaluate residual against the measured per-step loss and a fixed reference
    # l*=0. This is the direct strict-dissipativity residual for this experiment.
    if storage is not None:
        vv=torch.cat([storage(x0)[:,None], storage(pred)],1)
        residual=vv[:,1:]-vv[:,:-1]+0.08*err
        pos=F.relu(residual)
        frac=(residual>0).float().mean().item()
        mean_res=residual.mean().item(); cum_res=residual.sum(1).mean().item()
    else: frac=mean_res=cum_res=float('nan')
    return dict(mse10=err[:,:10].mean().item(), mse100=err.mean().item(), cum100=err.sum(1).mean().item(), positive_fraction=frac, mean_residual=mean_res, cumulative_residual=cum_res)

def math_check():
    # Direct numerical verification: sum of V_{k+1}-V_k telescopes exactly;
    # when residual <= 0, the proposed cumulative bound follows.
    rng=np.random.default_rng(SEED)
    B,N=16,40
    V=np.cumsum(np.abs(rng.normal(size=(B,N+1))),axis=1)
    ell=np.abs(rng.normal(size=(B,N)))+0.2
    lstar=np.zeros_like(ell); d2=np.abs(rng.normal(size=(B,N)))
    rho=0.1
    r=V[:,1:]-V[:,:-1]+rho*d2-ell+lstar
    identity=np.max(np.abs((V[:,1:]-V[:,:-1]).sum(1)-(V[:,-1]-V[:,0])))
    # Construct certified samples by setting ell equal to the RHS plus margin.
    ell_cert=V[:,1:]-V[:,:-1]+rho*d2+0.5
    lhs=(rho*d2).sum(1)
    rhs=(ell_cert).sum(1)+V[:,0]-V[:,-1]
    bound_violation=np.max(lhs-rhs)
    return dict(max_telescoping_error=float(identity), max_certified_bound_violation=float(bound_violation), random_mean_positive_residual=float(np.maximum(r,0).mean()))

def main():
    train_x=make_data(96); test_x=make_data(24)
    check=math_check()
    train_x=train_x.to(DEVICE); test_x=test_x.to(DEVICE)
    base,_=train(train_x,False); diss,store=train(train_x,True)
    result={'device':DEVICE,'seed':SEED,'math_check':check,
            'baseline':evaluate(base, None, test_x), 'dissipative':evaluate(diss,store,test_x)}
    print(json.dumps(result, indent=2))
    Path('results.json').write_text(json.dumps(result, indent=2))

if __name__=='__main__': main()