import json
from pathlib import Path
import numpy as np
from scipy.linalg import solve_discrete_are

SEED=7
rng=np.random.default_rng(SEED)
A=np.array([[1.,.2],[0.,1.]])
B=np.array([[.02],[.2]])
XMAX=np.array([2.,2.]); UMAX=1.
F=np.vstack([np.eye(2),-np.eye(2)]); fx=np.tile(XMAX,2)
Q=np.diag([3.,1.]); R=np.array([[.15]])
P=solve_discrete_are(A,B,Q,R)
K=-(np.linalg.inv(R+B.T@P@B)@(B.T@P@A)).reshape(1,2)
rho=max(abs(np.linalg.eigvals(A+B@K)))
Acl=A+B@K

def mpc(x): return float(np.clip((K@x)[0],-UMAX,UMAX))
def support(g): return g*np.sum(np.abs(F),axis=1)
def robust_ok(x,u,g):
    return bool(np.all(F@(A@x+B[:,0]*u)+support(g)<=fx+1e-10) and abs(u)<=UMAX+1e-10)
def shield(x,unn,g):
    um=mpc(x)
    if robust_ok(x,unn,g): return unn,0.,True
    # choose smallest tested correction to the explicit fallback
    for lam in np.linspace(0,1,41):
        u=(1-lam)*unn+lam*um
        if robust_ok(x,u,g): return u,lam,True
    return um,1.,robust_ok(x,um,g)
def rollout(policy,g,episodes=60,T=50,scale=1.):
    violations=0; fallbacks=0; steps=0; terminal=0
    for _ in range(episodes):
        x=rng.uniform(-1.2,1.2,2); reached=False
        for t in range(T):
            unn=float(np.clip(scale*mpc(x)+rng.normal(0,.28),-1.5,1.5))
            if policy=='shield': u,lam,ok=shield(x,unn,g); fallbacks+=lam>0
            else: u=unn
            w=rng.uniform(-g,g,2); xn=A@x+B[:,0]*u+w
            violations += np.any(np.abs(xn)>XMAX+1e-9); steps+=1
            x=xn
            if np.linalg.norm(x)<.25: reached=True
        terminal+=reached
    return {'violation_rate':violations/steps,'fallback_rate':fallbacks/steps,'terminal_fraction':terminal/episodes}
def verify_support():
    errs=[]
    for g in [.01,.1,.3]:
        for _ in range(3000):
            v=rng.uniform(-g,g,2)
            errs.append(np.max(np.abs(F@v)-support(g))) # <=0, equality occurs at vertices
    # independently maximize each face over box vertices
    max_err=0.
    for g in [.01,.1,.3]:
        verts=np.array([[a,b] for a in [-g,g] for b in [-g,g]])
        max_err=max(max_err,float(np.max(np.abs(np.max(verts@F.T,axis=0)-support(g)))))
    return max_err

def main():
    support_err=verify_support()
    # boundary prediction: robust state constraint at x=0 has gamma <= 2.
    # Exact disturbance-only predicted boundary is gamma*=2; test direct feasibility.
    boundary=[]
    for g in np.linspace(0,2.4,25): boundary.append((float(g),robust_ok(np.zeros(2),0.,g)))
    feasible=[g for g,ok in boundary if ok]
    observed=max(feasible)
    rows=[]
    for g in [0.,.1,.3,.6,1.0]:
        s=rollout('shield',g,scale=1.)
        u=rollout('raw',g,scale=1.)
        rows.append({'gamma':g,'shield':s,'unshielded':u})
    gain=[]
    for scale in [.5,1.,1.5,2.,3.]:
        r=rollout('shield',.1,scale=scale)
        gain.append({'policy_scale':scale,'fallback_rate':r['fallback_rate'],'violation_rate':r['violation_rate']})
    # contraction prediction for fallback/no disturbance: ||Acl^t x|| <= C rho^t ||x||.
    x=np.array([1.,1.]); norms=[]
    for _ in range(15): norms.append(float(np.linalg.norm(x))); x=Acl@x
    ratios=[norms[i+1]/norms[i] for i in range(len(norms)-1)]
    out={'seed':SEED,'K':K.tolist(),'closed_loop_spectral_radius':float(rho),
         'support_max_abs_error':support_err,'predicted_gamma_boundary':2.,'observed_grid_boundary':observed,
         'safety_sweep':rows,'fallback_sweep':gain,'contraction_norms':norms,'contraction_ratios':ratios}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()