# ============================================================================
#  04-08  Confidence Interval Estimator
#  Chapters 09-01, 09-02
#
#  Run from the project folder:   Rscript r/analysis.R
# ============================================================================

set.seed(2026)
battery <- read.csv("data/battery_life.csv")$hours
poll    <- read.csv("data/poll_responses.csv")$supports

# ── REUSABLE ESTIMATORS ─────────────────────────────────────────────────────
ci_mean_t <- function(x, conf = 0.95) {
  n <- length(x); se <- sd(x) / sqrt(n)
  tc <- qt(1 - (1 - conf) / 2, n - 1)
  c(n = n, mean = mean(x), sd = sd(x), se = se, df = n - 1,
    t_crit = tc, margin = tc * se,
    lower = mean(x) - tc * se, upper = mean(x) + tc * se,
    width = 2 * tc * se)
}

ci_mean_z <- function(x, sigma, conf = 0.95) {
  n <- length(x); se <- sigma / sqrt(n)
  zc <- qnorm(1 - (1 - conf) / 2)
  c(n = n, mean = mean(x), sigma = sigma, se = se,
    z_crit = zc, margin = zc * se,
    lower = mean(x) - zc * se, upper = mean(x) + zc * se)
}

ci_prop <- function(x, n, conf = 0.95) {
  ph <- x / n; se <- sqrt(ph * (1 - ph) / n)
  zc <- qnorm(1 - (1 - conf) / 2)
  c(x = x, n = n, p_hat = ph, se = se, z_crit = zc, margin = zc * se,
    lower = ph - zc * se, upper = ph + zc * se, width = 2 * zc * se)
}

# ── 1. CI FOR A MEAN ────────────────────────────────────────────────────────
cat("\n===== 1. CONFIDENCE INTERVAL FOR A MEAN — battery_life.csv =====\n")
cat("sigma is UNKNOWN (s was computed from the data)  ->  use t\n\n")
print(round(ci_mean_t(battery), 4))

cat("\nR's built-in t.test agrees:\n")
print(t.test(battery)$conf.int)

cat("\nIf sigma were known to be 96 (z-interval, for contrast):\n")
print(round(ci_mean_z(battery, sigma = 96), 4))

# ── 2. THREE CONFIDENCE LEVELS ──────────────────────────────────────────────
cat("\n===== 2. THE CONFIDENCE / PRECISION TRADE-OFF =====\n")

levels <- c(0.90, 0.95, 0.99)
tab <- t(sapply(levels, function(cl) {
  r <- ci_mean_t(battery, cl)
  c(conf = cl, t_crit = r[["t_crit"]], margin = r[["margin"]],
    lower = r[["lower"]], upper = r[["upper"]], width = r[["width"]])
}))
print(round(tab, 4))
cat("\nHigher confidence -> larger critical value -> WIDER interval.\n")
cat("You cannot have high confidence AND a narrow interval AND a small sample.\n")

# ── 3. CI FOR A PROPORTION ──────────────────────────────────────────────────
cat("\n===== 3. CONFIDENCE INTERVAL FOR A PROPORTION — poll_responses.csv =====\n")

x <- sum(poll == "Yes"); n <- length(poll)
cat(sprintf("condition check: n*p_hat = %.1f  and  n*(1-p_hat) = %.1f   (both >= 5)\n\n",
            x, n - x))
print(round(ci_prop(x, n), 5))

cat("\nAlternative methods on the SAME data:\n")
cat(sprintf("  Wald (above)     %.4f to %.4f\n",
            ci_prop(x, n)[["lower"]], ci_prop(x, n)[["upper"]]))
w <- prop.test(x, n, correct = FALSE)$conf.int
cat(sprintf("  Wilson score     %.4f to %.4f\n", w[1], w[2]))
b <- binom.test(x, n)$conf.int
cat(sprintf("  Clopper-Pearson  %.4f to %.4f\n", b[1], b[2]))
pf <- ci_prop(x + 2, n + 4)
cat(sprintf("  plus-four        %.4f to %.4f\n", pf[["lower"]], pf[["upper"]]))
cat("\nAt n = 600 all four agree closely. At small n they diverge sharply.\n")

# ── 4. SAMPLE SIZE ──────────────────────────────────────────────────────────
cat("\n===== 4. SAMPLE SIZE FOR A TARGET MARGIN OF ERROR =====\n")

n_mean <- function(sigma, E, conf = 0.95)
  ceiling((qnorm(1 - (1 - conf) / 2) * sigma / E)^2)
n_prop <- function(E, p = 0.5, conf = 0.95)
  ceiling(p * (1 - p) * (qnorm(1 - (1 - conf) / 2) / E)^2)

