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

# Discrete simulation of xdot = u + d, with safe set h(x)=1-x >= 0.
# The learned/reference model omits d: xhat_next=x+dt*u.
# For scalar u, the CBF-QP has the closed-form projection used below.

def filter_action(x, u_nom, ebar, *, dt=.1, alpha=1., kappa=1., umin=-2., umax=2.):
    h = 1.0 - x
    # hdot_model + alpha(h-kappa*ebar) - epsilon >= 0,
    # epsilon is the projected disturbance margin kappa*ebar/dt.
    upper = alpha * (h - kappa * ebar) - kappa * ebar / dt
    u = float(np.clip(min(u_nom, upper), umin, umax))
    return u, max(0., u_nom-u), upper

def smoothing_crossing(rho, target, raw, max_steps=10000):
    """First t (1-indexed) at which ebar >= target for constant raw error."""
    e = 0.
    for t in range(1, max_steps + 1):
        e = rho*e + (1-rho)*raw
        if e >= target - 1e-12:
            return t
    return None

def run_episode(kappa, rho=.8, d=.35, filtered=True, seed=0, steps=100):
    rng = np.random.default_rng(seed)
    x, ebar = 0., 0.
    violations, interventions, max_x = 0, 0., x
    crossed = None
    for t in range(steps):
        # A fixed neural-policy-like head drives toward the upper boundary.
        u_nom = .85 + .03*rng.normal()
        if filtered:
            u, intervention, _ = filter_action(x, u_nom, ebar, kappa=kappa)
        else:
            u, intervention = float(np.clip(u_nom, -2, 2)), 0.
        xhat = x + .1*u
        # Constant unknown velocity disturbance, observed after applying u.
        x = x + .1*(u + d)
        ebar = rho*ebar + (1-rho)*abs(x-xhat)
        if crossed is None and kappa*ebar/.1 >= d:
            crossed = t+1
        violations += int(x > 1.0 + 1e-10)
        interventions += intervention
        max_x = max(max_x, x)
    return dict(violations=violations, intervention=interventions, max_x=max_x,
                final_ebar=ebar, crossing=crossed)

def main():
    out = {}
    # Core math sanity: if epsilon bounds projected disturbance, robust residual is nonnegative.
    rng = np.random.default_rng(11)
    residuals = []
    for _ in range(10000):
        u = rng.uniform(-1, 1); d = rng.uniform(-.4, .4); h = rng.uniform(0, 1)
        eps = abs(d) + 1e-9
        # hdot=-u-d; model residual minus eps is a lower bound on true residual.
        model_res = -u + h - eps
        true_res = -u - d + h
        residuals.append(true_res - model_res)
    out['math_check'] = {'min_true_minus_model_lower_bound': float(min(residuals)),
                         'statement': 'true residual >= model residual when epsilon >= |d|'}

    # Prediction 1: protection turns on when kappa*ebar/dt reaches d.
    # Constant transition error is dt*d, so predicted crossing is
    # ceil(log(1-d*dt/(kappa*dt))/log(rho)) = ceil(log(1-1/kappa)/log(rho)) for kappa>1.
    rows = []
    for k in [0.8, 1., 1.2, 1.5, 2., 3.]:
        pred = None if k <= 1 else math.ceil(math.log(1-1/k)/math.log(.8))
        r = run_episode(k, rho=.8, d=.35, seed=0)
        rows.append({'kappa': k, 'predicted_margin_crossing_step': pred,
                     'observed_crossing_step': r['crossing'], 'violations': r['violations']})
    out['margin_threshold_sweep'] = rows

    # Prediction 2: ebar reaches a fixed fraction q of its asymptote with the
    # exponential law t=ceil(log(1-q)/log(rho)); verify several rho values.
    smooth = []
    for rho in [.5, .8, .9, .95]:
        q=.8; pred=math.ceil(math.log(1-q)/math.log(rho))
        obs=smoothing_crossing(rho, q*.1*.35, .1*.35)
        smooth.append({'rho':rho, 'target_fraction':q, 'predicted_step':pred, 'observed_step':obs})
    out['smoothing_law_sweep'] = smooth

    # Prediction 3: larger kappa monotonically increases intervention and decreases violations.
    pareto=[]
    for k in [0., .5, 1., 1.5, 2., 3., 4.]:
        vals=[run_episode(k, rho=.8, d=.35, seed=s) for s in range(8)]
        pareto.append({'kappa':k, 'mean_violations':float(np.mean([v['violations'] for v in vals])),
                       'mean_intervention':float(np.mean([v['intervention'] for v in vals])),
                       'mean_max_x':float(np.mean([v['max_x'] for v in vals]))})
    base=[run_episode(0, filtered=False, d=.35, seed=s) for s in range(8)]
    out['intervention_violation_sweep'] = pareto
    out['baseline_static_clipping'] = {'mean_violations':float(np.mean([v['violations'] for v in base])),
                                        'mean_max_x':float(np.mean([v['max_x'] for v in base]))}
    # Explicit mechanism verdict checks (allow one step numerical tolerance).
    out['checks'] = {
      'smoothing_matches_exact_recursion': all(abs(x['predicted_step']-x['observed_step'])<=1 for x in smooth),
      'intervention_monotone': all(pareto[i+1]['mean_intervention']+1e-9 >= pareto[i]['mean_intervention'] for i in range(len(pareto)-1)),
      'violations_nonincreasing': all(pareto[i+1]['mean_violations'] <= pareto[i]['mean_violations']+1e-9 for i in range(len(pareto)-1)),
      'adaptive_beats_clipping_at_kappa_2': pareto[4]['mean_violations'] < out['baseline_static_clipping']['mean_violations']
    }
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()