import json
import math
from pathlib import Path

import numpy as np


def required_rollouts(r_low: float, delta: float) -> int:
    """Smallest integer N for (1-r_low)^N <= delta."""
    if not 0.0 < r_low < 1.0 or not 0.0 < delta < 1.0:
        raise ValueError("r_low and delta must be in (0,1)")
    return int(math.ceil(math.log(delta) / math.log1p(-r_low)))


def beta_prior_lower_bound(events: int, total: int, alpha: float = 0.05) -> float:
    """Simple conservative one-sided lower estimate.

    The zero-event branch is the 1-alpha lower credible bound induced by a
    Beta(1,1) prior. For nonzero counts, Wilson's one-sided lower bound is used.
    """
    if total <= 0 or events < 0 or events > total:
        raise ValueError("invalid event counts")
    if events == 0:
        return 1.0 - alpha ** (1.0 / (total + 1.0))
    p, z = events / total, 1.6448536269514722
    denom = 1.0 + z * z / total
    center = (p + z * z / (2 * total)) / denom
    spread = z * math.sqrt(p * (1-p) / total + z*z / (4 * total * total)) / denom
    return max(0.0, center - spread)


def exact_miss_sweep(seed=7, trials=100_000):
    rng = np.random.default_rng(seed)
    rows = []
    for r in (0.01, 0.03, 0.10, 0.25):
        for n in (5, 20, 50):
            miss = rng.binomial(n, r, size=trials) == 0
            observed = float(np.mean(miss))
            predicted = (1-r) ** n
            rows.append({"r": r, "N": n, "observed_miss": observed,
                         "predicted_miss": predicted,
                         "abs_error": abs(observed - predicted)})
    return rows


def sizing_sweep(delta=0.05):
    rows = []
    for r in (0.005, 0.01, 0.02, 0.05, 0.10):
        n = required_rollouts(r, delta)
        exact = (1-r) ** n
        approximation = math.log(1/delta) / r
        rows.append({"r": r, "N_required": n, "exact_miss": exact,
                     "target_delta": delta, "small_r_approx_N": approximation,
                     "N_over_approx": n / approximation})
    return rows


def probe_sweep(seed=11, trials=100_000):
    """Matched-budget detection experiment.

    A random rollout sees the critical mode with probability r. A directed probe
    sees it with probability q. The proposed gate spends 30 random + 20 probes;
    the baseline spends all 50 randomly. Independent draws make the predicted
    proposed miss probability (1-r)^30 (1-q)^20.
    """
    rng = np.random.default_rng(seed)
    rows = []
    for r in (0.005, 0.01, 0.02, 0.05):
        q = min(0.60, 30*r)  # toy boundary score: probes amplify rare modes
        baseline_miss = float(np.mean(rng.binomial(50, r, trials) == 0))
        idea_miss = float(np.mean((rng.binomial(30, r, trials) +
                                   rng.binomial(20, q, trials)) == 0))
        predicted_idea_miss = (1-r)**30 * (1-q)**20
        rows.append({"r": r, "probe_rate_q": q,
                     "baseline_miss_observed": baseline_miss,
                     "baseline_miss_predicted": (1-r)**50,
                     "idea_miss_observed": idea_miss,
                     "idea_miss_predicted": predicted_idea_miss,
                     "detection_gain": (1-baseline_miss) - (1-idea_miss),
                     "probe_scaling_factor_q_over_r": q/r})
    return rows


def validate(rows, sizing, probes):
    max_exact_error = max(x["abs_error"] for x in rows)
    max_probe_error = max(abs(x["idea_miss_observed"] - x["idea_miss_predicted"])
                          for x in probes)
    assert max_exact_error < 0.01
    assert max_probe_error < 0.01
    assert all(x["exact_miss"] <= x["target_delta"] for x in sizing)
    # Predictions: exact exponential scaling; ceil-sized N meets delta; probes
    # reduce miss probability for every tested rare-event rate.
    assert all(x["idea_miss_predicted"] < x["baseline_miss_predicted"] for x in probes)
    return {"max_exact_abs_error": max_exact_error,
            "max_probe_abs_error": max_probe_error,
            "all_sizing_rows_meet_delta": True,
            "all_probe_rows_improve_over_baseline": True}


def main():
    exact = exact_miss_sweep()
    sizing = sizing_sweep()
    probes = probe_sweep()
    checks = validate(exact, sizing, probes)
    result = {"checks": checks, "exact_miss_sweep": exact,
              "required_count_sweep": sizing, "matched_budget_probe_sweep": probes}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()