"""
Supplementary File 1. Analysis code.
"The offer, not the odds: allocation of trial of labour after caesarean
 and the vaginal births it forgoes in 1020 women in central Kazakhstan"

Python 3.12; pandas, numpy, statsmodels, scikit-learn, scipy, matplotlib.
Input: de-identified cohort file, one row per woman, 1020 rows.
"""
import pandas as pd, numpy as np, statsmodels.api as sm
from sklearn.metrics import roc_auc_score
from scipy import stats

RNG = np.random.default_rng(20260827)
STRATA = [0, 1, 2, 3, 4]

# ---------------------------------------------------------------- load
COLS = {'Возраст_Age':'age','Группа_Group':'group','Предшествующие_ЕР_PVD':'pvd',
        'ИМТ_больше_30_BMI':'bmi30','Интервал_между_родами_мес':'idi',
        'Интервал_больше_18_мес':'idi18','Нерекуррентные_показания':'nonrec',
        'ГСД_GDM':'gdm','ГРП_HDP':'hdp','Кровотечение_PPH':'pph',
        'Гемотрансфузия':'transf','ПИТ_новорожденных_NICU':'nicu',
        'Разрыв_матки':'rupture','Несостоятельность_рубца':'dehisc'}

def load(path):
    df = pd.read_excel(path)
    v = df[df['Источник_данных'] == 'VBAC Study'].rename(columns=COLS)
    v['tolac'] = (v.group != 'PRCD').astype(int)
    v['vbac']  = (v.group == 'VBAC').astype(int)
    v['comp']  = ((v.pph + v.transf + v.rupture + v.dehisc) > 0).astype(int)
    # unweighted, guideline-derived favourability count (0-4); no weights fitted here
    v['fav'] = v.pvd + v.nonrec + v.idi18 + (1 - v.bmi30)
    return v

# ------------------------------------------- offer model vs outcome model
ANTENATAL = ['pvd', 'nonrec', 'idi18', 'bmi30', 'age']
# Variables generated AFTER the mode-of-birth decision (gestational age at
# delivery, onset of labour, Bishop score) are deliberately excluded: planned
# caesareans are scheduled at 39 weeks, so those are consequences of the plan.

def two_models(v):
    Xo = sm.add_constant(v[ANTENATAL].astype(float))
    m_offer = sm.Logit(v.tolac, Xo).fit(disp=0)
    t = v[v.tolac == 1]
    Xs = sm.add_constant(t[ANTENATAL].astype(float))
    m_succ = sm.Logit(t.vbac, Xs).fit(disp=0)
    out = pd.DataFrame({
        'aOR_offer': np.exp(m_offer.params),
        'lo_offer' : np.exp(m_offer.conf_int()[0]),
        'hi_offer' : np.exp(m_offer.conf_int()[1]),
        'p_offer'  : m_offer.pvalues,
        'aOR_vbac' : np.exp(m_succ.params),
        'lo_vbac'  : np.exp(m_succ.conf_int()[0]),
        'hi_vbac'  : np.exp(m_succ.conf_int()[1]),
        'p_vbac'   : m_succ.pvalues})
    c_offer = roc_auc_score(v.tolac, m_offer.predict(Xo))
    c_succ  = roc_auc_score(t.vbac,  m_succ.predict(Xs))
    return out, c_offer, c_succ

# ------------------------------------------------------ allocation analysis
def _clip(p): return min(max(p, 1e-3), 1 - 1e-3)

def stratum_rates(v):
    return {s: v[(v.fav == s) & (v.tolac == 1)].vbac.mean() for s in STRATA}

def stratum_sizes(v):
    return {s: int((v.fav == s).sum()) for s in STRATA}

def fixed_volume(n_trials, rates, sizes, delta=1.0):
    """Allocate n_trials best-first; delta discounts the odds of vaginal birth
    among women who were never offered labour (1.0 = exchangeable)."""
    births, left = 0.0, n_trials
    for s in sorted(STRATA, reverse=True):
        if left <= 0: break
        k = min(left, sizes[s])
        p = _clip(rates[s]); odds = p / (1 - p) * delta
        births += k * odds / (1 + odds); left -= k
    return births

def fixed_failures(budget, rates, sizes, delta=1.0):
    """Allocate best-first until the expected number of FAILED trials hits
    `budget`. Returns (trials, expected vaginal births)."""
    trials = births = failures = 0.0
    for s in sorted(STRATA, reverse=True):
        p = _clip(rates[s]); odds = p / (1 - p) * delta
        pa = odds / (1 + odds); fr = 1 - pa
        if failures + sizes[s] * fr <= budget:
            trials += sizes[s]; births += sizes[s] * pa; failures += sizes[s] * fr
        else:
            k = min((budget - failures) / fr, sizes[s]) if fr > 0 else sizes[s]
            trials += k; births += k * pa
            break
    return trials, births

def bootstrap(v, n_trials, fail_budget, B=2000):
    sizes = stratum_sizes(v); A, Bv, Bt = [], [], []
    for _ in range(B):
        idx = []
        for s in STRATA:
            g = v[(v.fav == s) & (v.tolac == 1)]
            idx.extend(RNG.choice(g.index, len(g), replace=True))
        bs = v.loc[idx]
        r = {s: bs[bs.fav == s].vbac.mean() for s in STRATA}
        A.append(fixed_volume(n_trials, r, sizes))
        t, b = fixed_failures(fail_budget, r, sizes); Bv.append(b); Bt.append(t)
    ci = lambda a: (np.percentile(a, 2.5), np.percentile(a, 97.5))
    return ci(A), ci(Bv), ci(Bt)

