"""Closed-form arithmetic for the section on screening: the level exceeded once in the sample, the limit tau / sigma^2 above which log-normal error and an exponential tail cannot both hold, the bound on sigma that the limit gives, and the modal ratio of calculated to true Tc from R ln R = sigma^2 y / tau_t. No sampling.

Run with the project environment: .venv/bin/python numerics/derived/screening_selection.py
The printed output is stored beside this file as screening_selection.out.txt.
"""
import json
import math
import pathlib

import numpy as np

ROOT = pathlib.Path(__file__).resolve().parents[2]
checks = json.load(open(ROOT / 'numerics' / 'tails_checks.json'))
tau = checks['tau'] if isinstance(checks['tau'], (int, float)) else checks['tau']['value']
u, n_exceed = 10.0, 355
print('tau from tails_checks.json (K)', tau)

# level exceeded once among the exceedances of the sample
print('u + tau ln 355 (K)', u + tau * math.log(n_exceed))

# limit y = tau / sigma^2, and the sigma at which the limit equals a given temperature
for s in (0.3, 0.4):
    print('limit tau/sigma^2 at sigma', s, '(K)', tau / s ** 2)
for y in (30.0, 20.0):
    print('sigma at which the limit is', y, 'K:', math.sqrt(tau / y))


def ratio(x):
    """Solve R ln R = x for R > 1 by bisection."""
    lo, hi = 1.0, 100.0
    for _ in range(200):
        mid = 0.5 * (lo + hi)
        if mid * math.log(mid) > x:
            hi = mid
        else:
            lo = mid
    return 0.5 * (lo + hi)


# modal ratio with tau_t set to the scale of the calculated values
for tt in (tau, 4.0):
    for s in (0.3,):
        for y in (10.0, 27.0):
            print('R at sigma', s, 'tau_t', round(tt, 3), 'y', y, ':', ratio(s * s * y / tt))
        for r in (1.2, 1.5):
            print('y at which R =', r, 'sigma', s, 'tau_t', round(tt, 3), ':', r * math.log(r) * tt / s ** 2)

# Under the two assumptions (true Tc exponential with scale tau_t, calculated y = t exp(eps),
# eps normal with zero mean and width sigma) the mean excess of y above a threshold is
#   m(v) = E[tau_t e^eps exp(-v e^-eps / tau_t)] / E[exp(-v e^-eps / tau_t)],
# evaluated here by Gauss-Hermite quadrature over eps. It exceeds tau_t at every threshold,
# so a calculated scale of 4.0 K implies a true scale below 4.0 K.
nodes, weights = np.polynomial.hermite_e.hermegauss(200)
for s in (0.15, 0.2, 0.3, 0.37, 0.4):
    for tt in (2.0, 3.0, 4.0):
        w = np.exp(s * nodes)
        row = []
        for v in (0.0, 2.0, 5.0, 10.0, 15.0):
            surv = np.exp(-v / (tt * w))
            row.append(float((weights * tt * w * surv).sum() / (weights * surv).sum()))
        print('mean excess of y over tau_t at sigma', s, 'tau_t', tt, 'thresholds 0, 2, 5, 10, 15 K:',
              [round(m / tt, 4) for m in row])