# Stage 2: preregistered analysis (prereg:eba635114e0f30e35505c4c97c11053dffd29f743b0933d4cf116ca1ed78c22c, plan/prereg.md).
# Sharp RD in single-year maternal age; local linear, triangular kernel, cells weighted by births; honest CI
# (bounded second derivative M): half-width = worst-case bias bound + 1.96 SE. Reads data/cells_*.csv, writes results/.
import json, os, numpy as np, pandas as pd
from scipy.stats import norm
os.makedirs('results', exist_ok=True)
D = 'data'
main = pd.read_csv(f'{D}/cells_main.csv'); sub = pd.read_csv(f'{D}/cells_sub.csv'); fd = pd.read_csv(f'{D}/cells_fd.csv')
M0 = 0.001                     # 0.1 pp per yr^2 (prereg), in probability units
Z = 1.959963984540054
R = lambda v: None if v is None else float(round(float(v), 8))

def weights(ages, n, cut, h, deg=1):
    x = ages - cut + 0.5                                       # bin centres relative to the cutoff (age a covers [a,a+1))
    k = np.maximum(0.0, 1 - np.abs(x) / h)
    w = np.zeros(len(ages))
    for side, sg in ((x < 0, -1.0), (x > 0, 1.0)):
        X = np.vander(x[side], deg + 1, increasing=True); W = k[side] * n[side]
        L = np.linalg.solve(X.T @ (X * W[:, None]), (X * W[:, None]).T)[0]
        w[side] = sg * L
    return x, w

def honest(est, se, bias, sesoi=None):
    hw = bias + Z * se; lo, hi = est - hw, est + hw
    t = (abs(est) - bias) / se if se > 0 else np.inf
    p = 1.0 if t <= 0 else float(2 * norm.sf(t))
    return dict(estimate=R(est), se=R(se), bias_bound=R(bias), ci_low=R(lo), ci_high=R(hi), p_honest=R(p))

def rd_series(ages, y, var, n, cut, h, M, deg=1, nw=None):
    # y: cell means by age; var: per-observation variance by age; n: counts (precision); nw: weighting counts (default n)
    s = (ages >= cut - h) & (ages < cut + h)
    a, yy, vv, nn = ages[s], y[s], var[s], n[s]
    x, w = weights(a, nn if nw is None else nw[s], cut, h, deg)
    est = float(w @ yy); se = float(np.sqrt(np.sum(w ** 2 * vv / nn))); bias = M / 2 * float(np.sum(np.abs(w) * x ** 2))
    return honest(est, se, bias)

def agg(df, by='age'):
    return df.groupby(by).sum(numeric_only=True)

def binary(cells, kcol, ncol, cut, h=3, M=M0, deg=1):
    g = agg(cells); ages = g.index.values.astype(float); n = g[ncol].values.astype(float); p = g[kcol].values / n
    return rd_series(ages, p, p * (1 - p), n, cut, h, M, deg)

def classify(r, sesoi, falsification_pass=True):
    if r['ci_low'] > 0 and r['estimate'] >= sesoi / 2 and falsification_pass: return 'confirm'
    if r['ci_low'] > -sesoi and r['ci_high'] < sesoi: return 'refute'
    return 'inconclusive'

def holm(ps):
    order = np.argsort(ps); m = len(ps); adj = np.empty(m); run = 0.0
    for i, j in enumerate(order): run = max(run, min(1.0, (m - i) * ps[j])); adj[j] = run
    return adj

def emp_M(series, cuts=(35, 40)):
    # deviation D3: curvature bound for non-probability covariates = max |second difference| of the variable's
    # own single-age series (ages 28-46), excluding second differences that straddle a cutoff
    a = series.index.values; v = series.values; out = 0.0
    for i in range(1, len(a) - 1):
        if not (28 <= a[i] <= 46): continue
        if any(a[i - 1] < c <= a[i + 1] for c in cuts): continue
        out = max(out, abs(v[i + 1] - 2 * v[i] + v[i - 1]))
    return float(out)