# ------------------------------- bootstrap comparison of model weights
def weight_divergence(v, B=4000, seed=7):
    """Bootstrap the difference in log odds ratios between the offer model and
    the vaginal-birth model, covariate by covariate, plus a joint Wald test."""
    import numpy as np
    from scipy.stats import chi2 as chi2dist
    rng = np.random.default_rng(seed)
    def fit(d):
        Xo = sm.add_constant(d[ANTENATAL].astype(float))
        mo = sm.Logit(d.tolac, Xo).fit(disp=0)
        t = d[d.tolac == 1]
        Xs = sm.add_constant(t[ANTENATAL].astype(float))
        ms = sm.Logit(t.vbac, Xs).fit(disp=0)
        return mo.params[ANTENATAL].values, ms.params[ANTENATAL].values
    o, s = fit(v); diff = o - s
    boots = []
    for _ in range(B):
        d = v.loc[rng.choice(v.index, len(v), replace=True)]
        try:
            a, b = fit(d); boots.append(a - b)
        except Exception:
            pass
    boots = np.array(boots)
    out = {}
    for j, c in enumerate(ANTENATAL):
        lo, hi = np.percentile(boots[:, j], [2.5, 97.5])
        p = 2 * min((boots[:, j] <= 0).mean(), (boots[:, j] >= 0).mean())
        out[c] = (diff[j], lo, hi, p)
    S = np.cov(boots.T)
    stat = diff @ np.linalg.pinv(S) @ diff
    return out, stat, chi2dist.sf(stat, len(ANTENATAL))

# ------------------------------------------------------------- E-values
def evalue(point, limit):
    """E-value (VanderWeele & Ding) for an odds ratio with a common outcome;
    OR is converted to an approximate risk ratio by the square-root rule."""
    import numpy as np
    def e(x):
        rr = np.sqrt(x)
        if rr < 1: rr = 1 / rr
        return rr + np.sqrt(rr * (rr - 1))
    return e(point), e(limit)

# ----------------------------------------------------------------- report
if __name__ == '__main__':
    v = load('VBAC_data.xlsx')
    rates, sizes = stratum_rates(v), stratum_sizes(v)
    obs_trials = int(v.tolac.sum()); obs_births = int(v.vbac.sum())
    obs_fail = obs_trials - obs_births

    print(f'Cohort {len(v)}; trials {obs_trials} ({100*v.tolac.mean():.1f}%); '
          f'vaginal births {obs_births}; failed {obs_fail}; '
          f'caesarean rate {100*(1-obs_births/len(v)):.1f}%')

    tbl, c_off, c_suc = two_models(v)
    print('\nTable 2 — offer vs outcome, identical covariates')
    print(tbl.round(3).to_string())
    print(f'C-statistic: offer {c_off:.3f}, vaginal birth {c_suc:.3f}')

    print('\nTable 5 — allocation')
    print(f'  observed                 : {obs_trials} trials -> {obs_births} births')
    print(f'  same volume, reordered   : {obs_trials} trials -> '
          f'{fixed_volume(obs_trials, rates, sizes):.0f} births')
    t, b = fixed_failures(obs_fail, rates, sizes)
    print(f'  same failed trials ({obs_fail})  : {t:.0f} trials -> {b:.0f} births; '
          f'caesarean rate {100*(1-b/len(v)):.1f}%')
    ciA, ciB, ciT = bootstrap(v, obs_trials, obs_fail)
    print(f'  bootstrap 95% CI: same volume {ciA[0]:.0f}-{ciA[1]:.0f}; '
          f'same failures {ciB[0]:.0f}-{ciB[1]:.0f} (trials {ciT[0]:.0f}-{ciT[1]:.0f})')

    print('\nBootstrap comparison of model weights (offer vs vaginal birth)')
    div, stat, pj = weight_divergence(v)
    for c, (d, lo, hi, p) in div.items():
        print(f'  {c:10} diff logOR {d:+.3f} [{lo:+.3f}, {hi:+.3f}]  p={p:.4f}')
    print(f'  joint Wald chi2({len(ANTENATAL)}) = {stat:.2f}, p = {pj:.4f}')

    print('\nE-values for composite maternal morbidity')
    print('  vaginal birth vs PRCD  aOR 0.56 (0.32-0.96): E = %.2f, at CI limit %.2f' % evalue(0.56, 0.96))
    print('  failed trial vs PRCD   aOR 1.35 (0.74-2.46): E = %.2f, at upper limit %.2f' % evalue(1.35, 2.46))

    print('\nTable 6 — sensitivity to unmeasured selection')
    for d in [1.0, 0.8, 0.6, 0.5, 0.4, 0.3]:
        _, b = fixed_failures(obs_fail, rates, sizes, d)
        print(f'  delta={d:.2f}  same volume {fixed_volume(obs_trials, rates, sizes, d):.0f}  '
              f'same failures {b:.0f}')
