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


def conformal_quantile(scores, alpha):
    scores = np.sort(np.asarray(scores, dtype=float))
    m = len(scores)
    k = int(math.ceil((m + 1) * (1 - alpha)))
    return float("inf") if k > m else float(scores[k - 1])


def kalman_run(y, q=0.04, r=0.25, gate=None, r_inflate=1.0):
    xhat, p = 0.0, 1.0
    estimates, innovations, triggered = [], [], []
    for obs in y:
        p_pred = p + q
        innov = float(obs - xhat)
        s = p_pred + r
        hit = gate is not None and abs(innov) / math.sqrt(s) > gate
        # Correct robust intervention: distrust the current suspicious observation.
        r_used = r * r_inflate if hit else r
        gain = p_pred / (p_pred + r_used)
        xhat = xhat + gain * innov
        p = (1.0 - gain) * p_pred
        estimates.append(xhat)
        innovations.append(innov / math.sqrt(s))
        triggered.append(bool(hit))
    return np.asarray(estimates), np.asarray(innovations), np.asarray(triggered)


def make_sequence(rng, n, q=0.04, r=0.25, outliers=False, outlier_prob=.08, outlier_scale=4.):
    x = np.zeros(n)
    for t in range(1, n):
        x[t] = x[t-1] + rng.normal(0, math.sqrt(q))
    y = x + rng.normal(0, math.sqrt(r), n)
    if outliers:
        mask = rng.random(n) < outlier_prob
        y[mask] += rng.normal(0, outlier_scale * math.sqrt(r), mask.sum())
    return x, y


def collect_scores(rng, n_scores, score_kind="state"):
    vals = []
    while len(vals) < n_scores:
        x, y = make_sequence(rng, 50)
        est, inn, _ = kalman_run(y)
        vals.extend(np.abs(x - est) if score_kind == "state" else np.abs(inn))
    return np.asarray(vals[:n_scores])


def coverage_for_q(rng, q, n_sequences=300):
    hits = []
    for _ in range(n_sequences):
        x, y = make_sequence(rng, 50)
        est, _, _ = kalman_run(y)
        hits.extend(np.abs(x - est) <= q)
    return float(np.mean(hits))


def run(seed=902):
    rng = np.random.default_rng(seed)
    alpha_values = (.05, .10, .20, .30)
    cal = collect_scores(rng, 5000)
    sigma = float(np.std(cal))
    # For |N(0,sigma^2)|, the exact reference is sigma*Phi^{-1}(1-alpha/2).
    zhalf = {0.05: 1.959964, 0.10: 1.644854, 0.20: 1.281552, 0.30: 1.036433}
    qvals = {str(a): conformal_quantile(cal, a) for a in alpha_values}
    coverage = {str(a): coverage_for_q(rng, qvals[str(a)]) for a in alpha_values}
    halfnormal_reference = {str(a): sigma * zhalf[a] for a in alpha_values}

    # Calibration-size prediction: rank quantile converges with approximately O(m^-1/2)
    # sampling variation, while expected coverage remains near 1-alpha.
    size_sweep = {}
    for m in (100, 500, 2000, 5000):
        qs = []
        covs = []
        for rep in range(20):
            local = np.random.default_rng(seed + 10000 + m * 10 + rep)
            scores = collect_scores(local, m)
            q = conformal_quantile(scores, .10)
            qs.append(q)
            covs.append(coverage_for_q(local, q, 100))
        size_sweep[str(m)] = {
            "q_mean": float(np.mean(qs)), "q_std": float(np.std(qs)),
            "coverage_mean": float(np.mean(covs)), "coverage_std": float(np.std(covs)),
            "predicted_q_std_scaling": 1.0 / math.sqrt(m)
        }

    # Deployment intervention: clean innovation gate, then contaminated observations.
    gate_scores = collect_scores(rng, 5000, "innovation")
    gate = conformal_quantile(gate_scores, .10)
    baseline, gated, triggers = [], [], []
    for _ in range(300):
        x, y = make_sequence(rng, 50, outliers=True, outlier_prob=.08, outlier_scale=4.)
        eb, _, _ = kalman_run(y)
        eg, _, tr = kalman_run(y, gate=gate, r_inflate=8.0)
        baseline.append(np.sqrt(np.mean((x-eb) ** 2)))
        gated.append(np.sqrt(np.mean((x-eg) ** 2)))
        triggers.append(np.mean(tr))
    result = {
        "seed": seed, "calibration_scores": len(cal), "alpha": list(alpha_values),
        "state_quantiles": qvals, "state_coverage": coverage,
        "residual_std": sigma, "halfnormal_reference": halfnormal_reference,
        "calibration_size_sweep": size_sweep,
        "innovation_threshold": gate,
        "contaminated_rmse_baseline": float(np.mean(baseline)),
        "contaminated_rmse_gated": float(np.mean(gated)),
        "contaminated_rmse_improvement_fraction": float(1 - np.mean(gated) / np.mean(baseline)),
        "trigger_rate": float(np.mean(triggers)),
        "predictions": {
            "coverage": "exchangeable marginal coverage is at least approximately 1-alpha",
            "alpha_scaling": "q decreases with alpha and follows the half-normal reference",
            "calibration_scaling": "quantile variability decreases approximately as m^(-1/2)",
            "robust_intervention": "R inflation on innovation-triggered observations reduces error under outliers"
        }
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    run()