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

# Robust HOCBF for p_dot=v, v_dot=u+w, h>=0, with linear alpha_i(s)=k_i s.
# For lower wall p>=pmin: u >= -what + eps -(k1+k2)v-k1*k2*(p-pmin)
# For upper wall p<=pmax: u <= -what - eps -(k1+k2)v+k1*k2*(pmax-p)

def interval(p, v, pmin, pmax, umax, what=0.0, eps=0.0, k1=2.0, k2=2.0):
    lo = -what + eps - (k1+k2)*v - k1*k2*(p-pmin)
    hi = -what - eps - (k1+k2)*v + k1*k2*(pmax-p)
    return max(-umax, lo), min(umax, hi)

def project(u, lo, hi):
    # Emergency action is the closest actuator endpoint if uncertainty makes I empty.
    return float(np.clip(u, lo, hi)) if lo <= hi else float(np.clip(u, [-1, 1][int(u < 0)], [1, -1][int(u < 0)]))

def nominal_policy(p, v, target=0.0):
    # Small neural-policy surrogate: bounded PD policy with deliberately aggressive targets.
    return float(np.clip(-2.5*(p-target)-1.3*v, -1.0, 1.0))

def exact_margin(p, v, **kw):
    lo, hi = interval(p, v, **kw)
    return hi-lo

def math_sweeps():
    # Use a wide actuator only for algebraic verification, avoiding saturation masking.
    pars=dict(pmin=-1.0,pmax=1.0,umax=10.0,k1=2.0,k2=2.0)
    # Prediction 1: M(eps)=M(0)-2 eps, until actuator clipping changes the slope.
    p,v=0.25,0.35
    m0=exact_margin(p,v,what=0,eps=0,**pars)
    eps_grid=np.linspace(0,2.4,13)
    margins=np.array([exact_margin(p,v,what=0,eps=e,**pars) for e in eps_grid])
    # Use only nonempty/non-clipped points to estimate slope; report all values.
    slope=float(np.polyfit(eps_grid[margins>0],margins[margins>0],1)[0])
    # Prediction 2: symmetric disturbance uncertainty becomes infeasible at eps=M0/2.
    epscrit=m0/2
    # Directly estimate the first epsilon at which the interval is empty.
    dense=np.linspace(0, 5.0, 5001)
    dense_m=np.array([exact_margin(p,v,what=0,eps=e,**pars) for e in dense])
    first_empty=dense[np.flatnonzero(dense_m < 0)[0]] if np.any(dense_m < 0) else None
    # Prediction 3: with no actuator clipping, uncertainty shifts each bound by eps.
    p2,v2=.0,0.
    l0,h0=interval(p2,v2,**pars,what=0,eps=0)
    shifts=[]
    for e in [.1,.2,.4]:
        l,h=interval(p2,v2,**pars,what=0,eps=e)
        shifts.append((e,l-l0,h-h0))
    return {'state':[p,v], 'M0':m0, 'eps_grid':eps_grid.tolist(), 'margins':margins.tolist(),
            'fitted_margin_slope':slope, 'predicted_margin_slope':-2.0,
            'predicted_eps_critical':epscrit, 'observed_eps_critical':first_empty,
            'bound_shift_samples':shifts, 'predicted_shifts': 'lower +eps, upper -eps'}

def simulate(mode, disturbance, eps_bound, seed=7, T=8.0, dt=.002):
    rng=np.random.default_rng(seed)
    p,v=.72,-.05
    pmin,pmax,umax=-1.,1.,1.
    max_violation=0.; interventions=[]; margins=[]; bound_misses=0
    n=int(T/dt)
    for i in range(n):
        t=i*dt
        # bounded time-varying disturbance, plus fixed observer estimate zero
        w=disturbance*(0.65*math.sin(1.7*t)+0.35*math.sin(4.1*t+0.3))
        un=nominal_policy(p,v,target=-.65 if t<3 else .65)
        if mode=='clip':
            lo,hi=-umax,umax
            us=float(np.clip(un,lo,hi))
            margin=2*umax
        elif mode=='cbf':
            lo,hi=interval(p,v,pmin,pmax,umax,what=0,eps=0)
            us=project(un,lo,hi); margin=hi-lo
        else:
            lo,hi=interval(p,v,pmin,pmax,umax,what=0,eps=eps_bound)
            us=project(un,lo,hi); margin=hi-lo
        # Euler is sufficiently fine for this sanity test; include continuous sampled states.
        interventions.append(abs(us-un)); margins.append(margin)
        p += dt*v
        v += dt*(us+w)
        max_violation=max(max_violation,max(0,pmin-p,p-pmax))
        if abs(w)>eps_bound+1e-9: bound_misses += 1
    return {'max_violation':float(max_violation), 'mean_intervention':float(np.mean(interventions)),
            'min_margin':float(np.min(margins)), 'bound_misses':bound_misses}

def run():
    random.seed(7); np.random.seed(7)
    math_result=math_sweeps()
    rows=[]
    for d in [0,.1,.2,.3,.4,.5,.6,.8,1.0]:
        rows.append({'disturbance_amplitude':d,
          'clip':simulate('clip',d,.35), 'ordinary_cbf':simulate('cbf',d,.35),
          'robust_cbf':simulate('robust',d,.35)})
    result={'math_verification':math_result,'closed_loop':rows,
      'notes':'Disturbance bound is eps=.35; actual sinusoid peak equals amplitude. Empty robust intervals use closest actuator endpoint emergency fallback.'}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': run()