"""Salar Cafe: Games–Howell in Python with statsmodels 0.15.

Run from the folder containing posthoc_training_scores.csv.
The script prints the analysis, saves Tukey/Games–Howell CSV output,
and recreates the four figures used in the tutorial.
"""

from pathlib import Path
import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats
from statsmodels import __version__ as statsmodels_version
from statsmodels.stats.oneway import anova_oneway
from statsmodels.stats.multicomp import pairwise_tukeyhsd

DATA = Path("posthoc_training_scores.csv")
OUT = Path("python_output")
OUT.mkdir(exist_ok=True)
GROUP_ORDER = ["Control", "Video", "Interactive", "Tutoring"]


def fmt_p(p):
    return "< .001" if p < .001 else f"= {p:.3f}"


def games_howell_manual(df, value="score", group="group", order=None, alpha=.05):
    """Transparent formula-based Games–Howell implementation for verification.

    Native statsmodels 0.15 usage is shown later. This function is included so
    readers can inspect the pair-specific Welch/Satterthwaite calculations.
    """
    levels = order or list(pd.unique(df[group]))
    k = len(levels)
    rows = []
    for i in range(k):
        for j in range(i + 1, k):
            g1, g2 = levels[i], levels[j]
            x1 = df.loc[df[group] == g1, value].to_numpy(float)
            x2 = df.loc[df[group] == g2, value].to_numpy(float)
            n1, n2 = len(x1), len(x2)
            m1, m2 = x1.mean(), x2.mean()
            v1, v2 = x1.var(ddof=1), x2.var(ddof=1)
            se2 = v1 / n1 + v2 / n2
            diff = m1 - m2
            t_abs = abs(diff) / math.sqrt(se2)
            dfw = (se2 ** 2) / (
                ((v1 / n1) ** 2) / (n1 - 1) + ((v2 / n2) ** 2) / (n2 - 1)
            )
            q = math.sqrt(2) * t_abs
            p_adj = float(stats.studentized_range.sf(q, k, dfw))
            qcrit = float(stats.studentized_range.ppf(1 - alpha, k, dfw))
            half = qcrit * math.sqrt(se2 / 2)
            rows.append({
                "group1": g1,
                "group2": g2,
                "meandiff": diff,
                "df": dfw,
                "p_adj": p_adj,
                "lower": diff - half,
                "upper": diff + half,
                "reject": p_adj < alpha,
            })
    return pd.DataFrame(rows)