nul = main[main.nullip == 1]; allp = main
res = dict(prereg='prereg:eba635114e0f30e35505c4c97c11053dffd29f743b0933d4cf116ca1ed78c22c', M_prereg_pp=0.1)

# ---- placebo cutoffs and the prespecified M calibration rule (deviation D1 operationalizes it) ----
plc = {str(c): binary(nul, 'ind_k', 'ind_n', c) for c in (31, 32, 33, 37, 38)}
nrej = sum(r['ci_low'] > 0 or r['ci_high'] < 0 for r in plc.values())
M_dec = M0
if nrej >= 2:
    for m in (0.0015, 0.002, 0.003, 0.005, 0.01, 0.02):
        if sum((lambda r: r['ci_low'] > 0 or r['ci_high'] < 0)(binary(nul, 'ind_k', 'ind_n', c, M=m)) for c in (31, 32, 33, 37, 38)) <= 1:
            M_dec = m; break
    else: M_dec = 0.02
res['placebo'] = dict(cutoffs=plc, n_excluding_zero=nrej, M_decision_pp=R(100 * M_dec),
                      placebo_at_M_decision={str(c): binary(nul, 'ind_k', 'ind_n', c, M=M_dec) for c in (31, 32, 33, 37, 38)})

# ---- covariate falsification (pre-treatment), density, missingness ----
def cov_block(cut):
    out = {}
    for name, k, n in (('ba_plus', 'ba_k', 'ba_n'), ('medicaid', 'med_k', 'pay_n'), ('infertility_tx', 'inf_k', 'inf_n'),
                       ('prepreg_diabetes', 'pdiab_k', 'pdiab_n'), ('chronic_htn', 'phype_k', 'phype_n'),
                       ('nh_white', 'nhw_k', 're_n'), ('nh_black', 'nhb_k', 're_n'), ('hispanic', 'hisp_k', 're_n'), ('nh_asian', 'nha_k', 're_n')):
        out[name] = binary(nul, k, n, cut)
    g = agg(allp); a = g.index.values.astype(float)
    ns = g['n'].values.astype(float); pn = agg(nul)['n'].reindex(g.index).values / ns
    out['nulliparity_share_allparity'] = rd_series(a, pn, pn * (1 - pn), ns, cut, 3, M0)
    gb = agg(nul); nb = gb['bmi_n'].values.astype(float); mb = gb['bmi_s10'].values / 10 / nb
    vb = gb['bmi_ss100'].values.astype(float) / 100 / nb - mb ** 2
    Mb = emp_M(pd.Series(mb, gb.index)); out['bmi'] = {**rd_series(gb.index.values.astype(float), mb, vb, nb, cut, 3, Mb), 'M_used': R(Mb)}
    ps = np.array([out[v]['p_honest'] for v in out]); adj = holm(ps)
    for v, q in zip(list(out), adj): out[v]['p_holm'] = R(q)
    return out, bool(np.all(adj >= 0.05))

def density(cut, sample):
    g = agg(sample); a = g.index.values.astype(float); n = g['n'].values.astype(float); y = np.log(n)
    Md = emp_M(pd.Series(y, g.index))
    return {**rd_series(a, y, np.ones_like(n), n, cut, 3, Md, nw=np.ones_like(n)), 'M_used': R(Md)}

fals = {}
for cut in (35, 40):
    cv, ok = cov_block(cut)
    g = agg(nul); nn = g['n'].values.astype(float); pm = g['ind_miss'].values / nn
    fals[str(cut)] = dict(covariates=cv, passes_holm=ok,
                          density_log_births_nulliparous=density(cut, nul), density_log_births_allparity=density(cut, allp),
                          induction_missing_share=rd_series(g.index.values.astype(float), pm, pm * (1 - pm), nn, cut, 3, M0))
res['falsification'] = fals

