"""Overprediction that follows from ranking candidates on a noisy calculated Tc.

True Tc, t, is drawn from an exponential distribution with scale tau_true. The calculated
value is y = t * exp(eps) with eps normal of standard deviation sigma. tau_true is set so
that the calculated values reproduce the observed tail scale above 10 K (4.03 K, from
tails.py). The script reports the median of y / t in bins of y, the expected true Tc of
the best calculated compound in a pool of 8,323, and the overlap between the calculated
and true top 25.
"""

import json
import math
from pathlib import Path

import numpy as np
from scipy.optimize import brentq

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


def observed_tau(tau_true, sigma, rng, n=4_000_000):
    t = rng.exponential(tau_true, n)
    y = t * np.exp(rng.normal(0.0, sigma, n))
    exc = y[y >= U] - U
    return exc.mean()


def run(sigma, seed):
    rng = np.random.default_rng(seed)
    tau_true = brentq(lambda x: observed_tau(x, sigma, np.random.default_rng(seed)) - TAU_OBS, 1.0, 4.03, xtol=1e-3)
    n = 60_000_000
    t = rng.exponential(tau_true, n)
    y = t * np.exp(rng.normal(0.0, sigma, n))
    rows = []
    for lo, hi in BINS:
        m = (y >= lo) & (y < hi)
        rows.append({"lo": lo, "hi": hi, "n": int(m.sum()), "median_ratio": float(np.median(y[m] / t[m])) if m.sum() > 50 else None})
    # pools: the sample has 4.27% above 10 K; scale an exponential population to match
    frac = np.mean(y >= U)
    best_y, best_t, overlap = [], [], []
    pool_n = int(POOL * 0.0427 / frac)
    for _ in range(300):
        tt = rng.exponential(tau_true, pool_n)
        yy = tt * np.exp(rng.normal(0.0, sigma, pool_n))
        i = np.argmax(yy)
        best_y.append(yy[i]); best_t.append(tt[i])
        overlap.append(len(set(np.argsort(yy)[-25:]) & set(np.argsort(tt)[-25:])) / 25)
    closed = []
    for yv in (12.5, 25.0, 55.0, 110.0):
        r = brentq(lambda x: x * math.log(x) - sigma**2 * yv / tau_true, 1.0001, 1e4)
        closed.append({"y": yv, "ratio_map": r})
    return {"sigma": sigma, "tau_true": tau_true, "bins": rows, "best_calculated": float(np.mean(best_y)),
            "true_tc_of_best": float(np.mean(best_t)), "top25_overlap": float(np.mean(overlap)), "closed_form": closed}


if __name__ == "__main__":
    out = [run(s, 11 + i) for i, s in enumerate((0.3, 0.4, 0.5))]
    (HERE / "selection.json").write_text(json.dumps(out, indent=1))
    for o in out:
        print(f"sigma={o['sigma']}: tau_true={o['tau_true']:.2f} K; best calculated {o['best_calculated']:.1f} K has true {o['true_tc_of_best']:.1f} K; top-25 overlap {o['top25_overlap']:.2f}")
        print("   median y/t by bin:", [(f"{b['lo']}-{b['hi']}", None if b['median_ratio'] is None else round(b['median_ratio'], 2), b['n']) for b in o["bins"]])
        print("   closed form R ln R = sigma^2 y / tau:", [(c['y'], round(c['ratio_map'], 2)) for c in o["closed_form"]])