# ============================================================================
#  04-10  Chi-Square & ANOVA Test Suite
#  Chapters 12-01, 12-02, 12-03
#
#  Run from the project folder:   Rscript r/analysis.R
# ============================================================================

survey <- read.csv("data/health_survey.csv", stringsAsFactors = TRUE)
fert   <- read.csv("data/fertilizer_yield.csv", stringsAsFactors = TRUE)

survey$exercise <- factor(survey$exercise, levels = c("Regularly", "Sometimes", "Never"))
survey$health   <- factor(survey$health,   levels = c("Excellent", "Good", "Poor"))

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

obs <- table(survey$exercise)
gof <- chisq.test(obs)

print(data.frame(level = names(obs),
                 observed = as.integer(obs),
                 expected = round(as.numeric(gof$expected), 2),
                 contribution = round(as.numeric(gof$residuals^2), 4)),
      row.names = FALSE)

cat(sprintf("\nchi-square = %.4f   df = %d   p = %.6f\n",
            gof$statistic, gof$parameter, gof$p.value))
cat(sprintf("critical chi-square(0.05, %d) = %.4f\n", gof$parameter,
            qchisq(0.95, gof$parameter)))
cat(sprintf("=> %s\n", ifelse(gof$p.value <= 0.05,
    "REJECT H0. The exercise levels are NOT equally common.",
    "FAIL TO REJECT H0. No evidence against equal proportions.")))
cat(sprintf("all expected counts >= 5?  %s\n", all(gof$expected >= 5)))

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

tab <- table(survey$exercise, survey$health)
cat("OBSERVED counts:\n"); print(addmargins(tab))

ind <- chisq.test(tab)

cat("\nEXPECTED counts if independent (row x column / grand):\n")
print(round(ind$expected, 2))
cat(sprintf("\nall expected counts >= 5?  %s   (smallest = %.2f)\n",
            all(ind$expected >= 5), min(ind$expected)))

cat("\nCELL CONTRIBUTIONS (O-E)^2/E  -- this is where the story is:\n")
print(round(ind$residuals^2, 4))
cat("\nSIGNED standardized residuals (blue/positive = more than expected):\n")
print(round(ind$residuals, 3))

cat(sprintf("\nchi-square = %.4f   df = %d   p = %.3e\n",
            ind$statistic, ind$parameter, ind$p.value))
cat(sprintf("critical chi-square(0.05, %d) = %.4f\n", ind$parameter,
            qchisq(0.95, ind$parameter)))

V <- sqrt(as.numeric(ind$statistic) / (sum(tab) * min(dim(tab) - 1)))
cat(sprintf("Cramer's V = %.4f   (%s association)\n", V,
            if (V < 0.10) "negligible" else if (V < 0.30) "small"
            else if (V < 0.50) "medium" else "large"))

cat(sprintf("\n=> %s\n", ifelse(ind$p.value <= 0.05,
    "REJECT H0. Exercise habit and health rating ARE associated.",
    "FAIL TO REJECT H0. No evidence of an association.")))

cat("\nRow percentages -- P(health | exercise):\n")
print(round(100 * prop.table(tab, 1), 1))

biggest <- which(ind$residuals^2 == max(ind$residuals^2), arr.ind = TRUE)
cat(sprintf("\nLargest contribution: %s / %s  (observed %d vs expected %.1f)\n",
            rownames(tab)[biggest[1]], colnames(tab)[biggest[2]],
            tab[biggest], ind$expected[biggest]))

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

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

grp <- split(fert$growth_cm, fert$fertilizer)
summ <- data.frame(
  fertilizer = names(grp),
  n    = sapply(grp, length),
  mean = round(sapply(grp, mean), 3),
  sd   = round(sapply(grp, sd), 3),
  var  = round(sapply(grp, var), 3)
)
print(summ, row.names = FALSE)

cat(sprintf("\nvariance ratio (max/min) = %.3f  -> %s\n",
            max(summ$var) / min(summ$var),
            ifelse(max(summ$var) / min(summ$var) < 4,
                   "OK, the pooled ANOVA is appropriate",
                   "unequal; prefer Welch's ANOVA")))

