"""Small-sample floor and simulations for /guides/kruskal-wallis-vs-anova/.
Python 3.13, SciPy 1.18.1. Run from anywhere: python small_n_and_simulation.py
Writes small_n_and_simulation_output.json beside this file.

1. Exact permutation distribution of H for k = 3 groups of n = 3, 4 and 5 (no ties): the smallest
   attainable exact p-value, and the chi-square approximation's p-value for the same H.
2. False-positive rate when all groups have the same mean and median but different spreads
   (normal data), for Kruskal-Wallis, classical ANOVA and Welch's ANOVA.
3. Power under normal data with a pure location shift (Kruskal-Wallis vs classical ANOVA).
Seed 20261110 throughout; 20,000 runs per scenario."""
import itertools
import json
import math
import os
import numpy as np
from scipy import stats

os.chdir(os.path.dirname(os.path.abspath(__file__)))
out = {}


def h_from_ranksums(sums, n, N):
    return 12 / (N * (N + 1)) * sum(s * s / n for s in sums) - 3 * (N + 1)


# 1. exact floor for equal group sizes, k = 3
floor = []
for n in (3, 4, 5):
    N = 3 * n
    ranks = list(range(1, N + 1))
    hs = []
    # enumerate group-1 sets, then group-2 sets from the rest (group 3 is the remainder)
    for g1 in itertools.combinations(ranks, n):
        rest = [r for r in ranks if r not in g1]
        s1 = sum(g1)
        for g2 in itertools.combinations(rest, n):
            s2 = sum(g2)
            s3 = N * (N + 1) // 2 - s1 - s2
            hs.append(h_from_ranksums((s1, s2, s3), n, N))
    hs = np.array(hs)
    hmax = hs.max()
    total = len(hs)
    p_exact = float(np.mean(hs >= hmax - 1e-9))
    floor.append({"n_per_group": n, "assignments": total, "H_max": float(hmax),
                  "p_exact_min": p_exact, "p_chisq_at_Hmax": float(stats.chi2.sf(hmax, 2))})
out["exact_floor_k3"] = floor

rng = np.random.default_rng(20261110)
R = 20_000


def rates(sample_fn):
    kw = cl = we = 0
    for _ in range(R):
        gs = sample_fn()
        kw += stats.kruskal(*gs).pvalue < 0.05
        cl += stats.f_oneway(*gs).pvalue < 0.05
        we += stats.f_oneway(*gs, equal_var=False).pvalue < 0.05
    return {"kruskal": kw / R, "classical_anova": cl / R, "welch_anova": we / R}


# 2. same mean and median (normal, centred at 0), different spreads
null_scenarios = [
    {"label": "equal n 8/8/8, SD 1/1/1", "n": (8, 8, 8), "sd": (1, 1, 1)},
    {"label": "equal n 8/8/8, SD 1/1/4", "n": (8, 8, 8), "sd": (1, 1, 4)},
    {"label": "n 12/8/4, SD 1/1/4 (smallest group most variable)", "n": (12, 8, 4), "sd": (1, 1, 4)},
    {"label": "n 4/8/12, SD 1/1/4 (largest group most variable)", "n": (4, 8, 12), "sd": (1, 1, 4)},
]
out["null_unequal_spread"] = []
for sc in null_scenarios:
    r = rates(lambda sc=sc: [rng.normal(0, s, n) for n, s in zip(sc["n"], sc["sd"])])
    out["null_unequal_spread"].append({**sc, **r})

# 3. power: normal data, shifts 0 / 0.5 / 1.0 SD, 8 per group
def power_shift():
    kw = cl = 0
    for _ in range(R):
        gs = [rng.normal(m, 1, 8) for m in (0, 0.5, 1.0)]
        kw += stats.kruskal(*gs).pvalue < 0.05
        cl += stats.f_oneway(*gs).pvalue < 0.05
    return {"kruskal": kw / R, "classical_anova": cl / R}


out["power_normal_shift_n8"] = power_shift()
out["runs_per_scenario"] = R
out["seed"] = 20261110
out["are_normal_3_over_pi"] = 3 / math.pi

with open("small_n_and_simulation_output.json", "w") as f:
    json.dump(out, f, indent=2)
print(json.dumps(out, indent=2))