"""
Net Expected-Lives Monte Carlo for the 1,500-hour rule.

Implements the LOCKED pre-registration in PRE-REGISTRATION.md / param-spec.json
(FRISA integrity-governed round, June 2026). Two competing hypotheses (H_A
monotonic experience-floor; H_B structure-and-selection-dominated) share every
cost-side parameter and differ ONLY in the benefit term lives_saved_per_decade.
A Gaussian copula couples the benefit to the rule's service-loss "bindingness"
(rho ~0.75) so the independent corners (max benefit + min cost, or vice versa)
are forbidden — the admissible region is the correlated diagonal.

Model (per decade):
  net_lives        = lives_saved_Part121  -  net_road_deaths
  net_road_deaths  = affected_communities * road_minus_air_deaths_per_comm_per_yr * 10yr
  affected_communities = communities_lost * rule_attributable_share * rule_binding_fraction
  road_minus_air_deaths_per_comm_per_yr =
        displaced_roundtrips * fraction_who_drive
        * (added_oneway_road_miles * 2)                       # round-trip pax-miles per driver
        * (road_rate - air_leg_rate) / 1e8                    # rates per 100M pax-miles

Positive net_lives => the rule saves lives on net; negative => it costs lives.

Outputs (all pre-specified, reported regardless of sign):
  - full net distribution quantiles 5/25/50/75/95, E[net], P(net<0), MC SE
    under EACH hypothesis (realized-evidence standard)
  - the precautionary/maximin reading (adopt the max-benefit hypothesis H_A)
  - one-at-a-time sensitivity of P(net<0)/E[net] to rule_attributable_share
    and the four linear cost dials
Run: python net_lives_mc.py
"""
import json
import numpy as np

N = 1_000_000
SEED = 20260613            # fixed for reproducibility (Date.now() unavailable; stamped from the date)
RHO = 0.75                 # benefit<->bindingness coupling (prereg rho_AB, mode 0.75)
rng = np.random.default_rng(SEED)


# ---- distribution helpers (parametrised from the locked priors) ----
def lognorm_from_median_p95(median, p95, size, u=None):
    """Lognormal from median and 95th percentile. u: optional uniform draws for copula."""
    sigma = np.log(p95 / median) / 1.6448536269514722
    mu = np.log(median)
    if u is None:
        return np.exp(rng.normal(mu, sigma, size))
    z = _norm_ppf(u)
    return np.exp(mu + sigma * z)


def tri(low, mode, high, size):
    return rng.triangular(low, mode, high, size)


def beta_mean(mean, conc, size, u=None):
    """Beta with given mean and concentration (a+b). u: optional uniforms for copula."""
    a = mean * conc
    b = (1 - mean) * conc
    if u is None:
        return rng.beta(a, b, size)
    # invert via the regularized incomplete beta (use scipy if present, else numeric)
    from math import lgamma
    # vectorised inverse-CDF through numpy's beta is not available; use acceptance via sorting trick:
    # draw independent beta then rank-match to u to impose the copula rank order.
    x = rng.beta(a, b, size)
    order = np.argsort(np.argsort(u))
    return np.sort(x)[order]


