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

SEED=7
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DT=0.12
P=np.diag([1.0, 0.35]).astype(np.float32)
C2=0.025

# True plant has an unmodeled bounded force; nominal model omits it.
def true_step(x,u, amp=0.22):
    x=np.asarray(x); u=np.asarray(u)
    return np.stack([x[...,0]+DT*x[...,1], x[...,1]+DT*(u+0.18*np.sin(1.7*x[...,0])+amp*np.sin(2.3*x[...,0]))],-1)
def nominal_step(x,u):
    x=np.asarray(x); u=np.asarray(u)
    return np.stack([x[...,0]+DT*x[...,1], x[...,1]+DT*(u+0.18*np.sin(1.7*x[...,0]))],-1)
def V(x): return 0.5*np.einsum('...i,ij,...j->...',x,P,x)
def gradV(x): return np.asarray(x)@P.T

def analytic_lipschitz(z,delta):
    # Valid on the uncertainty ball for quadratic V: ||Pz|| + ||P||_2 delta.
    return np.linalg.norm(gradV(z),axis=-1)+np.max(np.linalg.eigvalsh(P))*delta

def math_check():
    rng=np.random.default_rng(3); rows=[]
    for d in [0.,.03,.08,.15,.25]:
        z=rng.normal(size=(4000,2)); e=rng.normal(size=(4000,2)); e=e/np.linalg.norm(e,axis=1,keepdims=True)*d*rng.random((4000,1))
        lhs=np.abs(V(z+e)-V(z)); rhs=analytic_lipschitz(z,d)*np.linalg.norm(e,axis=1)
        rows.append((d,float(np.max(lhs-rhs)),float(np.mean(lhs/rhs.clip(1e-9)))))
    return rows

def policy_torch(K,x):
    return -(x@K.T).clamp(-3,3)
def train(robust, delta, n=1000, steps=180):
    rng=np.random.default_rng(SEED+n+int(delta*1000)+(1 if robust else 0))
    x=rng.uniform(-2,2,size=(n,2)).astype(np.float32)
    # held-out-like nominal transition samples, using the nominal model as learned dynamics
    xt=torch.tensor(x); K=torch.tensor([[1.0,1.0]],requires_grad=True)
    opt=torch.optim.Adam([K],lr=.045)
    for _ in range(steps):
        u=policy_torch(K,xt)
        # differentiable nominal transition for policy gradients
        z= torch.stack([xt[:,0]+DT*xt[:,1], xt[:,1]+DT*(u[:,0]+.18*torch.sin(1.7*xt[:,0]))],1)
        dv=0.5*torch.sum((z@torch.tensor(P))*z,1)-0.5*torch.sum((xt@torch.tensor(P))*xt,1)
        lv=torch.sqrt(torch.sum((z@torch.tensor(P))**2,1)+1e-8)+delta
        target=C2*torch.sum(xt*xt,1)
        margin=dv+(lv*delta if robust else 0)+target
        loss=torch.mean(u[:,0]**2)*.012 + torch.mean(torch.nn.functional.softplus(8*margin))/8 + .0005*torch.sum(K*K)
        opt.zero_grad(); loss.backward(); opt.step()
    return K.detach().numpy()

def evaluate(K, delta_eval, amp=.22):
    rng=np.random.default_rng(900+int(delta_eval*1000)); x=rng.uniform(-2,2,size=(1800,2)).astype(np.float32)
    u=-(x@K.T).clip(-3,3); z=nominal_step(x,u[:,0]); dv=V(z)-V(x); lv=analytic_lipschitz(z,delta_eval)
    robust_margin=dv+lv*delta_eval+C2*np.sum(x*x,1)
    nominal_margin=dv+C2*np.sum(x*x,1)
    # random bounded errors test the actual next state and certificate
    e=true_step(x,u[:,0],amp=amp)-z
    actual_margin=V(z+e)-V(x)+C2*np.sum(x*x,1)
    conv=0
    for j in range(400):
        xx=rng.uniform(-2,2,size=2)
        for t in range(60):
            uu=float(np.clip(-(K@xx)[0],-3,3)); xx=true_step(xx,uu,amp=amp)
            if np.linalg.norm(xx)<.12: conv+=1; break
    return dict(min_robust=float(np.min(-robust_margin)), median_robust=float(np.median(-robust_margin)), min_nominal=float(np.min(-nominal_margin)), actual_violation=float(np.mean(actual_margin>0)), convergence=conv/1500, K=K.tolist())

def mechanism_sweep():
    # Prediction 1: for a fixed policy, m(delta)=m(0)-L*delta-(lambda_max P)delta^2;
    # observed fit should have positive linear coefficient L and small quadratic term.
    K=np.array([[1.25,1.55]],dtype=np.float32); rng=np.random.default_rng(11); x=rng.uniform(-1.8,1.8,(12000,2)); u=-(x@K.T).clip(-3,3)[:,0]; z=nominal_step(x,u)
    base=V(z)-V(x)+C2*np.sum(x*x,1); L=np.linalg.norm(gradV(z),axis=1); lam=np.max(np.linalg.eigvalsh(P))
    ds=np.array([0,.02,.05,.10,.18,.28]); observed=[]; predicted=[]
    for d in ds:
        observed.append(float(np.min(-(base+(L+d)*d))))
        predicted.append(float(np.min(-base-L*d-lam*d*d)))
    # Prediction 2: actual bounded error is certified whenever predicted robust margin >= 0.
    # Sweep scalar c2 to locate the sign transition and compare actual violations.
    c2s=np.array([0,.01,.02,.03,.05,.08]); boundary=[]; actual=[]
    e=true_step(x,u,amp=.22)-z; ae=V(z+e)-V(x)
    for c in c2s:
        mm=ae+c*np.sum(x*x,1); boundary.append(float(np.min(-mm))); actual.append(float(np.mean(mm>0)))
    return {'delta_sweep':{'delta':ds.tolist(),'observed_min_certificate':observed,'predicted_min_certificate':predicted},'c2_sweep':{'c2':c2s.tolist(),'observed_min_actual_certificate':boundary,'actual_violation_fraction':actual}}

def main():
    check=math_check(); mech=mechanism_sweep()
    # Same data-size proxy: uncertainty radius grows with scarcity, calibrated conservatively.
    rounds=[]
    for frac,delta in [(1.0,.055),(.5,.11),(.2,.18)]:
        kb=train(False,delta,n=int(500*frac)); kr=train(True,delta,n=int(500*frac))
        rounds.append({'fraction':frac,'delta':delta,'baseline':evaluate(kb,delta),'robust':evaluate(kr,delta)})
    out={'math_check':check,'mechanism_predictions':mech,'training_comparison':rounds}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()