s_pilot <- sd(battery)
cat(sprintf("Using the pilot sd = %.2f hours\n\n", s_pilot))
cat("  MEAN                 90%%      95%%      99%%\n")
for (E in c(50, 25, 10, 5)) {
  cat(sprintf("  E = %3d hours     %6d   %6d   %6d\n", E,
              n_mean(s_pilot, E, 0.90), n_mean(s_pilot, E, 0.95),
              n_mean(s_pilot, E, 0.99)))
}

ph <- x / n
cat(sprintf("\n  PROPORTION        p=0.5    p=%.2f\n", ph))
for (E in c(0.05, 0.03, 0.02, 0.01)) {
  cat(sprintf("  E = %.2f          %6d   %6d\n", E, n_prop(E), n_prop(E, ph)))
}
cat("\nHalving E QUADRUPLES n -- precision is bought with the square of the sample.\n")

# ── 5. CI FOR THE VARIANCE ──────────────────────────────────────────────────
cat("\n===== 5. CONFIDENCE INTERVAL FOR sigma^2 AND sigma =====\n")

nb <- length(battery); s2 <- var(battery); df <- nb - 1
lo <- df * s2 / qchisq(0.975, df)
hi <- df * s2 / qchisq(0.025, df)
cat(sprintf("s^2 = %.3f   df = %d\n", s2, df))
cat(sprintf("95%% CI for sigma^2:  %.3f to %.3f\n", lo, hi))
cat(sprintf("95%% CI for sigma  :  %.3f to %.3f\n", sqrt(lo), sqrt(hi)))
cat(sprintf("\nDistance below s^2: %.2f    above: %.2f   -> NOT symmetric,\n",
            s2 - lo, hi - s2))
cat("because the chi-square distribution is right-skewed.\n")
cat("This interval also REQUIRES normality and is not robust to it.\n")

# ── 6. COVERAGE SIMULATION ──────────────────────────────────────────────────
cat("\n===== 6. DOES A 95% INTERVAL REALLY COVER 95% OF THE TIME? =====\n")

true_mu <- 1218; true_sigma <- 96; R <- 5000

covered <- replicate(R, {
  smp <- rnorm(45, true_mu, true_sigma)
  ci  <- ci_mean_t(smp)
  ci[["lower"]] <= true_mu && true_mu <= ci[["upper"]]
})
cat(sprintf("t-interval for a mean, n = 45:  coverage = %.4f  (target 0.95)\n",
            mean(covered)))

cat("\nWald proportion interval coverage, by n and p:\n")
cover_p <- function(n, p, R = 5000) {
  xs <- rbinom(R, n, p); ph <- xs / n
  e  <- qnorm(0.975) * sqrt(ph * (1 - ph) / n)
  mean(ph - e <= p & p <= ph + e)
}
grid <- expand.grid(n = c(20, 100, 600), p = c(0.05, 0.50))
grid$coverage <- round(mapply(cover_p, grid$n, grid$p), 4)
print(grid, row.names = FALSE)
cat("\nWald delivers its promise near p = 0.5, but UNDER-covers badly near 0 or 1.\n")
cat("That is why prop.test and statsmodels default to the WILSON interval.\n")

# ── 7. PLOTS ────────────────────────────────────────────────────────────────
png("ci_plots.png", width = 1100, height = 850)
par(mfrow = c(2, 2), mar = c(4.5, 4.5, 3, 1))

# 100 simulated intervals — the classic picture of "95% confidence"
K <- 100
lo_v <- hi_v <- numeric(K)
for (i in 1:K) {
  smp <- rnorm(45, true_mu, true_sigma)
  ci <- ci_mean_t(smp); lo_v[i] <- ci[["lower"]]; hi_v[i] <- ci[["upper"]]
}
hit <- lo_v <= true_mu & true_mu <= hi_v
plot(NA, xlim = range(c(lo_v, hi_v)), ylim = c(0, K + 1),
     xlab = "hours", ylab = "simulated study",
     main = sprintf("100 x 95%% CIs — %d missed", sum(!hit)))
segments(lo_v, 1:K, hi_v, 1:K, col = ifelse(hit, "#0FA3A3", "#b4122e"), lwd = 2)
abline(v = true_mu, col = "#5B2A86", lwd = 2)

barplot(tab[, "width"], names.arg = paste0(levels * 100, "%"),
        col = "#5B2A86", border = NA,
        main = "Interval width vs confidence", ylab = "width (hours)")

Es <- seq(5, 60, by = 1)
plot(Es, sapply(Es, function(E) n_mean(s_pilot, E)), type = "l", lwd = 2,
     col = "#0B7A7A", log = "y",
     main = "Sample size vs margin of error", xlab = "target E (hours)",
     ylab = "required n (log scale)")

hist(battery, breaks = 12, col = "#8A5FBF", border = "white",
     main = "Battery life with 95% CI for the mean", xlab = "hours")
ci <- ci_mean_t(battery)
abline(v = ci[["mean"]], col = "#5B2A86", lwd = 2)
abline(v = c(ci[["lower"]], ci[["upper"]]), col = "#0FA3A3", lwd = 2, lty = 2)

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