"""Simulations for /guides/should-you-test-for-normality/ (Python 3.13, NumPy, SciPy 1.18.1).

(a) How often Shapiro-Wilk (alpha = 0.05) flags a sample, by distribution and sample size.
(b) False-positive rate (no true difference in means) of Student's t-test used every time,
    - the two-stage strategy: Shapiro-Wilk on each group; Student's t-test if both pass,
      otherwise Mann-Whitney U (SciPy defaults), and
    - Welch's t-test used every time.
Seed 20261110, fixed before any run. Run: python simulation.py > simulation_output.json
"""
import json
import numpy as np
from scipy import stats

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


def draw(dist, n):
    if dist == "normal":
        return rng.normal(0.0, 1.0, n)
    if dist == "lognormal (sigma 0.5)":
        return rng.lognormal(0.0, 0.5, n)
    if dist == "t, 5 df (heavier tails)":
        return rng.standard_t(5, n)
    if dist == "lognormal (sigma 0.1, nearly normal)":
        return rng.lognormal(0.0, 0.1, n)
    raise ValueError(dist)


power = []
for dist in ["normal", "lognormal (sigma 0.5)", "t, 5 df (heavier tails)", "lognormal (sigma 0.1, nearly normal)"]:
    for n in [5, 10, 30, 100, 1000]:
        hits = sum(stats.shapiro(draw(dist, n)).pvalue < ALPHA for _ in range(REPS))
        power.append({"distribution": dist, "n": n, "flagged": hits / REPS})

# Two-stage vs always-Welch, true means equal in every scenario.
SCENARIOS = [
    # name, n1, n2, sd1, sd2, shape
    ("normal, equal SD, 5 vs 5", 5, 5, 1.0, 1.0, "normal"),
    ("normal, equal SD, 10 vs 10", 10, 10, 1.0, 1.0, "normal"),
    ("lognormal, same distribution, 10 vs 10", 10, 10, 1.0, 1.0, "lognormal"),
    ("normal, SD 2 vs 1, 10 vs 10", 10, 10, 2.0, 1.0, "normal"),
    ("normal, SD 2 vs 1, 5 vs 15 (smaller group more variable)", 5, 15, 2.0, 1.0, "normal"),
]


def group(n, sd, shape):
    if shape == "normal":
        return rng.normal(0.0, sd, n)
    return rng.lognormal(0.0, 0.5, n)  # same distribution in both groups, so means are equal


two_stage = []
for name, n1, n2, sd1, sd2, shape in SCENARIOS:
    fp_two = fp_welch = fp_student = passed = fp_t_given_pass = 0
    for _ in range(REPS):
        a, b = group(n1, sd1, shape), group(n2, sd2, shape)
        both_pass = stats.shapiro(a).pvalue >= ALPHA and stats.shapiro(b).pvalue >= ALPHA
        if both_pass:
            passed += 1
            p = stats.ttest_ind(a, b).pvalue  # SciPy default: Student (equal_var=True)
            fp_t_given_pass += p < ALPHA
        else:
            p = stats.mannwhitneyu(a, b).pvalue
        fp_two += p < ALPHA
        fp_welch += stats.ttest_ind(a, b, equal_var=False).pvalue < ALPHA
        fp_student += stats.ttest_ind(a, b).pvalue < ALPHA
    two_stage.append({
        "scenario": name,
        "share_sent_to_t_test": passed / REPS,
        "two_stage_false_positive": fp_two / REPS,
        "t_test_false_positive_when_chosen": fp_t_given_pass / passed if passed else None,
        "always_welch_false_positive": fp_welch / REPS,
        "always_student_false_positive": fp_student / REPS,
    })

print(json.dumps({"seed": 20261110, "reps": REPS, "alpha": ALPHA, "shapiro_flag_rate": power,
                  "two_stage_vs_welch": two_stage}, indent=1))