# Kaplan-Meier median and log-rank test used by analysis.py and simulation.py. Python 3.13, NumPy.
import numpy as np

def km_median(day, died):                            # Kaplan-Meier median, with the midpoint rule when S is exactly 0.5
    s, ts, ss = 1.0, [], []
    for t in np.unique(day):
        at_risk = np.sum(day >= t); deaths = np.sum((day == t) & (died == 1))
        if deaths: s *= 1 - deaths / at_risk; ts.append(t); ss.append(s)
    ts, ss = np.array(ts), np.array(ss)
    i = np.argmax(ss <= 0.5 + 1e-12)
    if not np.any(ss <= 0.5 + 1e-12): return np.nan
    return (ts[i] + ts[i + 1]) / 2 if abs(ss[i] - 0.5) < 1e-12 and i + 1 < len(ts) else float(ts[i])

def logrank(x, strata):                              # chi-squared (1 df), summed over strata as in survival::survdiff
    o = e = v = 0.0
    for _, g in x.groupby(strata):
        day, died, trt = g.day.to_numpy(), g.died.to_numpy(), (g.condition == "treated").to_numpy()
        for t in np.unique(day[died == 1]):
            r = day >= t; n, n1 = r.sum(), (r & trt).sum(); dd = ((day == t) & (died == 1)); m, m1 = dd.sum(), (dd & trt).sum()
            o += m1; e += m * n1 / n; v += m * (n1 / n) * (1 - n1 / n) * (n - m) / (n - 1) if n > 1 else 0
    return (o - e) ** 2 / v