import json
import math
import numpy as np
from scipy.linalg import solve_continuous_lyapunov

SEED = 7
rng = np.random.default_rng(SEED)
DT = 0.01
U_MAX = 1.0
SAFE_P, SAFE_V = 2.0, 2.0

# Double integrator: p_dot=v, v_dot=u.  The backup is a saturated PD law.
def backup(x, kp=2.0, kd=2.0):
    p, v = np.asarray(x)
    return float(np.clip(-kp*p-kd*v, -U_MAX, U_MAX))

def plant_step(x, u, dt=DT):
    p, v = x
    # exact constant-input step
    return np.array([p + dt*v + .5*dt*dt*u, v + dt*u])

def safe(x):
    return abs(x[0]) <= SAFE_P + 1e-9 and abs(x[1]) <= SAFE_V + 1e-9

def terminal(x):
    return x[0]*x[0] + x[1]*x[1] <= 0.06**2

def simulate_backup(x0, horizon, kp=2.0, kd=2.0):
    x = np.array(x0, dtype=float)
    n = int(round(horizon / DT))
    max_violation = 0.0
    for _ in range(n):
        max_violation = max(max_violation, max(abs(x[0])-SAFE_P, abs(x[1])-SAFE_V))
        x = plant_step(x, backup(x, kp, kd))
    max_violation = max(max_violation, max(abs(x[0])-SAFE_P, abs(x[1])-SAFE_V))
    return x, max_violation

# Quadratic certificate V=x'Px. P is the exact Lyapunov solution for the
# unsaturated closed-loop matrix.  On a small enough sublevel set saturation
# is inactive and -Vdot = ||x||^2 exactly.
def certificate(kp=2.0, kd=2.0):
    A = np.array([[0., 1.], [-kp, -kd]])
    P = solve_continuous_lyapunov(A.T, -np.eye(2))
    return A, P

def core_math_check():
    A, P = certificate()
    # numerical identity and rate bound on the certified ellipsoid V<=gamma
    Q = -(A.T @ P + P @ A)
    identity_err = float(np.max(np.abs(Q - np.eye(2))))
    gamma = 0.06**2 / np.max(np.linalg.eigvalsh(P))
    vals = []
    max_u = 0.
    for th in np.linspace(0, 2*np.pi, 2000, endpoint=False):
        z = np.array([math.cos(th), math.sin(th)])
        x = math.sqrt(gamma) * z / math.sqrt(z @ P @ z)
        u = -2*x[0]-2*x[1]
        max_u = max(max_u, abs(u))
        vals.append(float(x @ Q @ x))
    # Since Q=I, -dV/dt = ||x||^2 and normalized dissipation is 1/lambda_max(P).
    observed_rate = min(vals) / gamma
    predicted_rate = 1.0 / np.max(np.linalg.eigvalsh(P))
    return {"P": P.tolist(), "identity_error": identity_err,
            "gamma": gamma, "max_backup_u_on_certificate": max_u,
            "observed_min_neg_vdot_over_V": observed_rate,
            "predicted_rate_bound": predicted_rate,
            "certificate_saturation_free": bool(max_u <= U_MAX + 1e-10)}

def certificate_scaling():
    # For V=x'P x and terminal disk radius r, gamma=r^2/lambda_max(P)
    # and area(Omega_0)=pi*gamma/sqrt(det(P)); hence area scales as r^2.
    _, P = certificate()
    lam = np.max(np.linalg.eigvalsh(P))
    det = np.linalg.det(P)
    rows = []
    base = math.pi * (0.06**2 / lam) / math.sqrt(det)
    for r in [0.03, 0.06, 0.12, 0.18]:
        area = math.pi * (r*r / lam) / math.sqrt(det)
        rows.append({"terminal_radius": r, "certified_area": area,
                     "area_over_base": area/base, "predicted_ratio": (r/0.06)**2})
    return rows

