"""
04-10  Chi-Square & ANOVA Test Suite
Chapters 12-01, 12-02, 12-03

Run from the project folder:   python python/analysis.py
"""
import numpy as np
import pandas as pd
from scipy import stats
import statsmodels.api as sm
from statsmodels.formula.api import ols
from statsmodels.stats.multicomp import pairwise_tukeyhsd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

pd.set_option("display.width", 120)

survey = pd.read_csv("data/health_survey.csv")
fert = pd.read_csv("data/fertilizer_yield.csv")

ex_levels = ["Regularly", "Sometimes", "Never"]
h_levels = ["Excellent", "Good", "Poor"]
survey["exercise"] = pd.Categorical(survey["exercise"], categories=ex_levels, ordered=True)
survey["health"] = pd.Categorical(survey["health"], categories=h_levels, ordered=True)

# ════════════════════════════════════════════════════════════════════════════
#  1. GOODNESS OF FIT
# ════════════════════════════════════════════════════════════════════════════
print("\n===== 1. GOODNESS-OF-FIT — exercise level =====")
print("H0: the three exercise levels are equally likely (p = 1/3 each)")
print("H1: at least one level differs\n")

obs = survey["exercise"].value_counts().reindex(ex_levels)
exp = np.full(3, obs.sum() / 3)
gof = stats.chisquare(obs.to_numpy(), f_exp=exp)

print(pd.DataFrame({"level": ex_levels,
                    "observed": obs.to_numpy(),
                    "expected": exp.round(2),
                    "contribution": ((obs.to_numpy() - exp) ** 2 / exp).round(4)
                    }).to_string(index=False))

print(f"\nchi-square = {gof.statistic:.4f}   df = 2   p = {gof.pvalue:.6f}")
print(f"critical chi-square(0.05, 2) = {stats.chi2.ppf(0.95, 2):.4f}")
print("=> " + ("REJECT H0. The exercise levels are NOT equally common."
                if gof.pvalue <= 0.05 else
                "FAIL TO REJECT H0. No evidence against equal proportions."))
print(f"all expected counts >= 5?  {(exp >= 5).all()}")

# ════════════════════════════════════════════════════════════════════════════
#  2. TEST OF INDEPENDENCE
# ════════════════════════════════════════════════════════════════════════════
print("\n===== 2. TEST OF INDEPENDENCE — exercise x health =====")
print("H0: exercise habit and health rating are INDEPENDENT")
print("H1: they are ASSOCIATED\n")

tab = pd.crosstab(survey["exercise"], survey["health"])
margins = pd.crosstab(survey["exercise"], survey["health"], margins=True)
print("OBSERVED counts:")
print(margins.to_string())

chi2, p, dof, expected = stats.chi2_contingency(tab)
expected = pd.DataFrame(expected, index=tab.index, columns=tab.columns)

print("\nEXPECTED counts if independent (row x column / grand):")
print(expected.round(2).to_string())
print(f"\nall expected counts >= 5?  {(expected.values >= 5).all()}   "
      f"(smallest = {expected.values.min():.2f})")

contrib = (tab - expected) ** 2 / expected
resid = (tab - expected) / np.sqrt(expected)
print("\nCELL CONTRIBUTIONS (O-E)^2/E  -- this is where the story is:")
print(contrib.round(4).to_string())
print("\nSIGNED standardized residuals (positive = more than expected):")
print(resid.round(3).to_string())

print(f"\nchi-square = {chi2:.4f}   df = {dof}   p = {p:.3e}")
print(f"critical chi-square(0.05, {dof}) = {stats.chi2.ppf(0.95, dof):.4f}")

n_tot = tab.values.sum()
V = np.sqrt(chi2 / (n_tot * (min(tab.shape) - 1)))
strength = ("negligible" if V < 0.10 else "small" if V < 0.30
            else "medium" if V < 0.50 else "large")
print(f"Cramer's V = {V:.4f}   ({strength} association)")

print("\n=> " + ("REJECT H0. Exercise habit and health rating ARE associated."
                 if p <= 0.05 else
                 "FAIL TO REJECT H0. No evidence of an association."))

print("\nRow percentages -- P(health | exercise):")
print((100 * tab.div(tab.sum(axis=1), axis=0)).round(1).to_string())

i, j = np.unravel_index(contrib.values.argmax(), contrib.shape)
print(f"\nLargest contribution: {tab.index[i]} / {tab.columns[j]}  "
      f"(observed {tab.iloc[i, j]} vs expected {expected.iloc[i, j]:.1f})")

print("\nCAUTION: this is an OBSERVATIONAL survey. It shows ASSOCIATION only.")
print("Age, income, and pre-existing conditions plausibly drive both variables.")

# ════════════════════════════════════════════════════════════════════════════
#  3. ONE-WAY ANOVA
# ════════════════════════════════════════════════════════════════════════════
print("\n===== 3. ONE-WAY ANOVA — fertilizer vs growth =====")
print("H0: mu_A = mu_B = mu_C")
print("H1: at least one mean differs\n")

