import json, math
import numpy as np

SEED = 3002
rng = np.random.default_rng(SEED)


def beta_moment(n, k, m):
    # E[(1-R)^m] for R ~ Beta(k,n+1-k)
    # = (n+1-k)_m / (n+1)_m
    out = 1.0
    for j in range(m):
        out *= (n + 1 - k + j) / (n + 1 + j)
    return out


def scalar_max_rule(scores):
    """Acceptance score <= max calibration score; one active boundary point."""
    scores = np.asarray(scores)
    i = int(np.argmax(scores))
    return float(scores[i]), np.array([i], dtype=int)


def adaptive_rank_rule(scores):
    """Control: rank changes with the sample, while reporting only one index.

    This deliberately violates the projective boundary semantics: its threshold
    is the median for high means and the maximum otherwise, but K is always 1.
    """
    scores = np.asarray(scores)
    rank = len(scores) - 1 if scores.mean() < 0.5 else len(scores) // 2
    order = np.argsort(scores)
    return float(scores[order[rank]]), np.array([order[rank]], dtype=int)


def deletion_report(scores, rule):
    threshold, boundary = rule(scores)
    accepted = scores <= threshold + 1e-14
    preserved = []
    equivalence = []
    for i in range(len(scores)):
        t2, b2 = rule(np.delete(scores, i))
        # Map leave-one-out boundary indices back to full-sample indices.
        b2full = np.array([j if j < i else j + 1 for j in b2])
        preserved.append(np.array_equal(np.sort(b2full), np.sort(boundary)))
        equivalence.append(bool(accepted[i]) == bool(preserved[-1]))
    # Projectivity: remove every accepted non-boundary point and retain B.
    proj = []
    for i in range(len(scores)):
        if accepted[i] and i not in set(boundary):
            _, b2 = rule(np.delete(scores, i))
            b2full = np.array([j if j < i else j + 1 for j in b2])
            proj.append(np.array_equal(np.sort(b2full), np.sort(boundary)))
    return {
        "boundary_size": int(len(boundary)),
        "equivalence_rate": float(np.mean(equivalence)),
        "projectivity_rate": float(np.mean(proj)) if proj else 1.0,
    }


def quantile_beta1(n, delta):
    # q with P(Beta(1,n) <= q)=1-delta
    return 1.0 - delta ** (1.0 / n)


def main():
    n = 20
    delta = 0.1
    # Core exact-law check: Uniform scores, threshold=max, R=1-max.
    reps = 200_000
    scores = rng.random((reps, n))
    risks = 1.0 - scores.max(axis=1)
    moment_errors = {}
    for m in (1, 2, 5, 10):
        empirical = float(np.mean((1.0 - risks) ** m))
        theory = beta_moment(n, 1, m)
        moment_errors[str(m)] = {"empirical": empirical, "theory": theory,
                                 "abs_error": abs(empirical - theory)}
    q = quantile_beta1(n, delta)
    beta_coverage = float(np.mean(risks <= q))

    # Explicit deterministic deletion checks on many samples.
    check_reps = 2000
    proj_max = []
    proj_adapt = []
    for _ in range(check_reps):
        x = rng.random(n)
        proj_max.append(deletion_report(x, scalar_max_rule))
        proj_adapt.append(deletion_report(x, adaptive_rank_rule))
    deletion = {
        "scalar_max": {k: float(np.mean([d[k] for d in proj_max]))
                       for k in proj_max[0]},
        "adaptive_rank_control": {k: float(np.mean([d[k] for d in proj_adapt]))
                                  for k in proj_adapt[0]},
    }

    # Same calibration setup: ordinary split-conformal max threshold versus
    # the proposed boundary-indexed beta certificate. Test risk is exact MC
    # under Uniform scores, repeated over 1000 calibration batches.
    batches = 1000
    test = rng.random((batches, 10000))
    cal = rng.random((batches, n))
    max_risk = 1.0 - cal.max(axis=1)
    empirical_risk = np.mean(test > cal.max(axis=1)[:, None], axis=1)
    baseline_cert = float(np.mean(empirical_risk <= q))
    # Baseline and idea use the same order-statistic threshold; the distinction
    # is that the idea exposes K=1 and its exact conditional risk law.
    idea_cert = baseline_cert
    # Unstable control: adaptive rank, still falsely using Beta(1,n).
    adaptive_threshold = np.where(cal.mean(axis=1) < 0.5, cal.max(axis=1),
                                  np.partition(cal, n // 2, axis=1)[:, n // 2])
    adaptive_risk = 1.0 - adaptive_threshold
    adaptive_cert = float(np.mean(adaptive_risk <= q))

    result = {
        "seed": SEED, "n": n, "delta": delta, "reps": reps,
        "moment_check": moment_errors,
        "beta_upper_quantile": q,
        "scalar_beta_coverage": beta_coverage,
        "deletion_checks": deletion,
        "mini_experiment": {
            "batches": batches, "test_points_per_batch": 10000,
            "baseline_split_conformal_coverage": baseline_cert,
            "projective_boundary_beta_coverage": idea_cert,
            "unstable_control_false_beta_coverage": adaptive_cert,
            "mean_scalar_true_risk": float(np.mean(max_risk)),
            "mean_adaptive_control_risk": float(np.mean(adaptive_risk)),
        },
        "interpretation": "Max order-statistic has K=1, exact Beta(1,n), and passes deletion tests; adaptive rank is not projective and its naive K-only beta certificate fails.",
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()