"""Simulations for /guides/welch-anova/: false-positive rate and power of the classical F-test and
Welch's F-test for three independent groups with normal data.
Python 3.13, NumPy 2.5, SciPy 1.18.1. 100,000 simulated experiments per scenario, seed 20261110.
Run: python simulation.py  -> prints a table and writes simulation_output.json
"""
import json

import numpy as np
from scipy import stats

REPS = 100_000
ALPHA = 0.05
rng = np.random.default_rng(20261110)


def tests(samples):
    """samples: list of arrays shaped (REPS, n_i). Returns p-values (classical, welch)."""
    n = np.array([s.shape[1] for s in samples], dtype=float)
    k = len(samples)
    means = np.stack([s.mean(axis=1) for s in samples], axis=1)
    vars_ = np.stack([s.var(axis=1, ddof=1) for s in samples], axis=1)
    N = n.sum()
    grand = (means * n).sum(axis=1) / N
    ssb = ((means - grand[:, None]) ** 2 * n).sum(axis=1)
    ssw = (vars_ * (n - 1)).sum(axis=1)
    f_classic = (ssb / (k - 1)) / (ssw / (N - k))
    p_classic = stats.f.sf(f_classic, k - 1, N - k)
    w = n / vars_
    sw = w.sum(axis=1)
    mw = (w * means).sum(axis=1) / sw
    a = ((w * (means - mw[:, None]) ** 2).sum(axis=1)) / (k - 1)
    tmp = (((1 - w / sw[:, None]) ** 2) / (n - 1)).sum(axis=1)
    b = 1 + 2 * (k - 2) / (k ** 2 - 1) * tmp
    f_welch = a / b
    df2 = (k ** 2 - 1) / (3 * tmp)
    p_welch = stats.f.sf(f_welch, k - 1, df2)
    return p_classic, p_welch


def run(name, n, sd, mu=(0, 0, 0)):
    samples = [rng.normal(m, s, size=(REPS, ni)) for ni, s, m in zip(n, sd, mu)]
    pc, pw = tests(samples)
    return {"scenario": name, "n": list(n), "sd": list(sd), "means": list(mu),
            "classical": float((pc < ALPHA).mean()), "welch": float((pw < ALPHA).mean())}


# Check the vectorised Welch implementation against SciPy on a few draws.
for _ in range(5):
    xs = [rng.normal(0, s, size=(1, ni)) for ni, s in zip((6, 10, 14), (2, 1, 1))]
    pc, pw = tests(xs)
    sp = stats.f_oneway(*[x[0] for x in xs], equal_var=False).pvalue
    sc = stats.f_oneway(*[x[0] for x in xs]).pvalue
    assert abs(pw[0] - sp) < 1e-10 and abs(pc[0] - sc) < 1e-10

results = [
    run("equal SD, equal n (8 each), no effect", (8, 8, 8), (1, 1, 1)),
    run("smallest group most variable, no effect", (6, 10, 14), (2, 1, 1)),
    run("largest group most variable, no effect", (6, 10, 14), (1, 1, 2)),
    run("equal n (10 each), one SD three times the others, no effect", (10, 10, 10), (1, 1, 3)),
    run("example-like: 10, 9, 5 animals, SDs 7, 3.5, 11.6, no effect", (10, 9, 5), (7, 3.5, 11.6)),
    run("small groups (4 each), equal SD, no effect", (4, 4, 4), (1, 1, 1)),
    run("small groups (3 each), equal SD, no effect", (3, 3, 3), (1, 1, 1)),
    run("small groups (4 each), one SD three times the others, no effect", (4, 4, 4), (1, 1, 3)),
    run("power: equal SD, 8 each, one group shifted by 1.4 SD", (8, 8, 8), (1, 1, 1), (0, 0, 1.4)),
    run("power: equal SD, 5 each, one group shifted by 1.8 SD", (5, 5, 5), (1, 1, 1), (0, 0, 1.8)),
]
for r in results:
    print(f"{r['scenario']:<66} classical {r['classical']:.4f}  welch {r['welch']:.4f}")
with open("simulation_output.json", "w") as f:
    json.dump({"reps": REPS, "alpha": ALPHA, "seed": 20261110, "results": results}, f, indent=2)