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

SEED = 2068

def update_radius(radius, delta, gamma, alpha):
    """Adaptive conformal recursion with the optional nonnegative projection."""
    return max(0.0, radius + gamma * (float(delta > radius) - alpha))

def conservative_radius(ensemble_radii, beta):
    """Empirical upper (1-beta) quantile across estimators."""
    return float(np.quantile(np.asarray(ensemble_radii), 1.0 - beta, method="linear"))

def stationary_sweep():
    rng = np.random.default_rng(SEED)
    rows = []
    for alpha in (0.10, 0.20, 0.30, 0.40):
        for gamma in (0.01, 0.03):
            r = 0.5
            hits = []
            for t in range(120000):
                d = rng.exponential(1.0)
                hit = d > r
                r = update_radius(r, d, gamma, alpha)
                if t >= 20000:
                    hits.append(hit)
            observed = float(np.mean(hits))
            rows.append({"alpha": alpha, "gamma": gamma,
                         "predicted_exceedance": alpha,
                         "observed_exceedance": observed,
                         "abs_error": abs(observed - alpha)})
    return rows

def ramp_sweep():
    # If every error exceeds r, dr/dt = gamma(1-alpha).
    rows = []
    alpha, r0, target = 0.2, 0.2, 1.4
    for gamma in (0.02, 0.04, 0.08, 0.16):
        r, n = r0, 0
        while r < target and n < 10000:
            r = update_radius(r, 100.0, gamma, alpha)
            n += 1
        predicted = math.ceil((target - r0) / (gamma * (1 - alpha)))
        rows.append({"gamma": gamma, "predicted_steps": predicted,
                     "observed_steps": n,
                     "relative_error": abs(n - predicted) / predicted})
    return rows

def recovery_sweep():
    # If every error is below r, dr/dt = -gamma*alpha.
    rows = []
    alpha, r0, target = 0.2, 1.8, 0.4
    for gamma in (0.02, 0.04, 0.08, 0.16):
        r, n = r0, 0
        while r > target and n < 10000:
            r = update_radius(r, 0.0, gamma, alpha)
            n += 1
        predicted = math.ceil((r0 - target) / (gamma * alpha))
        rows.append({"gamma": gamma, "predicted_steps": predicted,
                     "observed_steps": n,
                     "relative_error": abs(n - predicted) / predicted})
    return rows

def crossing_comparison():
    # Predicted clearance is 1.0 m. A shifted predictor error of .75 m makes
    # actual clearance .25 m, below the .55 m physical threshold. Benign error
    # is .15 m, giving .85 m actual clearance. Strictly use < for collisions.
    alpha, gamma = 0.2, 0.08
    errors = np.r_[np.full(60, .15), np.full(80, .75), np.full(60, .15)]
    actual_clearance = 1.0 - errors
    unsafe = actual_clearance < .55
    out = {}
    for name, kind in (("fixed_0.20", "fixed"), ("fixed_0.80", "fixed"),
                       ("adaptive", "adaptive")):
        r = .20
        accepted, radii = [], []
        for e in errors:
            if kind == "adaptive":
                radius_used = r
                r = update_radius(r, e, gamma, alpha)
            else:
                radius_used = float(name.split("_")[1])
            accepted.append(1.0 > .55 + radius_used)
            radii.append(radius_used)
        accepted = np.asarray(accepted, dtype=bool)
        out[name] = {
            "unsafe_frames": int(unsafe.sum()),
            "false_safe_frames": int(np.sum(accepted & unsafe)),
            "safe_rejections": int(np.sum((~accepted) & (~unsafe))),
            "mean_radius": float(np.mean(radii)),
            "accept_rate": float(np.mean(accepted))}
    # A direct check of the ensemble conservatism formula.
    out["quantile_check"] = {"radii": [0.2, 0.4, 0.8, 1.0], "beta": 0.25,
                             "upper_quantile": conservative_radius([0.2, 0.4, 0.8, 1.0], 0.25)}
    return out

def main():
    stationary, ramp, recovery = stationary_sweep(), ramp_sweep(), recovery_sweep()
    crossing = crossing_comparison()
    coverage_ok = max(x["abs_error"] for x in stationary) < .02
    ramp_ok = max(x["relative_error"] for x in ramp) == 0.0
    recovery_ok = max(x["relative_error"] for x in recovery) == 0.0
    result = {"seed": SEED, "stationary_coverage": stationary,
              "post_shift_ramp": ramp, "recovery_decay": recovery,
              "crossing_comparison": crossing,
              "checks": {"stationary_within_0.02": coverage_ok,
                         "ramp_formula_exact": ramp_ok,
                         "recovery_formula_exact": recovery_ok,
                         "mechanism_manifested": coverage_ok and ramp_ok and recovery_ok}}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()