# 5,000 simulated studies with the colonies.csv design (5 experiments, 4 doses, 3 dishes per dose and condition).
# Python 3.13, NumPy, SciPy. Seed fixed before running.
import numpy as np
from scipy import stats
from model import experiment, DOSES, SEEDED, DISHES
rng = np.random.default_rng(20261112); R = 5_000

def poisson_dev(y, off, X):
    b = np.linalg.lstsq(X, np.log(y + 0.5) - off, rcond=None)[0]
    for _ in range(100):
        mu = np.exp(off + X @ b); step = np.linalg.solve(X.T @ (X * mu[:, None]), X.T @ (y - mu)); b += step
        if np.max(np.abs(step)) < 1e-10: break
    mu = np.exp(off + X @ b); return 2 * np.sum(np.where(y > 0, y * np.log(np.where(y > 0, y, 1) / mu), 0) - (y - mu)), np.sum((y - mu) ** 2 / mu)
def one(alpha_ratio):
    rows = np.array([[e, r[0] == "sensitiser", r[1], r[3], r[4]] for e in range(5) for r in experiment(rng, alpha_ratio=alpha_ratio)], float)
    exp_, s, D, seeded, y = rows.T
    lr_sf2 = []
    for e in range(5):
        sf2 = []
        for c in (0, 1):
            k = (exp_ == e) & (s == c); m = np.array([(y[k & (D == dd)] / seeded[k & (D == dd)]).mean() for dd in DOSES])
            sf = m[1:] / m[0]; A = np.column_stack([DOSES[1:], DOSES[1:] ** 2])
            (a, b), *_ = np.linalg.lstsq(A, -np.log(np.maximum(sf, 1e-9)), rcond=None); sf2.append(np.exp(-(2 * a + 4 * b)))
        lr_sf2.append(np.log(sf2[1] / sf2[0]))
    p_exp = stats.ttest_1samp(lr_sf2, 0).pvalue
    off = np.log(seeded); P = (exp_[:, None] * 2 + s[:, None] == np.arange(10)[None, :]).astype(float)   # PE per experiment and condition
    X0 = np.column_stack([P, D, D ** 2]); X1 = np.column_stack([X0, s * D, s * D ** 2])
    (d0, _), (d1, p1) = poisson_dev(y, off, X0), poisson_dev(y, off, X1); lr = d0 - d1; dfr = len(y) - X1.shape[1]
    return p_exp < 0.05, stats.chi2.sf(lr, 2) < 0.05, stats.f.sf((lr / 2) / (p1 / dfr), 2, dfr) < 0.05
for lab, ar in (("Sensitiser has no effect", 1.0), ("Sensitiser raises alpha by 30%", 1.3)):
    r = np.array([one(ar) for _ in range(R)])
    print(f"{lab}: p<0.05 with SF2 per experiment (paired) {r[:, 0].mean():.1%}; all dishes pooled, Poisson {r[:, 1].mean():.1%}; "
          f"all dishes pooled, quasi-Poisson {r[:, 2].mean():.1%}")