# C. elegans lifespan, treated vs control: Kaplan-Meier medians, log-rank tests pooled and stratified by trial, and
# analyses with plates or trials as the units. Python 3.13, NumPy, pandas, SciPy.
import numpy as np, pandas as pd
from scipy import stats
from analysis_functions import km_median, logrank
d = pd.read_csv("lifespan.csv")                      # one row per worm: trial, condition, plate, day (of death or censoring), died (1) or censored (0)

for t, g in d.groupby("trial"):
    mc, mt = [km_median(x.day.to_numpy(), x.died.to_numpy()) for _, x in g.groupby("condition")]
    print(f"Trial {t}: median lifespan control {mc:g} days, treated {mt:g} days ({g.died.eq(0).sum()} worms censored)")
c = logrank(d, np.zeros(len(d)))
print(f"A. All worms pooled ({(d.condition == 'treated').sum()} vs {(d.condition == 'control').sum()}), log-rank test: chi2(1) = {c:.2f}, p = {stats.chi2.sf(c, 1):.2g}")
c = logrank(d, d.trial)
print(f"B. Log-rank test stratified by trial: chi2(1) = {c:.2f}, p = {stats.chi2.sf(c, 1):.2g}")
pl = d.groupby(["trial", "condition", "plate"]).apply(lambda x: 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) - X.shape[1]
se = np.sqrt(res @ res / dfr * np.linalg.inv(X.T @ X)[1, 1]); tq = stats.t.ppf(0.975, dfr)
print(f"C. Plate medians as n ({len(pl) // 2} vs {len(pl) // 2}), log scale with trial as a block: treated / control {np.exp(beta[1]):.2f} "
      f"(95% CI {np.exp(beta[1] - tq * se):.2f} to {np.exp(beta[1] + tq * se):.2f}), p = {2 * stats.t.sf(abs(beta[1] / se), dfr):.2g}")
tr = d.groupby(["trial", "condition"]).apply(lambda x: km_median(x.day.to_numpy(), x.died.to_numpy()), include_groups=False).unstack()
r = np.log(tr.treated / tr.control); t = stats.ttest_1samp(r, 0); ci = t.confidence_interval()
print(f"D. Trial medians, log ratio per trial (paired, n = {len(r)}): treated / control {np.exp(r.mean()):.2f} (95% CI {np.exp(ci.low):.2f} to {np.exp(ci.high):.2f}), p = {t.pvalue:.2g}")