import json
import math
import time
import numpy as np


def kl_bits(q, p):
    return float(np.sum(q * (np.log(q) - np.log(p)) / np.log(2.0)))


def l1(q, p):
    return float(np.sum(np.abs(q - p)))


def poisson_race(p, q, rng, max_candidates=100000):
    """Exact marked-Poisson race with theorem stopping rule."""
    logp, logq = np.log(p), np.log(q)
    logw_max = float(np.max(logq - logp))
    # Draw the first arrival; subsequent arrivals are carried into the next iteration.
    arrival = -math.log(max(rng.random(), np.finfo(float).tiny))
    best_score, best_x = math.inf, -1
    for i in range(1, max_candidates + 1):
        x = int(rng.choice(len(p), p=p))
        score = math.log(arrival) + logp[x] - logq[x]
        if score < best_score:
            best_score, best_x = score, x
        next_arrival = arrival - math.log(max(rng.random(), np.finfo(float).tiny))
        if math.log(next_arrival) - logw_max >= best_score:
            return best_x, i
        arrival = next_arrival
    raise RuntimeError("candidate cap exceeded")


def rejection_sample(p, q, rng, max_candidates=100000):
    """Standard rejection sampler with envelope M=max_x Q(x)/P(x)."""
    m = float(np.max(q / p))
    for i in range(1, max_candidates + 1):
        x = int(rng.choice(len(p), p=p))
        if rng.random() <= q[x] / (m * p[x]):
            return x, i
    raise RuntimeError("rejection candidate cap exceeded")


def run_case(name, p, q, n=30000, seed=1234):
    rng = np.random.default_rng(seed)
    counts = np.zeros(len(p), dtype=np.int64)
    race_idx = np.empty(n, dtype=np.int64)
    t0 = time.perf_counter()
    for k in range(n):
        x, i = poisson_race(p, q, rng)
        counts[x] += 1
        race_idx[k] = i
    race_seconds = time.perf_counter() - t0

    rng = np.random.default_rng(seed + 1)
    rej_idx = np.empty(n, dtype=np.int64)
    rej_counts = np.zeros(len(p), dtype=np.int64)
    t0 = time.perf_counter()
    for k in range(n):
        x, i = rejection_sample(p, q, rng)
        rej_counts[x] += 1
        rej_idx[k] = i
    rej_seconds = time.perf_counter() - t0

    empirical_race = counts / n
    empirical_rej = rej_counts / n
    return {
        "name": name,
        "kl_bits": kl_bits(q, p),
        "l1": l1(q, p),
        "upper_bound_bits": kl_bits(q, p) + 1.45 * l1(q, p),
        "theorem_lower_bound_bits": 0.5 * max(kl_bits(q, p), l1(q, p)),
        "race_tv": 0.5 * float(np.sum(np.abs(empirical_race - q))),
        "rejection_tv": 0.5 * float(np.sum(np.abs(empirical_rej - q))),
        "race_mean_log2_index": float(np.mean(np.log2(race_idx))),
        "race_mean_index": float(np.mean(race_idx)),
        "rejection_mean_index": float(np.mean(rej_idx)),
        "race_p95_index": float(np.quantile(race_idx, 0.95)),
        "rejection_p95_index": float(np.quantile(rej_idx, 0.95)),
        "race_seconds": race_seconds,
        "rejection_seconds": rej_seconds,
        "race_target_evals_per_sample": 1.0,
        "rejection_target_evals_per_sample": 1.0,
    }


def main():
    # All probabilities are positive, making the finite-support KL well defined.
    v = 16
    p = np.full(v, 1.0 / v)
    cases = []
    for name, alpha in [("identical", 0.0), ("close", 0.10), ("moderate", 0.30), ("farther", 0.60)]:
        # Q is a smooth tilt of P toward the first symbol; alpha controls distance.
        q = np.full(v, (1.0 - alpha) / (v - 1))
        q[0] = alpha if alpha > 0 else 1.0 / v
        if alpha == 0:
            q = p.copy()
        cases.append(run_case(name, p, q))
    result = {
        "description": "Finite categorical verification of KL-budgeted Poisson-race sampling",
        "n_samples_per_case": 30000,
        "vocab": v,
        "cases": cases,
        "notes": [
            "Race samples use iid P marks and Exp(1) arrivals; stopping uses known max Q/P.",
            "TV is halved L1, while the theorem's displayed L1 is reported separately.",
            "Target evaluation count is not a speed advantage here: q is explicitly available; this is an exactness/candidate-index test."
        ]
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()