# built by hand
N <- nrow(fert); k <- length(grp); grand <- mean(fert$growth_cm)
SSB <- sum(sapply(names(grp), function(g) length(grp[[g]]) * (mean(grp[[g]]) - grand)^2))
SSW <- sum(sapply(grp, function(x) sum((x - mean(x))^2)))
SST <- SSB + SSW
MSB <- SSB / (k - 1); MSW <- SSW / (N - k); Fstat <- MSB / MSW

cat("\nANOVA table, computed by hand:\n")
print(data.frame(
  Source = c("Between", "Within", "Total"),
  SS = round(c(SSB, SSW, SST), 4),
  df = c(k - 1, N - k, N - 1),
  MS = c(round(MSB, 4), round(MSW, 4), NA),
  F  = c(round(Fstat, 4), NA, NA),
  p  = c(signif(pf(Fstat, k - 1, N - k, lower.tail = FALSE), 4), NA, NA)
), row.names = FALSE)

cat(sprintf("\ncritical F(0.05, %d, %d) = %.4f\n", k - 1, N - k, qf(0.95, k - 1, N - k)))
cat(sprintf("eta squared = SSB/SST = %.4f  -> %.1f%% of the variation explained\n",
            SSB / SST, 100 * SSB / SST))

cat("\nSoftware check:\n")
model <- aov(growth_cm ~ fertilizer, data = fert)
print(summary(model))

# assumptions
cat("\nASSUMPTION CHECKS\n")
cat(sprintf("  Shapiro-Wilk on the RESIDUALS: p = %.4f\n",
            shapiro.test(residuals(model))$p.value))
cat(sprintf("  Bartlett test for equal variances: p = %.4f\n",
            bartlett.test(growth_cm ~ fertilizer, data = fert)$p.value))
cat("  (car::leveneTest is the robust alternative if you have the package)\n")

# post-hoc, ONLY because F was significant
if (pf(Fstat, k - 1, N - k, lower.tail = FALSE) <= 0.05) {
  cat("\nF IS significant -> run the post-hoc test.\n\n")
  print(TukeyHSD(model))

  q <- 3.44                                    # q(0.05, k=3, df=42), from a table
  HSD <- q * sqrt(MSW / mean(summ$n))
  cat(sprintf("\nTukey HSD by hand: q x sqrt(MSW/n) = %.2f x sqrt(%.4f/%.0f) = %.4f\n",
              q, MSW, mean(summ$n), HSD))
  cat("Any pair of means differing by more than that is significant.\n")
} else {
  cat("\nF is NOT significant -> DO NOT run a post-hoc test.\n")
}

cat("\nWelch's ANOVA (does not assume equal variances):\n")
print(oneway.test(growth_cm ~ fertilizer, data = fert, var.equal = FALSE))
cat("Kruskal-Wallis (nonparametric):\n")
print(kruskal.test(growth_cm ~ fertilizer, data = fert))

# ── PLOTS ───────────────────────────────────────────────────────────────────
png("categorical_plots.png", width = 1150, height = 900)
par(mfrow = c(2, 2), mar = c(4.5, 4.5, 3, 1))

barplot(prop.table(tab, 1), beside = TRUE,
        col = c("#5B2A86", "#8A5FBF", "#0FA3A3"), border = NA, ylim = c(0, 0.6),
        main = "Health rating by exercise habit",
        ylab = "proportion within exercise group")
legend("topright", rownames(tab), fill = c("#5B2A86", "#8A5FBF", "#0FA3A3"), bty = "n")

mosaicplot(tab, shade = TRUE, main = "Mosaic plot with residual shading",
           xlab = "exercise", ylab = "health")

boxplot(growth_cm ~ fertilizer, data = fert,
        col = c("#5B2A86", "#0FA3A3", "#0B7A7A"),
        main = "Growth by fertilizer", ylab = "growth (cm)")
points(1:k, summ$mean, pch = 4, cex = 1.5, lwd = 2)

plot(TukeyHSD(model), las = 1)

par(mfrow = c(1, 1))
dev.off()
cat("\nWrote categorical_plots.png\n")