def reachable_sweep():
    # Grid estimate of the finite-horizon backward reachable set to the
    # terminal safe disk, with safety checked throughout the backup rollout.
    points = np.array([(p,v) for p in np.linspace(-1.8,1.8,61)
                              for v in np.linspace(-1.8,1.8,61)])
    cell = (3.6/60.)**2
    out = []
    for T in [0.5, 1.0, 2.0, 3.0, 4.0]:
        good = 0
        for x in points:
            xf, viol = simulate_backup(x, T)
            if viol <= 1e-8 and terminal(xf):
                good += 1
        out.append({"T": T, "volume": good*cell, "fraction": good/len(points)})
    # input limit transition: longer horizons cannot rescue points requiring
    # excessive initial acceleration; report saturation prevalence.
    sat = []
    for kp in [0.5, 1., 2., 4., 8.]:
        count = 0
        for x in points[::20]:
            if abs(backup(x, kp, 2.0)) >= U_MAX-1e-12: count += 1
        sat.append({"kp": kp, "saturated_fraction": count/len(points[::20])})
    return out, sat

def certified_rollout_audit(n=500, horizon=4.0):
    _, P = certificate()
    rg = np.random.default_rng(SEED + 2)
    gamma = 0.06**2
    safe_count = terminal_count = 0
    max_violation = 0.0
    for _ in range(n):
        z = rg.normal(size=2)
        z /= np.linalg.norm(z)
        x = z * math.sqrt(gamma * rg.random() / (z @ P @ z))
        xf, viol = simulate_backup(x, horizon)
        safe_count += int(viol <= 1e-8)
        terminal_count += int(terminal(xf))
        max_violation = max(max_violation, viol)
    return {"n": n, "safe_rate": safe_count / n,
            "terminal_rate": terminal_count / n,
            "max_violation": max_violation, "horizon": horizon}

def nominal(x, rng):
    # Fixed seeded proxy for a poorly trained learned policy: it tends to
    # accelerate toward the positive position boundary and has small noise.
    p, v = x
    return float(np.clip(0.35 + 0.10*p + 0.10*v + rng.normal(0, .08), -1, 1))

def shielded_rollout(x0, steps=700, rng=None):
    x = np.array(x0, dtype=float); switches=0; violation=0.
    for _ in range(steps):
        h = SAFE_P*SAFE_P - x[0]*x[0]
        # h_dot + alpha(h), matching the stated barrier feasibility test.
        up = nominal(x, rng)
        margin = -2*x[0]*x[1] + .5*h
        A, P = certificate()
        inside = x @ P @ x <= 0.06**2 + 1e-12
        if margin >= 0 and inside:
            u = up
        else:
            u = backup(x); switches += 1
        x = plant_step(x, u)
        violation = max(violation, max(0., abs(x[0])-SAFE_P, abs(x[1])-SAFE_V))
    return violation, switches

def baseline_rollout(x0, steps=700, rng=None):
    x=np.array(x0,float); violation=0.
    for _ in range(steps):
        x=plant_step(x, nominal(x,rng))
        violation=max(violation,max(0.,abs(x[0])-SAFE_P,abs(x[1])-SAFE_V))
    return violation

def shield_experiment():
    # Initial states in a ring near the safety boundary, same draws for both.
    rg=np.random.default_rng(SEED+1)
    xs=[]
    while len(xs)<1000:
        x=rg.uniform([-1.9,-1.5],[1.9,1.5])
        if abs(x[0])>1.4: xs.append(x)
    bviol=[]; sviol=[]; sw=[]
    for x in xs:
        bviol.append(baseline_rollout(x, rng=np.random.default_rng(100+len(bviol))))
        z,n=shielded_rollout(x, rng=np.random.default_rng(100+len(sviol)))
        sviol.append(z); sw.append(n)
    return {"n":len(xs), "baseline_violation_rate":float(np.mean(np.array(bviol)>1e-8)),
            "shield_violation_rate":float(np.mean(np.array(sviol)>1e-8)),
            "baseline_mean_violation":float(np.mean(bviol)),
            "shield_mean_violation":float(np.mean(sviol)),
            "shield_mean_switches":float(np.mean(sw))}

def main():
    reach, sat = reachable_sweep()
    result={"seed":SEED, "math_check":core_math_check(),
            "prediction_1_dissipation": "-Vdot/V >= 1/lambda_max(P) on unsaturated certificate",
            "prediction_2_horizon": "backward reachable volume is nondecreasing with T",
            "prediction_3_scaling": "quadratic certificate area scales as terminal_radius^2",
            "certificate_scaling_sweep": certificate_scaling(),
            "reachability_sweep":reach, "input_saturation_sweep":sat,
            "certified_rollout_audit":certified_rollout_audit(),
            "policy_comparison":shield_experiment()}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__ == '__main__': main()