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

SEEDS=list(range(8)); DEVICE='cuda' if torch.cuda.is_available() else 'cpu'

def seed(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

def plant(x,u):
    # locally controlled nonlinear pendulum: [angle, angular velocity]
    th,om=x[...,0],x[...,1]
    return torch.stack((th+0.12*om, 0.985*om-0.16*torch.sin(th)+0.10*u.squeeze(-1)), -1)

def data(s,n=400,T=10,d=4):
    rng=np.random.default_rng(s)
    xs=[]; qs=[]; ys=[]
    A=np.array([[1,.12],[-.16,.985]]); B=np.array([[0],[.10]])
    K=np.array([[-1.05,-.72]])
    for _ in range(n):
        x=rng.uniform([-1.1,-1.0],[1.1,1.0]); queue=[np.zeros(1) for _ in range(d)]
        for t in range(T):
            # target is stabilizing action for the state at actuation time
            target=float(np.clip(K@x, -2,2))
            hist=np.asarray(queue,dtype=np.float32).reshape(-1)
            xs.append(x.astype(np.float32)); qs.append(hist); ys.append(target)
            applied=queue.pop(0); queue.append(np.array([target]))
            x=A@x+B[:,0]*applied[0] + rng.normal(0,.008,2)
            x[0]=((x[0]+np.pi)%(2*np.pi))-np.pi
    X=np.asarray(xs); Q=np.asarray(qs); Y=np.asarray(ys)[:,None]
    # deterministic split by generated order; independent test seed
    return X,Q,Y

class Controller(nn.Module):
    def __init__(self,d):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2,24),nn.Tanh(),nn.Linear(24,1),nn.Tanh())
    def forward(self,x): return 2*self.net(x)

def predict(x,q,d):
    # frozen local linearization, chronological queued inputs
    A=x.new_tensor([[1.,.12],[-.16,.985]])
    B=x.new_tensor([[0.],[.10]])
    p=x
    for j in range(d): p=p@A.T + q[:,j:j+1]@B.T
    return p

def fit(s,lr,idea,d=4,epochs=18):
    seed(s); X,Q,Y=data(s,400,10,d); Xt,Qt,Yt=data(s+10000,120,10,d)
    X=torch.tensor(X,device=DEVICE); Q=torch.tensor(Q,device=DEVICE); Y=torch.tensor(Y,device=DEVICE)
    Xt=torch.tensor(Xt,device=DEVICE); Qt=torch.tensor(Qt,device=DEVICE); Yt=torch.tensor(Yt,device=DEVICE)
    m=Controller(d).to(DEVICE); opt=torch.optim.Adam(m.parameters(),lr=lr)
    for _ in range(epochs):
        perm=torch.randperm(len(X),device=DEVICE)
        for ix in perm.split(128):
            inp=predict(X[ix],Q[ix],d) if idea else X[ix]
            loss=((m(inp)-Y[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        inp=predict(Xt,Qt,d) if idea else Xt
        mse=float(((m(inp)-Yt)**2).mean().cpu())
        # closed-loop rollout metric on fresh initial conditions with delayed actuation
        rng=np.random.default_rng(s+20000); vals=[]; pred_err=[]
        for z in range(30):
            x=torch.tensor(rng.uniform([-1.,-.8],[1.,.8]),dtype=torch.float32,device=DEVICE).unsqueeze(0)
            q=torch.zeros((1,d),device=DEVICE); accum=0.
            for t in range(35):
                p=predict(x,q,d); u=m(p if idea else x).clamp(-2,2)
                if idea: pred_err.append(float(torch.linalg.norm(p-x).cpu()))
                applied=q[:,0:1]; q=torch.cat((q[:,1:],u),1); x=plant(x,applied)
                accum += float((x*x).sum().cpu())
            vals.append(accum/35)
        return {'test_mse':mse,'rollout_mse':float(np.mean(vals)), 'predictor_shift':float(np.mean(pred_err)) if pred_err else 0.0}

def perm_p(a,b):
    a=np.asarray(a); b=np.asarray(b); obs=float(np.mean(b-a)); rng=np.random.default_rng(991)
    cnt=0; N=20000
    for _ in range(N):
        signs=rng.choice([-1,1],len(a)); v=float(np.mean((b-a)*signs))
        cnt += abs(v)>=abs(obs)
    return (cnt+1)/(N+1),obs

def main():
    # same union of learning rates on both sides; baseline sweep and idea sweep are identical.
    lrs=[0.001,0.003,0.009]; d=4
    base={str(lr):[fit(s,lr,False,d) for s in SEEDS] for lr in lrs}
    idea={str(lr):[fit(s,lr,True,d) for s in SEEDS] for lr in lrs}
    score=lambda r: np.mean([x['test_mse'] for x in r])
    bestb=min(lrs,key=lambda x:score(base[str(x)])); besti=min(lrs,key=lambda x:score(idea[str(x)]))
    br=base[str(bestb)]; ir=idea[str(besti)]
    p,delta=perm_p([x['test_mse'] for x in br],[x['test_mse'] for x in ir])
    # Signature measured on trained systems: prediction shift and observed one-step model mismatch.
    sig={'delay':d,'trained_models':True,'predicted_effect':'finite-horizon state differs from stale state and should reduce delayed rollout error',
         'observed_baseline_rollout_mse':float(np.mean([x['rollout_mse'] for x in br])),
         'observed_idea_rollout_mse':float(np.mean([x['rollout_mse'] for x in ir])),
         'observed_mean_predictor_shift':float(np.mean([x['predictor_shift'] for x in ir])),
         'confirmed':float(np.mean([x['rollout_mse'] for x in ir])) < float(np.mean([x['rollout_mse'] for x in br]))}
    out={'track':'dynamics_fallback','device':DEVICE,'baseline_sweep':{str(k):{'mean_test_mse':score(v),'per_seed':v} for k,v in base.items()},'idea_sweep':{str(k):{'mean_test_mse':score(v),'per_seed':v} for k,v in idea.items()},'best_baseline_lr':bestb,'best_idea_lr':besti,'paired_delta_mean':delta,'permutation_p':p,'mechanism_signature':sig,'custom_track':{'name':'delayed_pendulum_dynamics','file':'stage2_bench.py','domain':'dynamics'}}
    Path('bench_report.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()