# ---- primary H1 and secondary H2 ----
prim = binary(nul, 'ind_k', 'ind_n', 35)
prim_dec = binary(nul, 'ind_k', 'ind_n', 35, M=M_dec)
res['primary_H1'] = dict(outcome='labor induction', sample='singleton nulliparous 37-44 wk', cutoff=35, h=3,
                         at_M_prereg=prim, verdict_at_M_prereg=classify(prim, 0.01, fals['35']['passes_holm']),
                         M_decision_pp=R(100 * M_dec), at_M_decision=prim_dec,
                         verdict=classify(prim_dec, 0.01, fals['35']['passes_holm']), sesoi=0.01, falsification_pass=fals['35']['passes_holm'])
h2 = binary(nul, 'ind_k', 'ind_n', 40); h2d = binary(nul, 'ind_k', 'ind_n', 40, M=M_dec)
res['secondary_H2'] = dict(cutoff=40, at_M_prereg=h2, verdict_at_M_prereg=classify(h2, 0.02, fals['40']['passes_holm']),
                           at_M_decision=h2d, verdict=classify(h2d, 0.02, fals['40']['passes_holm']), sesoi=0.02)

# ---- key secondary benefit outcomes, Holm within cutoff ----
BEN = (('cesarean', 'ces_k', 'ces_n'), ('delivery_39w', 'w39_k', 'n'), ('nicu', 'nicu_k', 'nicu_n'),
       ('apgar5_lt7', 'apg_k', 'apg_n'), ('maternal_morbidity', 'mmc_k', 'mmc_n'))
ben = {}
for cut in (35, 40):
    b = {o: binary(nul, k, n, cut) for o, k, n in BEN}
    adj = holm(np.array([b[o]['p_honest'] for o in b]))
    for o, q in zip(b, adj): b[o]['p_holm'] = R(q)
    ben[str(cut)] = b
res['key_secondary_benefit'] = ben

# ---- other secondaries: stillbirth, all-parity, early term ----
def stillbirth(cut, nullip_only):
    f = fd[fd.nullip == 1] if nullip_only else fd
    lb = main[(main.yr >= 2018) & ((main.nullip == 1) if nullip_only else True)]
    g = pd.concat([agg(f)['fd'], agg(lb)['n']], axis=1).fillna(0); n = (g['fd'] + g['n']).values.astype(float); p = g['fd'].values / n
    r = rd_series(g.index.values.astype(float), p, p * (1 - p), n, cut, 3, M0)
    r['no_benefit_vs_jama_hf_minus_0.39pp'] = bool(r['ci_low'] > -0.0039)
    return r
res['secondary_other'] = {str(cut): dict(
    stillbirth_nulliparous=stillbirth(cut, True), stillbirth_allparity=stillbirth(cut, False),
    induction_allparity=binary(allp, 'ind_k', 'ind_n', cut), cesarean_allparity=binary(allp, 'ces_k', 'ces_n', cut),
    early_term_37_38w=binary(nul, 'early_k', 'n', cut)) for cut in (35, 40)}

# ---- sensitivities for the primary ----
sens = {f'h{h}': binary(nul, 'ind_k', 'ind_n', 35, h=h) for h in (2, 4, 5)}
sens['quadratic_h5'] = binary(nul, 'ind_k', 'ind_n', 35, h=5, deg=2)
sens['M0.05pp'] = binary(nul, 'ind_k', 'ind_n', 35, M=0.0005); sens['M0.2pp'] = binary(nul, 'ind_k', 'ind_n', 35, M=0.002)
def year_fe(cut=35, h=3, M=M0):
    c = nul[(nul.age >= cut - h) & (nul.age < cut + h)].copy()
    x = c.age.values - cut + 0.5; d = (x > 0).astype(float); k = np.maximum(0, 1 - np.abs(x) / h)
    yrs = sorted(c.yr.unique()); Y = np.c_[[(c.yr.values == y).astype(float) for y in yrs]].T
    X = np.c_[d, x, d * x, Y]; n = c.ind_n.values.astype(float); W = k * n; p = c.ind_k.values / n
    L = np.linalg.solve(X.T @ (X * W[:, None]), (X * W[:, None]).T)[0]
    return honest(float(L @ p), float(np.sqrt(np.sum(L ** 2 * p * (1 - p) / n))), M / 2 * float(np.sum(np.abs(L) * x ** 2)))
