import json, random
from pathlib import Path
import numpy as np

# Scalar robust control-affine plant: xdot = u + w, |w| <= wbar.
# V=x^2/2 and h=1-x^2. The QP uses nonnegative CLF/CBF slacks.

def quantities(x, wbar, cV=0.4, ch=1.0):
    V = .5*x*x
    # CLF: gradV*f + margin + cV*V <= sV
    aV, bV = x, wbar*abs(x) + cV*V
    # CBF: gradH*f - margin + alpha(h) >= -sH.
    # Multiplying by -1 gives aH*u+bH <= sH.
    aH, bH = 2*x, 2*wbar*abs(x) - ch*(1-x*x)
    return aV, bV, aH, bH, V

def shield(x, unet, wbar, umax=1.0, rho=100.0, cV=.4, ch=1.0):
    a,b,c,d,V = quantities(x,wbar,cV,ch)
    cuts=[-umax, umax]
    for q,r in ((a,b),(c,d)):
        if abs(q)>1e-12:
            z=-r/q
            if -umax <= z <= umax: cuts.append(z)
    cuts=sorted(set(cuts)); candidates=list(cuts)
    for lo,hi in zip(cuts[:-1],cuts[1:]):
        mid=(lo+hi)/2
        active1=(a*mid+b)>0; active2=(c*mid+d)>0
        den=1.; num=unet
        if active1: den += 2*rho*a*a; num -= 2*rho*a*b
        if active2: den += 2*rho*c*c; num -= 2*rho*c*d
        candidates.append(min(hi,max(lo,num/den)))
    def obj(u): return .5*(u-unet)**2 + rho*max(0,a*u+b)**2 + rho*max(0,c*u+d)**2
    u=min(candidates,key=obj)
    return float(u), float(max(0,a*u+b)), float(max(0,c*u+d)), float(obj(u))

def check_scaling():
    x=.6; cV=.4; rows=[]
    for wb in np.linspace(0,.8,9):
        pred=wb*abs(x)+cV*.5*x*x
        # Direct numerical robust CLF residual at u=0.
        a,b,_,_,_=quantities(x,float(wb),cV)
        observed=a*0+b
        rows.append({'wbar':float(wb),'predicted_margin':float(pred),'observed_margin':float(observed)})
    slope=float(np.polyfit([r['wbar'] for r in rows],[r['observed_margin'] for r in rows],1)[0])
    return rows,slope

def check_boundary():
    # At x=-1, CLF requires u >= wbar+cV/2; CBF requires u >= wbar.
    # Thus zero-slack feasibility ends at wbar=umax-cV/2=.8.
    umax=1.; cV=.4; predicted=umax-cV/2; rows=[]
    for wb in np.linspace(.4,1.2,17):
        u,sv,sh,_=shield(-1.,0.,float(wb),umax=umax,rho=1e6,cV=cV)
        rows.append({'wbar':float(wb),'u':u,'total_slack':float(sv+sh),
                     'zero_slack':bool(sv+sh<1e-5)})
    onset=next((r['wbar'] for r in rows if not r['zero_slack']),None)
    return predicted,rows,onset

def envelope():
    # For feasible robust CLF, adversarial V derivative satisfies Vdot <= -cV V.
    x=.8; wb=.2; dt=.002; T=2.; cV=.4; vs=[]; sl=[]; residual=[]
    for _ in range(int(T/dt)):
        u,s,h,_=shield(x,1.,wb,umax=1.,rho=1e6,cV=cV)
        w=wb*(1 if x>=0 else -1)
        residual.append(x*(u+w)+cV*.5*x*x)
        x += dt*(u+w); vs.append(.5*x*x); sl.append(s)
    t=np.arange(1,len(vs)+1)*dt; v0=.5*.8*.8
    bound=v0*np.exp(-cV*t)
    return {'max_V_over_zero_disturbance_bound':float(max(np.array(vs)/(bound+1e-12))),
            'max_robust_clf_residual':float(max(residual)), 'max_clf_slack':float(max(sl)),
            'final_V':float(vs[-1]), 'initial_V':v0, 'target_rate':cV}

def baseline_episode(wb=.2,n=1000,dt=.002):
    x=.8; violations=0
    for _ in range(n):
        x += dt*(1.+wb*(1 if x>=0 else -1)); violations += abs(x)>1
    return {'violation_fraction':violations/n,'final_abs_x':abs(x)}

def shield_episode(wb=.2,n=1000,dt=.002):
    x=.8; violations=0; dev=[]; sl=[]
    for _ in range(n):
        u,s,h,_=shield(x,1.,wb,rho=1e6)
        x += dt*(u+wb*(1 if x>=0 else -1)); violations += abs(x)>1
        dev.append(abs(u-1.)); sl.append(s+h)
    return {'violation_fraction':violations/n,'final_abs_x':abs(x),
            'mean_action_deviation':float(np.mean(dev)),'mean_slack':float(np.mean(sl))}

def main():
    np.random.seed(7); random.seed(7)
    scale,slope=check_scaling(); pred,rows,onset=check_boundary()
    out={'scaling':{'rows':scale,'observed_slope':slope,'predicted_slope':.6},
         'boundary':{'predicted_wbar_onset':pred,'observed_grid_onset':onset,'rows':rows},
         'envelope':envelope(),'baseline_episode':baseline_episode(),'shield_episode':shield_episode()}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()