# 10,000 simulated experiments per scenario: 12 control and 12 treated mice, one trial each, cut-off 60 s.
# Same model as the worked example (model.py). Python 3.13, NumPy, SciPy. Seed fixed before running.
import numpy as np
from scipy import stats, optimize
from model import simulate
rng = np.random.default_rng(20261111); R = 10_000; n = 12

def logrank_p(t, e, grp):
    times = np.unique(t[e == 1]); o1 = e1 = v = 0.0
    for tt in times:
        r = t >= tt; r1 = (r & grp).sum(); rn = r.sum(); dd = (t == tt) & (e == 1); dn = dd.sum()
        o1 += (dd & grp).sum(); e1 += dn * r1 / rn
        if rn > 1: v += dn * (r1 / rn) * (1 - r1 / rn) * (rn - dn) / (rn - 1)
    return stats.chi2.sf((o1 - e1) ** 2 / v, 1)

def aft(t, e, grp):                                   # log-normal AFT; returns log time ratio and its p-value
    y = np.log(t); x = grp.astype(float)
    def nll(p):
        z = (y - p[0] - p[1] * x) / np.exp(p[2])
        return -(e * (stats.norm.logpdf(z) - p[2]) + (1 - e) * stats.norm.logsf(z)).sum()
    f = optimize.minimize(nll, [np.log(30), 0, np.log(0.6)], method="BFGS", options={"gtol": 1e-8})
    E = np.eye(3) * 1e-4; p0 = f.x                    # observed information (numerical Hessian), as in analysis.py
    H = np.array([[(nll(p0 + a + b) - nll(p0 + a - b) - nll(p0 - a + b) + nll(p0 - a - b)) / 4e-8 for b in E] for a in E])
    return f.x[1], 2 * stats.norm.sf(abs(f.x[1] / np.sqrt(np.linalg.inv(H)[1, 1])))

for label, ratio in (("No treatment effect", 1.0), ("Treatment multiplies latency by 1.6", 1.6)):
    sig = np.zeros(4); cens = 0; diffs = []; ratios = []
    for _ in range(R):
        tc, ec = simulate(rng, n, 25); tt, et = simulate(rng, n, 25 * ratio)
        t = np.r_[tc, tt]; e = np.r_[ec, et]; grp = np.r_[np.zeros(n, bool), np.ones(n, bool)]
        sig[0] += stats.ttest_ind(tt, tc, equal_var=False).pvalue < 0.05
        sig[1] += stats.mannwhitneyu(tt, tc, method="asymptotic").pvalue < 0.05
        sig[2] += logrank_p(t, e, grp) < 0.05
        b, p = aft(t, e, grp); sig[3] += p < 0.05; ratios.append(np.exp(b))
        diffs.append(tt.mean() - tc.mean()); cens += (1 - et).sum()
    print(f"{label}: p<0.05 with Welch's t-test on recorded values {sig[0] / R:.1%}, Mann-Whitney {sig[1] / R:.1%}, "
          f"log-rank {sig[2] / R:.1%}, log-normal AFT {sig[3] / R:.1%}; treated mice censored {cens / (R * n):.1%}; "
          f"median estimated latency ratio (AFT) {np.median(ratios):.2f}; mean recorded difference {np.mean(diffs):.1f} s")
# true mean difference without the cut-off, for the second scenario
m = lambda med: med * np.exp(0.6 ** 2 / 2)
print(f"True mean difference without a cut-off in the second scenario: {m(40) - m(25):.1f} s")