# Latency with a 60-s cut-off: responded = 1 if the mouse responded before the cut-off, 0 if the trial ended at 60 s (censored).
# Python 3.13, NumPy, SciPy (survival methods written out, so they can be checked against R's survival package).
import numpy as np, pandas as pd
from scipy import stats, optimize

d = pd.read_csv("latency.csv")
g = {k: d[d.group == k] for k in ["control", "treated"]}
for k, x in g.items():
    print(f"{k}: n = {len(x)}, reached the cut-off (censored) = {(x.responded == 0).sum()}, mean of recorded latencies = {x.latency_s.mean():.2f} s")

t = stats.ttest_ind(g["treated"].latency_s, g["control"].latency_s, equal_var=False)
print(f"Welch's t-test on recorded latencies (cut-off values treated as real): difference {g['treated'].latency_s.mean() - g['control'].latency_s.mean():.2f} s, p = {t.pvalue:.4f}")
u = stats.mannwhitneyu(g["treated"].latency_s, g["control"].latency_s, method="asymptotic", use_continuity=True)
print(f"Mann-Whitney test (cut-off values tied at the top): p = {u.pvalue:.4f}")

def km_median(x):                                     # Kaplan-Meier median, with R's convention: if the curve
    s = 1.0; times = np.sort(x.latency_s[x.responded == 1].unique())   # sits exactly at 0.5, take the midpoint to the next drop
    for i, tt in enumerate(times):
        at_risk = (x.latency_s >= tt).sum(); died = ((x.latency_s == tt) & (x.responded == 1)).sum()
        s *= 1 - died / at_risk
        if abs(s - 0.5) < 1e-12: return (tt + times[i + 1]) / 2 if i + 1 < len(times) else tt
        if s < 0.5: return tt
    return np.nan
print("Kaplan-Meier median latency: " + ", ".join(f"{k} {km_median(x):.1f} s" for k, x in g.items()))

times = np.sort(d.latency_s[d.responded == 1].unique()); o1 = e1 = v = 0.0   # log-rank test
for tt in times:
    r = (d.latency_s >= tt); r1 = (r & (d.group == "treated")).sum(); rn = r.sum()
    dd = ((d.latency_s == tt) & (d.responded == 1)); dn = dd.sum(); d1 = (dd & (d.group == "treated")).sum()
    o1 += d1; e1 += dn * r1 / rn
    v += dn * (r1 / rn) * (1 - r1 / rn) * (rn - dn) / (rn - 1) if rn > 1 else 0
chi = (o1 - e1) ** 2 / v
print(f"Log-rank test: chi-square = {chi:.2f}, p = {stats.chi2.sf(chi, 1):.4f}")

y = np.log(d.latency_s.to_numpy()); ev = d.responded.to_numpy(); xg = (d.group == "treated").to_numpy().astype(float)
def nll(p):                                           # log-normal accelerated failure time model, right-censored
    b0, b1, ls = p; s = np.exp(ls); z = (y - b0 - b1 * xg) / s
    return -(ev * (stats.norm.logpdf(z) - ls) + (1 - ev) * stats.norm.logsf(z)).sum()
fit = optimize.minimize(nll, [np.log(30), 0, np.log(0.6)], method="BFGS", options={"gtol": 1e-10})
H = np.array([optimize.approx_fprime(fit.x, lambda p: optimize.approx_fprime(p, nll, 1e-6)[i], 1e-6) for i in range(3)])
se = np.sqrt(np.diag(np.linalg.inv((H + H.T) / 2)))
b1, s1 = fit.x[1], se[1]
print(f"Log-normal survival (AFT) model: treated/control latency ratio {np.exp(b1):.2f} "
      f"(95% CI {np.exp(b1 - 1.96 * s1):.2f} to {np.exp(b1 + 1.96 * s1):.2f}), p = {2 * stats.norm.sf(abs(b1 / s1)):.4f}")