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

SEED=7
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
except Exception:
    device=torch.device('cpu')

def pend_step(x, dt=.08):
    # normalized damped pendulum, x=[angle, angular velocity]
    th, om=x[...,0], x[...,1]
    return np.stack([th+dt*om, om+dt*(-np.sin(th)-.12*om)], axis=-1)

def make_data(ntraj=180, length=35):
    xs=[]
    for _ in range(ntraj):
        x=np.array([np.random.uniform(-2.5,2.5), np.random.uniform(-2.0,2.0)])
        for _ in range(length):
            y=pend_step(x); xs.append((x.copy(),y.copy())); x=y
    return np.asarray([a for a,b in xs],np.float32), np.asarray([b for a,b in xs],np.float32)

class MLP(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,2))
    def forward(self,x): return self.net(x)

def cap_spectral(A,r):
    rho=max(abs(np.linalg.eigvals(A)))
    return A*(r/rho) if rho>r else A.copy(), float(rho)

def prediction_mats(A,B,N):
    n,m=B.shape; AA=np.zeros((N*n,n)); BB=np.zeros((N*n,N*m))
    # Stack Z=[z_0,...,z_{N-1}].  Then z_k=A^k z_0+
    # sum_{j<k} A^(k-1-j) B v_j.
    p=np.eye(n)
    for k in range(N):
        AA[k*n:(k+1)*n]=p
        for j in range(k):
            BB[k*n:(k+1)*n,j*m:(j+1)*m]=(np.linalg.matrix_power(A,k-1-j)@B)
        p=p@A
    return AA,BB

def mpc_correction(z, neural_next, A, N=8, vmax=.20, delta=1.0):
    # B=I; Q tracks the stable Koopman prediction, R regularizes correction.
    n=2; B=np.eye(n); AA,BB=prediction_mats(A,B,N)
    target=AA@z
    # The neural proposal is the initial state of the controlled rollout;
    # optimize corrections toward the adapted Koopman reference trajectory.
    # This avoids the degenerate zero correction obtained by targeting AA z
    # from the same AA z initial condition.
    base=np.tile(neural_next, N)
    H=BB.T@BB+.08*np.eye(N*n); g=BB.T@(base-target)
    try: V=np.linalg.solve(H,-g)
    except np.linalg.LinAlgError: V=np.zeros(N*n)
    v=np.clip(V[:n],-vmax,vmax)
    raw=neural_next+v
    # trust region around current observed latent state; radial projection is the
    # small-QP feasibility projection for this two-dimensional toy.
    d=raw-z; norm=np.linalg.norm(d)
    if norm>delta: raw=z+d*(delta/norm)
    return raw, float(np.linalg.norm(v))

def main():
    X,Y=make_data(); tx=torch.tensor(X,device=device); ty=torch.tensor(Y,device=device)
    model=MLP().to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    for epoch in range(260):
        p=model(tx); loss=((p-ty)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad(): train_mse=float(((model(tx)-ty)**2).mean().cpu())
    # least-squares adapted linear latent dynamics
    A=np.linalg.lstsq(X,Y,rcond=None)[0].T
    A_cap,rho=cap_spectral(A,.95)
    # verify condensed formula against iterative dynamics
    B=np.eye(2); z=np.array([.3,-.4]); V=np.random.randn(6*2)*.03
    AA,BB=prediction_mats(A_cap,B,6); zs=[]; q=z.copy()
    # The standard matrices stack z_0,...,z_{N-1}; the first row is z_0.
    for k in range(6):
        zs.append(q.copy())
        q=A_cap@q+B@V[k*2:(k+1)*2]
    mat_err=float(np.max(np.abs(np.asarray(zs).reshape(-1)-(AA@z+BB@V))))
    # spectral growth check, same initial norm and 100 steps
    z0=np.array([1.,1.]); u=z0.copy(); c=z0.copy()
    for _ in range(100): u=A@u; c=A_cap@c
    growth_un=float(np.linalg.norm(u)/np.linalg.norm(z0)); growth_cap=float(np.linalg.norm(c)/np.linalg.norm(z0))
    # long rollout errors from held-out initial conditions
    starts=[]
    for _ in range(24): starts.append(np.array([np.random.uniform(-2.5,2.5),np.random.uniform(-2,2)],np.float32))
    horizons=60; errs_base=[]; errs_mpc=[]; norms_base=[]; norms_mpc=[]
    for s in starts:
        truth=s.copy(); zb=s.copy(); zm=s.copy()
        for _ in range(horizons):
            truth=pend_step(truth)
            with torch.no_grad(): nnnext=model(torch.tensor(zb,device=device,dtype=torch.float32)).cpu().numpy()
            zb=nnnext
            with torch.no_grad(): nnnext2=model(torch.tensor(zm,device=device,dtype=torch.float32)).cpu().numpy()
            zm,_=mpc_correction(zm,nnnext2,A_cap,N=8,vmax=.20,delta=.55)
            errs_base.append(np.linalg.norm(zb-truth)); errs_mpc.append(np.linalg.norm(zm-truth))
            norms_base.append(np.linalg.norm(zb)); norms_mpc.append(np.linalg.norm(zm))
    # trust-radius sweep exposes bias vs extrapolation in this implementation
    sweep=[]
    for delta in [.15,.3,.55,1.0,2.0]:
        ee=[]
        for s in starts[:12]:
            truth=s.copy(); z=s.copy()
            for _ in range(40):
                truth=pend_step(truth)
                with torch.no_grad(): nnn=model(torch.tensor(z,device=device,dtype=torch.float32)).cpu().numpy()
                z,_=mpc_correction(z,nnn,A_cap,delta=delta)
                ee.append(np.linalg.norm(z-truth))
        sweep.append([delta,float(np.mean(ee))])
    out={'device':str(device),'train_one_step_mse':train_mse,'rho_unconstrained':rho,
         'condensed_matrix_max_error':mat_err,'growth_100_unconstrained':growth_un,
         'growth_100_capped':growth_cap,'baseline_mean_60step_error':float(np.mean(errs_base)),
         'idea_mean_60step_error':float(np.mean(errs_mpc)),'baseline_mean_pred_norm':float(np.mean(norms_base)),
         'idea_mean_pred_norm':float(np.mean(norms_mpc)),'trust_radius_sweep':sweep,
         'config':{'n_train_pairs':len(X),'n_test':len(starts),'horizon':60,'spectral_cap':.95,'vmax':.20}}
    open('results.json','w').write(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()