"""Pairwise Brunner-Munzel tests with Holm adjustment, for /guides/kruskal-wallis-vs-anova/.
Python 3.13, SciPy 1.18.1, statsmodels 0.15.0. Run from anywhere: python brunner_munzel.py
Writes brunner_munzel_output.json beside this file.
1. The worked grades example (pairwise Brunner-Munzel, Holm-adjusted).
2. Familywise false-positive rate (any pair p < 0.05 after Holm) when all groups have the same mean and
   median but different spreads: pairwise Brunner-Munzel vs Dunn (both Holm), 10,000 runs, seed 20261110.
The grades example is reported for completeness; with heavily tied 0-4 scores and 8 animals per group the
Brunner-Munzel t approximation is not used on the page."""
import json
import warnings
warnings.filterwarnings('ignore')
import os
import numpy as np
import pandas as pd
from scipy import stats
from statsmodels.stats.multitest import multipletests

os.chdir(os.path.dirname(os.path.abspath(__file__)))
out = {}
g = pd.read_csv("grades.csv")
pairs = [("low", "vehicle"), ("high", "vehicle"), ("high", "low")]
res = []
for a, b in pairs:
    r = stats.brunnermunzel(g.loc[g.dose == a, "grade"], g.loc[g.dose == b, "grade"])
    res.append({"pair": f"{a} - {b}", "W": float(r.statistic), "p_unadj": float(r.pvalue)})
adj = multipletests([x["p_unadj"] for x in res], method="holm")[1]
for x, p in zip(res, adj):
    x["p_holm"] = float(p)
out["grades_brunner_munzel_holm"] = res


def dunn_p(groups):
    vals = np.concatenate(groups)
    labs = np.concatenate([[i] * len(x) for i, x in enumerate(groups)])
    ranks = stats.rankdata(vals)
    N = len(vals)
    ps = []
    for i, j in [(0, 1), (0, 2), (1, 2)]:
        ri, rj = ranks[labs == i], ranks[labs == j]
        se = np.sqrt(N * (N + 1) / 12 * (1 / len(ri) + 1 / len(rj)))
        ps.append(2 * stats.norm.sf(abs((ri.mean() - rj.mean()) / se)))
    return ps


rng = np.random.default_rng(20261110)
R = 10_000
scen = [{"label": "n 12/8/4, SD 1/1/4", "n": (12, 8, 4), "sd": (1, 1, 4)},
        {"label": "n 8/8/8, SD 1/1/4", "n": (8, 8, 8), "sd": (1, 1, 4)},
        ]
out["familywise_null"] = []
for sc in scen:
    bm = du = 0
    for _ in range(R):
        gs = [rng.normal(0, s, n) for n, s in zip(sc["n"], sc["sd"])]
        pb = [stats.brunnermunzel(gs[i], gs[j]).pvalue for i, j in [(0, 1), (0, 2), (1, 2)]]
        bm += np.nanmin(pb) < 0.05 / 3   # Holm rejects any pair iff the smallest p < alpha/m
        du += min(dunn_p(gs)) < 0.05 / 3
    out["familywise_null"].append({**sc, "brunner_munzel_holm": bm / R, "dunn_holm": du / R})
out["runs"] = R
out["seed"] = 20261110
with open("brunner_munzel_output.json", "w") as f:
    json.dump(out, f, indent=2)
print(json.dumps(out, indent=2))