summ = fert.groupby("fertilizer")["growth_cm"].agg(
    n="count", mean="mean", sd=lambda x: x.std(ddof=1), var=lambda x: x.var(ddof=1))
print(summ.round(3).to_string())

ratio = summ["var"].max() / summ["var"].min()
print(f"\nvariance ratio (max/min) = {ratio:.3f}  -> "
      + ("OK, the pooled ANOVA is appropriate" if ratio < 4
         else "unequal; prefer Welch's ANOVA"))

groups = [g["growth_cm"].to_numpy() for _, g in fert.groupby("fertilizer")]
N, k = len(fert), len(groups)
grand = fert["growth_cm"].mean()
SSB = sum(g.size * (g.mean() - grand) ** 2 for g in groups)
SSW = sum(((g - g.mean()) ** 2).sum() for g in groups)
SST = SSB + SSW
MSB, MSW = SSB / (k - 1), SSW / (N - k)
F = MSB / MSW
pF = stats.f.sf(F, k - 1, N - k)

print("\nANOVA table, computed by hand:")
print(pd.DataFrame({
    "Source": ["Between", "Within", "Total"],
    "SS": [round(SSB, 4), round(SSW, 4), round(SST, 4)],
    "df": [k - 1, N - k, N - 1],
    "MS": [round(MSB, 4), round(MSW, 4), np.nan],
    "F": [round(F, 4), np.nan, np.nan],
    "p": [float(f"{pF:.4g}"), np.nan, np.nan],
}).to_string(index=False))

print(f"\ncritical F(0.05, {k-1}, {N-k}) = {stats.f.ppf(0.95, k-1, N-k):.4f}")
print(f"eta squared = SSB/SST = {SSB/SST:.4f}  -> {100*SSB/SST:.1f}% of the "
      "variation explained")

print("\nSoftware check:")
model = ols("growth_cm ~ C(fertilizer)", data=fert).fit()
print(sm.stats.anova_lm(model, typ=2).round(6).to_string())

print("\nASSUMPTION CHECKS")
print(f"  Shapiro-Wilk on the RESIDUALS: p = {stats.shapiro(model.resid).pvalue:.4f}")
print(f"  Levene test for equal variances: p = {stats.levene(*groups).pvalue:.4f}")
print(f"  Bartlett test (assumes normality): p = {stats.bartlett(*groups).pvalue:.4f}")

if pF <= 0.05:
    print("\nF IS significant -> run the post-hoc test.\n")
    tukey = pairwise_tukeyhsd(fert["growth_cm"], fert["fertilizer"], alpha=0.05)
    print(tukey)
    q = 3.44                               # q(0.05, k=3, df=42), from a table
    HSD = q * np.sqrt(MSW / summ["n"].mean())
    print(f"\nTukey HSD by hand: q x sqrt(MSW/n) = {q:.2f} x "
          f"sqrt({MSW:.4f}/{summ['n'].mean():.0f}) = {HSD:.4f}")
    print("Any pair of means differing by more than that is significant.")
else:
    print("\nF is NOT significant -> DO NOT run a post-hoc test.")

print("\nKruskal-Wallis (nonparametric):")
print(stats.kruskal(*groups))

# ── PLOTS ───────────────────────────────────────────────────────────────────
fig, axes = plt.subplots(2, 2, figsize=(11.5, 9))

row_pct = tab.div(tab.sum(axis=1), axis=0)
w = 0.25
xs = np.arange(len(h_levels))
for i, ex in enumerate(ex_levels):
    axes[0, 0].bar(xs + (i - 1) * w, row_pct.loc[ex], w, label=ex,
                   color=["#5B2A86", "#8A5FBF", "#0FA3A3"][i])
axes[0, 0].set(title="Health rating by exercise habit",
               ylabel="proportion within exercise group",
               xticks=xs, xticklabels=h_levels)
axes[0, 0].legend()

im = axes[0, 1].imshow(resid, cmap="PuOr", vmin=-4, vmax=4)
axes[0, 1].set(xticks=range(3), yticks=range(3),
               xticklabels=h_levels, yticklabels=ex_levels,
               title="Standardized residuals")
for a in range(3):
    for b in range(3):
        axes[0, 1].text(b, a, f"{resid.iloc[a, b]:+.2f}", ha="center", va="center")
fig.colorbar(im, ax=axes[0, 1])

axes[1, 0].boxplot(groups, tick_labels=list(summ.index), showmeans=True)
axes[1, 0].set(title="Growth by fertilizer", ylabel="growth (cm)")

if pF <= 0.05:
    tukey.plot_simultaneous(ax=axes[1, 1])
    axes[1, 1].set_title("Tukey HSD simultaneous intervals")
else:
    axes[1, 1].axis("off")

plt.tight_layout()
plt.savefig("categorical_plots.png", dpi=110)
print("\nWrote categorical_plots.png")