# How often each analysis gives p < 0.05: simulated studies of 3 trials x 3 plates x 35 worms per condition, with no
# treatment effect and with a 12% longer lifespan. Python 3.13, NumPy, pandas, SciPy. Seed 20261112.
import numpy as np, pandas as pd
from scipy import stats
from model import study
import analysis_functions as af
rng = np.random.default_rng(20261112); R = 2_000

def one(ratio):
    d = pd.DataFrame(study(rng, ratio), columns=["trial", "condition", "plate", "day", "died"])
    pa = stats.chi2.sf(af.logrank(d, np.zeros(len(d))), 1); pb = stats.chi2.sf(af.logrank(d, d.trial), 1)
    pl = d.groupby(["trial", "condition", "plate"]).apply(lambda x: af.km_median(x.day.to_numpy(), x.died.to_numpy()), include_groups=False).rename("med").reset_index()
    X = np.column_stack([np.ones(len(pl)), pl.condition == "treated", pl.trial == 2, pl.trial == 3]).astype(float); y = np.log(pl.med.to_numpy())
    beta, *_ = np.linalg.lstsq(X, y, rcond=None); res = y - X @ beta; dfr = len(y) - 4
    pc = 2 * stats.t.sf(abs(beta[1] / np.sqrt(res @ res / dfr * np.linalg.inv(X.T @ X)[1, 1])), dfr)
    tr = d.groupby(["trial", "condition"]).apply(lambda x: af.km_median(x.day.to_numpy(), x.died.to_numpy()), include_groups=False).unstack()
    pd_ = stats.ttest_1samp(np.log(tr.treated / tr.control), 0).pvalue
    return np.array([pa, pb, pc, pd_]) < 0.05
for lab, ratio in (("No effect", 1.0), ("Treatment extends lifespan by 12%", 1.12)):
    r = np.array([one(ratio) for _ in range(R)]).mean(axis=0)
    print(f"{lab} ({R:,} studies): A log-rank, all worms pooled {r[0]:.1%}; B log-rank stratified by trial {r[1]:.1%}; "
          f"C plate medians, trial as block {r[2]:.1%}; D trial medians, paired {r[3]:.1%}")