# How often each test gives p < 0.05: 5,000 simulated breeding studies per row, from 6 to 48 heterozygous x
# heterozygous litters genotyped at weaning, with no lethality and with knockouts surviving at 75% or 50% of the rate
# of their littermates. A last block compares false positives with only 2 litters. Python 3.13, NumPy, SciPy.
# Seed 20261112.
import numpy as np
from scipy import stats
from model import counts
rng = np.random.default_rng(20261112); R = 5_000; P0 = np.array([0.25, 0.5, 0.25])

def multinom_exact(x):
    n = int(x.sum()); px = stats.multinomial.pmf(x, n, P0); tot = 0.0
    for a in range(n + 1):
        b = np.arange(n - a + 1); pr = stats.multinomial.pmf(np.column_stack([np.full(len(b), a), b, n - a - b]), n, P0)
        tot += pr[pr <= px * (1 + 1e-7)].sum()
    return tot

def tests(x, exact=False):
    n = x.sum()
    pa = stats.chisquare(x, n * P0).pvalue
    pc = stats.binomtest(int(x[2]), int(n), 0.25).pvalue
    pd_ = stats.binomtest(int(x[2]), int(x[0] + x[2]), 0.5).pvalue if x[0] + x[2] > 0 else 1.0
    out = [pa, pc, pd_] + ([multinom_exact(x)] if exact else [])
    return np.array(out) < 0.05

for via, lab in ((1.0, "No lethality (false positives)"), (0.75, "Knockouts survive at 75%"), (0.5, "Knockouts survive at 50%")):
    for L in (6, 12, 24, 48):
        sims = [counts(rng, L, via) for _ in range(R)]
        r = np.array([tests(x) for x in sims]).mean(axis=0); pups = np.mean([x.sum() for x in sims])
        print(f"{lab}, {L} litters (mean {pups:.0f} pups): A chi-square 1:2:1 {r[0]:.1%}; C binomial knockouts vs 1/4 {r[1]:.1%}; D knockouts vs wild-type {r[2]:.1%}")
sims = [counts(rng, 2) for _ in range(R)]
r = np.array([tests(x, exact=True) for x in sims]).mean(axis=0)
print(f"No lethality, 2 litters (mean {np.mean([x.sum() for x in sims]):.0f} pups): A chi-square 1:2:1 {r[0]:.1%}; B exact multinomial {r[3]:.1%}; C binomial knockouts vs 1/4 {r[1]:.1%}; D knockouts vs wild-type {r[2]:.1%}")