sens['year_fixed_effects'] = year_fe()
for c2 in (35, 40): pass
res['sensitivity_primary'] = sens
res['sensitivity_H2'] = {**{f'h{h}': binary(nul, 'ind_k', 'ind_n', 40, h=h) for h in (2, 4, 5)},
                         'M0.05pp': binary(nul, 'ind_k', 'ind_n', 40, M=0.0005), 'M0.2pp': binary(nul, 'ind_k', 'ind_n', 40, M=0.002)}

# ---- subgroups (primary outcome, cutoff 35) ----
sg = {}
for (dim, lvl), d in sub.groupby(['dim', 'lvl']):
    sg.setdefault(dim, {})[lvl] = binary(d, 'ind_k', 'ind_n', 35)
res['subgroups_35'] = sg

# ---- descriptive cell rates (pooled), for the figures ----
g = agg(nul)
rates = pd.DataFrame({'age': g.index, 'n': g['n'], 'induction': g.ind_k / g.ind_n, 'cesarean': g.ces_k / g.ces_n,
                      'nicu': g.nicu_k / g.nicu_n, 'apgar5_lt7': g.apg_k / g.apg_n, 'maternal_morbidity': g.mmc_k / g.mmc_n,
                      'delivery_39w': g.w39_k / g.n}).round(8)
rates.to_csv('results/rates_by_age_nulliparous.csv', index=False)
json.dump(res, open('results/results.json', 'w'), indent=1, sort_keys=True)
s = res['primary_H1']; print('PRIMARY', s['verdict'], s['at_M_decision'], 'M_dec_pp', s['M_decision_pp'])

# ---- declared results for claims.json (percentage points, 2 decimals) ----
pp = lambda v: float(round(100 * v, 2))
def blk(r): return dict(estimate_pp=pp(r['estimate']), se_pp=pp(r['se']), bias_bound_pp=pp(r['bias_bound']), ci_low_pp=pp(r['ci_low']), ci_high_pp=pp(r['ci_high']))
f35 = fals['35']['covariates']
nl = res['primary_H1']
json.dump(dict(**blk(nl['at_M_decision']), verdict=nl['verdict'], falsification_pass=nl['falsification_pass'],
               n_nulliparous_age32_37=int(agg(nul).loc[32:37, 'ind_n'].sum()),
               nulliparity_share_jump_pp=pp(f35['nulliparity_share_allparity']['estimate']),
               nulliparity_share_p_holm=float(round(f35['nulliparity_share_allparity']['p_holm'], 3)),
               placebo_rejections=int(nrej), M_decision_pp=float(nl['M_decision_pp'])), open('results/R1.json', 'w'), indent=1, sort_keys=True)
h2r = res['secondary_H2']
json.dump(dict(**blk(h2r['at_M_decision']), verdict=h2r['verdict'], falsification_pass=fals['40']['passes_holm'],
               n_nulliparous_age37_42=int(agg(nul).loc[37:42, 'ind_n'].sum())), open('results/R2.json', 'w'), indent=1, sort_keys=True)
json.dump({'age'+c: {o: blk(ben[c][o]) for o in ben[c]} for c in ben}, open('results/R3.json', 'w'), indent=1, sort_keys=True)
json.dump({'age'+c: {o: blk(res['secondary_other'][c][o]) for o in ('stillbirth_nulliparous', 'induction_allparity')} for c in ('35', '40')},
          open('results/R4.json', 'w'), indent=1, sort_keys=True)
json.dump({k: blk(v) for k, v in sens.items()}, open('results/R5.json', 'w'), indent=1, sort_keys=True)