"""Extra outputs for the selection calculation: smaller error widths, the true value of a
compound calculated at 53 K, the best true value in a pool of a million, and the limit
tau / sigma^2 above which an exponential tail in calculated values and log-normal error
cannot both hold. Same model as selection.py."""

import json
import math
from pathlib import Path

import numpy as np
from scipy.optimize import brentq

HERE = Path(__file__).parent
TAU_OBS, U = 4.03, 10.0
BINS = [(10, 15), (15, 20), (20, 30), (30, 45), (45, 70)]


def obs_tau(tt, s, seed):
    r = np.random.default_rng(seed); t = r.exponential(tt, 4_000_000); y = t * np.exp(r.normal(0, s, t.size)); e = y[y >= U] - U
    return e.mean()


out = []
for i, s in enumerate((0.15, 0.2, 0.3, 0.4, 0.5)):
    tt = brentq(lambda x: obs_tau(x, s, 21 + i) - TAU_OBS, 1.0, 4.03, xtol=1e-3)
    r = np.random.default_rng(100 + i); n = 80_000_000
    t = r.exponential(tt, n); y = t * np.exp(r.normal(0, s, n))
    row = {"sigma": s, "tau_true": tt, "limit_tau_over_sigma2": TAU_OBS / s**2, "bins": []}
    for lo, hi in BINS:
        m = (y >= lo) & (y < hi)
        row["bins"].append({"lo": lo, "hi": hi, "n": int(m.sum()), "median_ratio": float(np.median(y[m] / t[m])) if m.sum() > 200 else None})
    m = (y >= 49) & (y < 57)
    row["true_given_calc_53"] = {"n": int(m.sum()), "median_true": float(np.median(t[m])) if m.sum() > 50 else None}
    frac = float(np.mean(y >= U))
    n_pool = 1e6 * 0.0427 / frac        # pool in which 4.27% of calculated values exceed 10 K
    row["best_true_in_million_mode"] = tt * math.log(n_pool)
    row["frac_calc_above_10"] = frac
    out.append(row)
(HERE / "selection_checks.json").write_text(json.dumps(out, indent=1))
for o in out:
    print(f"sigma={o['sigma']}: tau_true={o['tau_true']:.2f}; limit tau/sigma^2={o['limit_tau_over_sigma2']:.0f} K; ratios", [None if b['median_ratio'] is None else round(b['median_ratio'], 2) for b in o["bins"]], "| calc 53 K -> true", o["true_given_calc_53"], f"| best true in 1e6: {o['best_true_in_million_mode']:.0f} K")