def main():
    print(f"statsmodels version: {statsmodels_version}")
    df = pd.read_csv(DATA)
    df["group"] = pd.Categorical(df["group"], categories=GROUP_ORDER, ordered=True)

    print("\nDATA CHECK")
    print(df.head())
    print("rows:", len(df))
    print("missing values:\n", df.isna().sum())
    print("group counts:\n", df["group"].value_counts(sort=False))

    desc = (
        df.groupby("group", observed=True)["score"]
          .agg(n="count", mean="mean", sd="std", median="median", minimum="min", maximum="max")
    )
    desc["se"] = desc["sd"] / np.sqrt(desc["n"])
    print("\nDESCRIPTIVES\n", desc.round(3))
    desc.to_csv(OUT / "descriptives.csv")

    arrays = [df.loc[df["group"] == g, "score"].to_numpy() for g in GROUP_ORDER]
    lev = stats.levene(*arrays, center="median")
    print(f"\nMedian-centered Levene: F(3, 97) = {lev.statistic:.3f}, p {fmt_p(lev.pvalue)}")

    classic = anova_oneway(df["score"], df["group"], use_var="equal")
    welch = anova_oneway(df["score"], df["group"], use_var="unequal")
    print(f"Classic ANOVA: F({classic.df[0]:.0f}, {classic.df[1]:.0f}) = {classic.statistic:.3f}, p {fmt_p(classic.pvalue)}")
    print(f"Welch ANOVA: F({welch.df[0]:.0f}, {welch.df[1]:.2f}) = {welch.statistic:.3f}, p {fmt_p(welch.pvalue)}")

    grand = df["score"].mean()
    ss_between = sum(
        len(df.loc[df["group"] == g]) * (df.loc[df["group"] == g, "score"].mean() - grand) ** 2
        for g in GROUP_ORDER
    )
    ss_total = ((df["score"] - grand) ** 2).sum()
    ss_within = ss_total - ss_between
    df_between = len(GROUP_ORDER) - 1
    df_within = len(df) - len(GROUP_ORDER)
    ms_within = ss_within / df_within
    eta2 = ss_between / ss_total
    omega2 = (ss_between - df_between * ms_within) / (ss_total + ms_within)
    print(f"Eta squared = {eta2:.4f}; Omega squared = {omega2:.4f}")

    # Tukey HSD: pooled/equal-variance path.
    tukey = pairwise_tukeyhsd(
        endog=df["score"], groups=df["group"], alpha=.05, use_var="equal"
    )
    tukey_frame = tukey.summary_frame()
    print("\nTUKEY HSD\n", tukey_frame)
    tukey_frame.to_csv(OUT / "tukey_results_statsmodels.csv", index=False)

    # Games–Howell: native statsmodels 0.15 path.
    gh = pairwise_tukeyhsd(
        endog=df["score"], groups=df["group"], alpha=.05, use_var="unequal"
    )
    gh_frame = gh.summary_frame()
    print("\nGAMES–HOWELL (STATSMODELS 0.15)\n", gh_frame)
    gh_frame.to_csv(OUT / "games_howell_results_statsmodels.csv", index=False)

    # Transparent verification of the same Games–Howell logic.
    gh_manual = games_howell_manual(df, order=GROUP_ORDER)
    gh_manual.to_csv(OUT / "games_howell_manual_verification.csv", index=False)

    # Figure 1: boxplots plus raw points.
    fig, ax = plt.subplots(figsize=(9, 6))
    data = [df.loc[df["group"] == g, "score"].to_numpy() for g in GROUP_ORDER]
    ax.boxplot(data, tick_labels=GROUP_ORDER, showmeans=True)
    for idx, y in enumerate(data, 1):
        jitter = np.linspace(-.13, .13, len(y))
        ax.scatter(np.full(len(y), idx) + jitter, y, s=20, alpha=.6)
    ax.set(title="Training scores by group: unequal spread is visible", xlabel="Training group", ylabel="Score")
    ax.grid(axis="y", alpha=.2)
    fig.tight_layout()
    fig.savefig(OUT / "figure-1-group-distributions.png", dpi=220, bbox_inches="tight")
    plt.close(fig)

    # Figure 2: mean 95% confidence intervals.
    means, lo_err, hi_err = [], [], []
    for g in GROUP_ORDER:
        x = df.loc[df["group"] == g, "score"].to_numpy(float)
        n, mean, sd = len(x), x.mean(), x.std(ddof=1)
        se = sd / math.sqrt(n)
        tcrit = stats.t.ppf(.975, n - 1)
        low, high = mean - tcrit * se, mean + tcrit * se
        means.append(mean); lo_err.append(mean - low); hi_err.append(high - mean)
    fig, ax = plt.subplots(figsize=(9, 5.5))
    xpos = np.arange(len(GROUP_ORDER))
    ax.errorbar(xpos, means, yerr=[lo_err, hi_err], fmt="o", capsize=7)
    ax.set_xticks(xpos, GROUP_ORDER)
    ax.set(title="Group means with 95% confidence intervals", xlabel="Training group", ylabel="Mean score")
    ax.grid(axis="y", alpha=.2)
    fig.tight_layout()
    fig.savefig(OUT / "figure-2-means-95ci.png", dpi=220, bbox_inches="tight")
    plt.close(fig)

    # Figure 3: Games–Howell forest plot from transparent verification table.
    labels = [f'{r.group1} − {r.group2}' for r in gh_manual.itertuples()]
    diffs = gh_manual["meandiff"].to_numpy()
    lower = gh_manual["lower"].to_numpy(); upper = gh_manual["upper"].to_numpy()
    ypos = np.arange(len(labels))
    fig, ax = plt.subplots(figsize=(10, 6.5))
    ax.errorbar(diffs, ypos, xerr=[diffs-lower, upper-diffs], fmt="o", capsize=5)
    ax.axvline(0, linewidth=1)
    ax.set_yticks(ypos, labels); ax.invert_yaxis()
    ax.set(title="Games–Howell pairwise mean differences with simultaneous 95% CIs",
           xlabel="Mean difference (first group − second group)")
    ax.grid(axis="x", alpha=.2)
    fig.tight_layout()
    fig.savefig(OUT / "figure-3-games-howell-forest.png", dpi=220, bbox_inches="tight")
    plt.close(fig)

    # Figure 4: adjusted p-value comparison.
    pair_labels = []
    tukey_p = []
    gh_p = []
    for r in gh_manual.itertuples():
        pair_labels.append(f"{r.group1} vs {r.group2}")
        gh_p.append(max(r.p_adj, 1e-12))
        # Match unordered pair in statsmodels Tukey frame.
        mask = (
            ((tukey_frame["group_t"] == r.group1) & (tukey_frame["group_c"] == r.group2)) |
            ((tukey_frame["group_t"] == r.group2) & (tukey_frame["group_c"] == r.group1))
        )
        tukey_p.append(max(float(tukey_frame.loc[mask, "p-adj"].iloc[0]), 1e-12))
    y = np.arange(len(pair_labels))
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.scatter(-np.log10(tukey_p), y, label="Tukey HSD", marker="o")
    ax.scatter(-np.log10(gh_p), y, label="Games–Howell", marker="x")
    ax.axvline(-math.log10(.05), linewidth=1)
    ax.set_yticks(y, pair_labels); ax.invert_yaxis()
    ax.set(title="Why the post-hoc choice changes conclusions",
           xlabel="−log10(adjusted p-value); right of line = p < .05")
    ax.legend(); ax.grid(axis="x", alpha=.2)
    fig.tight_layout()
    fig.savefig(OUT / "figure-4-tukey-vs-games-howell.png", dpi=220, bbox_inches="tight")
    plt.close(fig)

    print(f"\nSaved outputs to: {OUT.resolve()}")


if __name__ == "__main__":
    main()