"""Numbers for /guides/kruskal-wallis-vs-anova/ (constructed data).
Python 3.13, SciPy 1.18.1, statsmodels 0.15.0, scikit-posthocs 0.17.1.
Run from this folder: python example.py
Writes example_output.json, grades.csv and il6.csv (the CSVs are also published under
public/guides/kruskal-wallis-vs-anova/)."""
import json
import os
import numpy as np
import pandas as pd
from scipy import stats
from statsmodels.stats.multitest import multipletests
from statsmodels.stats.oneway import anova_oneway
import scikit_posthocs as sp

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

# 1. Ordinal example: histology damage grade (0-4), 8 mice per group. Values written by hand.
grades = {
    "vehicle": [0, 1, 1, 1, 2, 1, 0, 2],
    "low":     [1, 2, 2, 1, 3, 2, 2, 1],
    "high":    [2, 3, 3, 4, 2, 3, 4, 3],
}
g = pd.DataFrame([(k, v) for k, vals in grades.items() for v in vals], columns=["dose", "grade"])
g.to_csv("grades.csv", index=False)
kw = stats.kruskal(*grades.values())
out["grades_kw"] = {"H": kw.statistic, "df": 2, "p": kw.pvalue}
out["grades_medians"] = {k: float(np.median(v)) for k, v in grades.items()}
out["grades_iqr"] = {k: [float(np.percentile(v, 25)), float(np.percentile(v, 75))] for k, v in grades.items()}

# Monte Carlo permutation p-value for H (exact enumeration of 24 animals is too large)
rng = np.random.default_rng(20261110)
allv = g["grade"].to_numpy()
labels = g["dose"].to_numpy()
B = 100_000


def h_stat(vals, labs):
    return stats.kruskal(*[vals[labs == k] for k in grades]).statistic


h_obs = h_stat(allv, labels)
count = 0
for _ in range(B):
    count += h_stat(allv, rng.permutation(labels)) >= h_obs - 1e-12
out["grades_perm_p"] = (count + 1) / (B + 1)
out["grades_perm_resamples"] = B

# Dunn's test, own construction (pooled ranks, tie-corrected variance), Holm-adjusted
ranks = stats.rankdata(allv)
N = len(allv)
_, tie_counts = np.unique(allv, return_counts=True)
tie = (tie_counts**3 - tie_counts).sum() / (12 * (N - 1))
pairs = [("low", "vehicle"), ("high", "vehicle"), ("high", "low")]
dunn = []
for a, b in pairs:
    ra, rb = ranks[labels == a], ranks[labels == b]
    se = np.sqrt((N * (N + 1) / 12 - tie) * (1 / len(ra) + 1 / len(rb)))
    z = (ra.mean() - rb.mean()) / se
    dunn.append({"pair": f"{a} - {b}", "mean_rank_diff": ra.mean() - rb.mean(), "z": z,
                 "p_unadj": 2 * stats.norm.sf(abs(z))})
adj = multipletests([d["p_unadj"] for d in dunn], method="holm")[1]
for d, p in zip(dunn, adj):
    d["p_holm"] = p
out["grades_dunn"] = dunn
sp_holm = sp.posthoc_dunn(g, val_col="grade", group_col="dose", p_adjust="holm")
out["grades_dunn_scikit_posthocs_holm"] = {f"{a} - {b}": float(sp_holm.loc[a, b]) for a, b in pairs}

# Separate Mann-Whitney tests re-rank each pair: different p-values from Dunn
mw = []
for a, b in pairs:
    r = stats.mannwhitneyu(grades[a], grades[b], alternative="two-sided")
    mw.append({"pair": f"{a} - {b}", "U": r.statistic, "p_unadj": r.pvalue,
               "A": r.statistic / (len(grades[a]) * len(grades[b]))})   # Vargha-Delaney A
madj = multipletests([m["p_unadj"] for m in mw], method="holm")[1]
for m, p in zip(mw, madj):
    m["p_holm"] = p
out["grades_mannwhitney"] = mw

# 2. Skewed continuous example: IL-6 (pg/mL), 6 mice per group, log-normal.
# The seed was fixed in advance; the data were not selected for their result.
rng2 = np.random.default_rng(20261110)
il6 = {"vehicle": np.round(np.exp(rng2.normal(np.log(20), 0.6, 6)), 1),
       "low":     np.round(np.exp(rng2.normal(np.log(30), 0.6, 6)), 1),
       "high":    np.round(np.exp(rng2.normal(np.log(45), 0.6, 6)), 1)}
d = pd.DataFrame([(k, v) for k, vals in il6.items() for v in vals], columns=["dose", "il6"])
d.to_csv("il6.csv", index=False)
out["il6_values"] = {k: v.tolist() for k, v in il6.items()}
cl = stats.f_oneway(*il6.values())
we = stats.f_oneway(*il6.values(), equal_var=False)
logs = [np.log(v) for v in il6.values()]
wl = stats.f_oneway(*logs, equal_var=False)
kw2 = stats.kruskal(*il6.values())
out["il6_tests"] = {
    "classical_anova_raw": {"F": cl.statistic, "p": cl.pvalue},
    "welch_anova_raw": {"F": we.statistic, "p": we.pvalue},
    "welch_anova_log": {"F": wl.statistic, "p": wl.pvalue,
                        "df_denom": float(anova_oneway(logs, use_var="unequal").df_denom)},
    "kruskal": {"H": kw2.statistic, "p": kw2.pvalue},
}
out["il6_geometric_means"] = {k: float(np.exp(np.log(v).mean())) for k, v in il6.items()}
# ratio high/vehicle with Welch 95% CI on the log scale (unadjusted)
lh, lv = np.log(il6["high"]), np.log(il6["vehicle"])
t = stats.ttest_ind(lh, lv, equal_var=False)
ci = t.confidence_interval(0.95)
out["il6_ratio_high_vehicle"] = {"ratio": float(np.exp(lh.mean() - lv.mean())),
                                 "ci95_unadjusted": [float(np.exp(ci.low)), float(np.exp(ci.high))],
                                 "p_unadj": t.pvalue}

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