def _norm_ppf(u):
    # Acklam-free: use numpy via the inverse error function
    from numpy import sqrt
    import math
    # vectorised rational approximation (Acklam)
    a = [-3.969683028665376e+01, 2.209460984245205e+02, -2.759285104469687e+02,
         1.383577518672690e+02, -3.066479806614716e+01, 2.506628277459239e+00]
    b = [-5.447609879822406e+01, 1.615858368580409e+02, -1.556989798598866e+02,
         6.680131188771972e+01, -1.328068155288572e+01]
    c = [-7.784894002430293e-03, -3.223964580411365e-01, -2.400758277161838e+00,
         -2.549732539343734e+00, 4.374664141464968e+00, 2.938163982698783e+00]
    d = [7.784695709041462e-03, 3.224671290700398e-01, 2.445134137142996e+00,
         3.754408661907416e+00]
    u = np.clip(u, 1e-12, 1 - 1e-12)
    plow, phigh = 0.02425, 1 - 0.02425
    x = np.empty_like(u)
    lo = u < plow
    hi = u > phigh
    mid = ~(lo | hi)
    q = np.sqrt(-2 * np.log(u[lo]))
    x[lo] = (((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) / ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1)
    q = u[mid] - 0.5
    r = q*q
    x[mid] = (((((a[0]*r+a[1])*r+a[2])*r+a[3])*r+a[4])*r+a[5])*q / (((((b[0]*r+b[1])*r+b[2])*r+b[3])*r+b[4])*r+1)
    q = np.sqrt(-2 * np.log(1 - u[hi]))
    x[hi] = -(((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) / ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1)
    return x


def benefit_draw(hypothesis, u):
    """lives_saved_per_decade. u: uniform copula draws (coupled to bindingness)."""
    if hypothesis == 'H_A':       # median ~2, P95 ~8
        return lognorm_from_median_p95(2.0, 8.0, N, u=u)
    else:                         # H_B: median ~0.3, P95 ~3, P5≈0, negligible left tail
        return lognorm_from_median_p95(0.3, 3.0, N, u=u)


def sample_cost(overrides=None, u_bind=None):
    """Shared cost-side draws -> net_road_deaths per decade. Returns (net_road, dials)."""
    o = overrides or {}
    communities = o.get('communities', tri(14, 35, 50, N))
    # bindingness channel (coupled to benefit through u_bind): attribution & binding fraction
    attr = o.get('attr', beta_mean(0.25, 8.0, N, u=u_bind))            # Beta(2,6)
    binding = o.get('binding', beta_mean(0.40, 10.0, N, u=u_bind))     # Beta(4,6)
    displaced = o.get('displaced', lognorm_from_median_p95(3000, 10000, N))
    drive = o.get('drive', beta_mean(0.50, 10.0, N))                   # Beta(5,5)
    miles = o.get('miles', lognorm_from_median_p95(150, 250, N))
    road_rate = o.get('road_rate', lognorm_from_median_p95(0.9, 1.1, N))
    air_rate = lognorm_from_median_p95(0.03, 0.2, N)
    air_rate = np.minimum(air_rate, road_rate)                         # air leg never worse than road

    affected = communities * attr * binding
    road_minus_air_per_comm_yr = (displaced * drive * (miles * 2.0) * (road_rate - air_rate)) / 1e8
    net_road = affected * road_minus_air_per_comm_yr * 10.0            # per decade
    return net_road


def run(hypothesis, overrides=None):
    # shared latent u drives benefit and bindingness on the correlated diagonal
    zc = rng.normal(size=N)
    z_ben = RHO * zc + np.sqrt(1 - RHO**2) * rng.normal(size=N)
    u_ben = _norm_cdf(z_ben)
    u_bind = _norm_cdf(zc)
    benefit = benefit_draw(hypothesis, u_ben)
    net_road = sample_cost(overrides, u_bind=u_bind)
    net = benefit - net_road
    return net, benefit, net_road


def _norm_cdf(z):
    return 0.5 * (1 + _erf(z / np.sqrt(2)))


def _erf(x):
    # Abramowitz-Stegun 7.1.26
    t = 1 / (1 + 0.3275911 * np.abs(x))
    y = 1 - (((((1.061405429*t - 1.453152027)*t) + 1.421413741)*t - 0.284496736)*t + 0.254829592)*t*np.exp(-x*x)
    return np.sign(x) * y


def summarize(net):
    q = np.percentile(net, [5, 25, 50, 75, 95])
    p_neg = float(np.mean(net < 0))
    se = float(np.sqrt(p_neg * (1 - p_neg) / len(net)))
    return {
        'E_net': float(net.mean()),
        'q05': float(q[0]), 'q25': float(q[1]), 'median': float(q[2]),
        'q75': float(q[3]), 'q95': float(q[4]),
        'P_net_lt_0': p_neg, 'P_net_lt_0_SE': se,
    }


def main():
    out = {'N': N, 'seed': SEED, 'rho_benefit_cost': RHO, 'hypotheses': {}, 'sensitivity': {}}

    for h in ['H_A', 'H_B']:
        net, benefit, net_road = run(h)
        s = summarize(net)
        s['E_benefit'] = float(benefit.mean())
        s['E_net_road_cost'] = float(net_road.mean())
        out['hypotheses'][h] = s

    # precautionary / maximin: adopt the max-benefit hypothesis (H_A)
    net_prec, _, _ = run('H_A')
    out['precautionary_maximin_adopts_H_A'] = summarize(net_prec)

    # ---- pre-specified one-at-a-time sensitivity (under H_B, the shape-supported case) ----
    dials = {
        'rule_attributable_share': [('low', {'attr': np.full(N, 0.05)}),
                                    ('central', {'attr': np.full(N, 0.25)}),
                                    ('high', {'attr': np.full(N, 0.50)})],
        'communities_lost': [('low', {'communities': np.full(N, 14.0)}),
                             ('central', {'communities': np.full(N, 35.0)}),
                             ('high', {'communities': np.full(N, 50.0)})],
        'displaced_roundtrips': [('low', {'displaced': np.full(N, 3000.0)}),
                                 ('high', {'displaced': np.full(N, 10000.0)})],
        'fraction_who_drive': [('low', {'drive': np.full(N, 0.3)}),
                               ('high', {'drive': np.full(N, 0.7)})],
        'added_road_miles': [('low', {'miles': np.full(N, 75.0)}),
                             ('high', {'miles': np.full(N, 250.0)})],
    }
    for dial, settings in dials.items():
        out['sensitivity'][dial] = {}
        for label, ov in settings:
            net, _, _ = run('H_B', overrides=ov)
            out['sensitivity'][dial][label] = {
                'E_net': float(net.mean()), 'P_net_lt_0': float(np.mean(net < 0)),
            }

    with open('research/montecarlo/results.json', 'w') as f:
        json.dump(out, f, indent=2)

    # ---- plot ----
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(9, 5.2))
    for h, color in [('H_A', '#2C3A48'), ('H_B', '#BC6E2A')]:
        net, _, _ = run(h)
        ax.hist(np.clip(net, -15, 25), bins=200, density=True, alpha=0.55, color=color,
                label=f"{h}: E[net]={out['hypotheses'][h]['E_net']:+.2f}, P(net<0)={out['hypotheses'][h]['P_net_lt_0']:.0%}")
    ax.axvline(0, color='#111', lw=1.2, ls='--')
    ax.set_xlabel('Net lives per decade  (positive = rule saves lives; negative = rule costs lives)')
    ax.set_ylabel('density')
    ax.set_title('1,500-hour rule: net expected-lives Monte Carlo (1,000,000 draws)\nshared cost priors; hypotheses differ only in the benefit term; benefit–cost coupled at rho=0.75')
    ax.legend(loc='upper right', fontsize=9)
    ax.set_xlim(-15, 25)
    fig.tight_layout()
    fig.savefig('research/montecarlo/net_lives_distribution.png', dpi=130)

    print(json.dumps(out, indent=2))


if __name__ == '__main__':
    main()