# How often each analysis gives p < 0.05: 5,000 simulated studies of 10 + 10 mice, with no knockout effect and with a
# knockout that extinguishes more slowly (decline per tone halved on the logit scale). Python 3.13, NumPy, SciPy. Seed 20261112.
import numpy as np
from scipy import stats
from model import study, TRIALS, SECONDS
rng = np.random.default_rng(20261112); R = 5_000; x = np.arange(1, TRIALS + 1); H = np.eye(TRIALS) - 1 / TRIALS

def one(ko_decline):
    s = study(rng, ko_decline=ko_decline); wt = 100 * np.array([m[2] for m in s if m[0] == "wild-type"]) / SECONDS
    ko = 100 * np.array([m[2] for m in s if m[0] == "knockout"]) / SECONDS; n, k = wt.shape
    pa = min(stats.ttest_ind(ko[:, j], wt[:, j]).pvalue for j in range(k))
    pb = stats.ttest_ind(ko.ravel(), wt.ravel()).pvalue
    pc = stats.ttest_ind(ko.mean(axis=1), wt.mean(axis=1), equal_var=False).pvalue
    pd_ = stats.ttest_ind(np.polyfit(x, ko.T, 1)[0], np.polyfit(x, wt.T, 1)[0], equal_var=False).pvalue
    Y = np.vstack([wt, ko]); gm = Y.mean(); mg = np.vstack([np.tile(wt.mean(0), (n, 1)), np.tile(ko.mean(0), (n, 1))])
    rg = np.repeat([wt.mean(), ko.mean()], n)[:, None]
    ss_int = np.sum((mg - rg - Y.mean(0)[None, :] + gm) ** 2); ss_err = np.sum((Y - Y.mean(1)[:, None] - mg + rg) ** 2)
    df1, df2 = k - 1, (2 * n - 2) * (k - 1); F = (ss_int / df1) / (ss_err / df2)
    Sd = H @ ((np.cov(wt, rowvar=False) + np.cov(ko, rowvar=False)) / 2) @ H; eps = np.trace(Sd) ** 2 / ((k - 1) * np.sum(Sd * Sd))
    return np.array([pa, pb, pc, pd_, stats.f.sf(F, df1, df2), stats.f.sf(F, eps * df1, eps * df2)]) < 0.05
for lab, dec in (("Knockout has no effect", 0.12), ("Knockout extinguishes more slowly", 0.06)):
    r = np.array([one(dec) for _ in range(R)]).mean(axis=0)
    print(f"{lab}: p<0.05 with A any of 20 per-tone t-tests {r[0]:.1%}; B tones as n {r[1]:.1%}; C mean per mouse {r[2]:.1%}; "
          f"D slope per mouse {r[3]:.1%}; E RM-ANOVA genotype x tone {r[4]:.1%}, Greenhouse-Geisser {r[5]:.1%}")