import json
import numpy as np
from trajectory_tube import TrajectoryTube, conformal_quantile

SEED = 2834
rng = np.random.default_rng(SEED)
alpha = 0.1
alpha_L = 0.1
T = np.linspace(0.0, 1.0, 1001)
obs = np.array([0.1, 0.3, 0.5, 0.75])

# A continuous piecewise-linear trajectory makes the Lipschitz proof exact.
def ramp(t, slope):
    return slope * np.maximum(np.asarray(t) - 0.75, 0.0)

# Calibration residuals are zero because the predictor is zero and observations
# stop at the kink. Dense calibration gives the separate slope quantile.
cal_slopes = np.linspace(0.4, 0.95, 20)
cal_x = [ramp(T, a) for a in cal_slopes]
cal_obs = [ramp(obs, a) for a in cal_slopes]
q = conformal_quantile([np.max(np.abs(x)) for x in cal_obs], alpha)
Lhat = conformal_quantile(cal_slopes, alpha_L)

# Prediction 1: radius must be affine in nearest-observation distance.
gamma = Lhat
r, delta = TrajectoryTube(q, gamma).radius(T, obs)
fit_slope, fit_intercept = np.polyfit(delta, r, 1)
linear_max_error = np.max(np.abs(r - (q + gamma * delta)))

# Prediction 2: conditional transition occurs at multiplier one for a truth
# whose slope equals Lhat. A deliberately sublinear multiplier must fail in
# the largest gap (at the first post-kink point).
test_boundary = ramp(T, Lhat)
mults = np.array([0.0, 0.25, 0.5, 0.75, 0.9, 1.0, 1.1, 1.25, 1.5])
cover_boundary = []
for m in mults:
    tube = TrajectoryTube(q, m * Lhat)
    cover_boundary.append(np.mean(tube.contains(test_boundary, np.zeros_like(T), T, obs)))
cover_boundary = np.asarray(cover_boundary)
full = np.where(cover_boundary >= 1 - 1e-12)[0]
first_full = float(mults[full[0]]) if len(full) else None
sub_m = 0.75
inside_sub = TrajectoryTube(q, sub_m * Lhat).contains(test_boundary, np.zeros_like(T), T, obs)
first_failure = float(T[np.where(~inside_sub)[0][0]]) if np.any(~inside_sub) else None

# Prediction 3: marginal coverage is controlled by the slope event. Test slopes
# are iid from the same range as calibration; at multiplier 1 the expected
# slope event is 1-alpha_L, while a safety multiplier covers all this range.
test_slopes = rng.uniform(0.4, 0.95, 1000)
cover_m1 = []
cover_m11 = []
for a in test_slopes:
    truth = ramp(T, a)
    cover_m1.append(np.mean(TrajectoryTube(q, Lhat).contains(truth, np.zeros_like(T), T, obs)))
    cover_m11.append(np.mean(TrajectoryTube(q, 1.1 * Lhat).contains(truth, np.zeros_like(T), T, obs)))
marginal_m1 = float(np.mean(cover_m1))
marginal_m11 = float(np.mean(cover_m11))
slope_event = float(np.mean(test_slopes <= Lhat))

# Direct numerical triangle-inequality check for random linear truth/predictor.
violations = []
for _ in range(100):
    tt = np.linspace(0, 1, 101)
    lt, lp = rng.uniform(0.1, 2.0), rng.uniform(0.1, 1.0)
    truth, pred = lt * tt, lp * tt
    idx = np.arange(0, 101, 20)
    d = np.min(np.abs(tt[:, None] - tt[idx][None, :]), axis=1)
    q2 = np.max(np.abs(truth[idx] - pred[idx]))
    violations.append(np.max(np.abs(truth - pred) - (q2 + (lt + lp) * d)))

# Standard baseline: no between-observation inflation.
baseline = float(np.mean(np.abs(test_boundary) <= q + 1e-12))
out = {
    "seed": SEED, "alpha": alpha, "alpha_L": alpha_L, "q": q, "Lhat": Lhat,
    "predictions": {
        "radius_slope_predicted": float(gamma),
        "radius_slope_observed": float(fit_slope),
        "radius_intercept_predicted": float(q),
        "radius_intercept_observed": float(fit_intercept),
        "max_affine_error": float(linear_max_error),
        "conditional_transition_predicted_multiplier": 1.0,
        "conditional_transition_observed_first_full_multiplier": first_full,
        "sublinear_multiplier": sub_m,
        "sublinear_first_failure_time": first_failure,
        "slope_event_predicted": 1 - alpha_L,
        "slope_event_observed": slope_event,
        "marginal_coverage_at_multiplier_1": marginal_m1,
        "marginal_coverage_at_multiplier_1.1": marginal_m11,
        "triangle_max_violation": float(max(violations))
    },
    "sweep_boundary": [{"multiplier": float(m), "coverage": float(c)} for m, c in zip(mults, cover_boundary)],
    "comparison": {"baseline_constant_q_coverage": baseline,
                   "idea_boundary_coverage_at_Lhat": float(cover_boundary[mults == 1][0]),
                   "idea_marginal_coverage_at_Lhat": marginal_m1}
}
print(json.dumps(out, indent=2))