# ============================================================================
#  04-05  Binomial & Poisson Calculator
#  Chapters 06-01, 06-02, 12-01
#
#  Run from the project folder:   Rscript r/analysis.R
# ============================================================================

defects <- read.csv("data/batch_defects.csv")
calls   <- read.csv("data/calls_per_hour.csv")

# ── 1. REUSABLE CALCULATORS ─────────────────────────────────────────────────
binom_calc <- function(n, p, x, a = NULL, b = NULL) {
  res <- c(
    `P(X = x)`  = dbinom(x, n, p),
    `P(X <= x)` = pbinom(x, n, p),
    `P(X >= x)` = pbinom(x - 1, n, p, lower.tail = FALSE),
    `P(X < x)`  = pbinom(x - 1, n, p),
    `P(X > x)`  = pbinom(x, n, p, lower.tail = FALSE),
    mean        = n * p,
    sd          = sqrt(n * p * (1 - p))
  )
  if (!is.null(a) && !is.null(b))
    res[["P(a <= X <= b)"]] <- sum(dbinom(a:b, n, p))
  res
}

pois_calc <- function(lambda, x, interval = 1) {
  L <- lambda * interval
  c(`lambda'`   = L,
    `P(X = x)`  = dpois(x, L),
    `P(X <= x)` = ppois(x, L),
    `P(X >= x)` = ppois(x - 1, L, lower.tail = FALSE),
    mean        = L,
    variance    = L,
    sd          = sqrt(L))
}

cat("\n===== 1. BINOMIAL CALCULATOR  (n = 20, p = 0.06) =====\n")
print(round(binom_calc(20, 0.06, x = 3, a = 1, b = 4), 6))

cat("\n===== 2. POISSON CALCULATOR  (lambda = 3.4 per hour) =====\n")
cat("Full hour:\n");        print(round(pois_calc(3.4, 4), 6))
cat("\n20-minute window (lambda rescaled by 20/60):\n")
print(round(pois_calc(3.4, 1, interval = 20/60), 6))

# ── 3. FIT THE BINOMIAL TO THE INSPECTION DATA ──────────────────────────────
cat("\n===== 3. BINOMIAL FIT — batch_defects.csv =====\n")

n_items <- defects$items_inspected[1]
p_hat   <- sum(defects$defective) / sum(defects$items_inspected)

cat(sprintf("batches inspected      %d\n", nrow(defects)))
cat(sprintf("items per batch        %d\n", n_items))
cat(sprintf("p-hat (pooled)         %.5f\n", p_hat))
cat(sprintf("observed mean          %.4f    binomial np  = %.4f\n",
            mean(defects$defective), n_items * p_hat))
cat(sprintf("observed variance      %.4f    binomial npq = %.4f\n",
            var(defects$defective), n_items * p_hat * (1 - p_hat)))

obs_b <- table(factor(defects$defective, levels = 0:n_items))
exp_b <- dbinom(0:n_items, n_items, p_hat) * nrow(defects)

fit_b <- data.frame(defective = 0:n_items,
                    observed  = as.integer(obs_b),
                    expected  = round(exp_b, 2))
print(fit_b[fit_b$observed > 0 | fit_b$expected > 0.05, ], row.names = FALSE)

# ── 4. FIT THE POISSON TO THE CALL DATA ─────────────────────────────────────
cat("\n===== 4. POISSON FIT — calls_per_hour.csv =====\n")

lam_hat <- mean(calls$calls)
cat(sprintf("hours observed         %d\n", nrow(calls)))
cat(sprintf("lambda-hat (mean)      %.4f\n", lam_hat))
cat(sprintf("observed variance      %.4f\n", var(calls$calls)))
cat(sprintf("dispersion index       %.4f   (near 1 supports Poisson)\n",
            var(calls$calls) / lam_hat))

# ── 5. GOODNESS-OF-FIT TEST FOR THE POISSON ─────────────────────────────────
cat("\n===== 5. CHI-SQUARE GOODNESS OF FIT (Poisson) =====\n")

N <- nrow(calls)
kmax <- max(calls$calls)
obs <- as.integer(table(factor(calls$calls, levels = 0:kmax)))
p_k <- dpois(0:kmax, lam_hat)
p_k[length(p_k)] <- p_k[length(p_k)] + ppois(kmax, lam_hat, lower.tail = FALSE)
exp_k <- p_k * N

# combine sparse tail classes until every EXPECTED count is >= 5
lab <- as.character(0:kmax)
while (length(exp_k) > 2 && tail(exp_k, 1) < 5) {
  exp_k[length(exp_k) - 1] <- exp_k[length(exp_k) - 1] + tail(exp_k, 1)
  obs[length(obs) - 1]     <- obs[length(obs) - 1] + tail(obs, 1)
  lab[length(lab) - 1]     <- paste0(lab[length(lab) - 1], "+")
  exp_k <- head(exp_k, -1); obs <- head(obs, -1); lab <- head(lab, -1)
}
while (length(exp_k) > 2 && exp_k[1] < 5) {
  exp_k[2] <- exp_k[2] + exp_k[1]; obs[2] <- obs[2] + obs[1]
  lab[2] <- paste0("<=", lab[2])
  exp_k <- exp_k[-1]; obs <- obs[-1]; lab <- lab[-1]
}

contrib <- (obs - exp_k)^2 / exp_k
chi2 <- sum(contrib)
df   <- length(obs) - 1 - 1        # -1 for the ESTIMATED lambda
pval <- pchisq(chi2, df, lower.tail = FALSE)

print(data.frame(calls = lab, observed = obs,
                 expected = round(exp_k, 2),
                 contribution = round(contrib, 4)), row.names = FALSE)

cat(sprintf("\nchi-square = %.4f   df = %d   p-value = %.4f\n", chi2, df, pval))
cat(sprintf("critical chi-square(0.05, %d) = %.4f\n", df, qchisq(0.95, df)))
cat(if (pval > 0.05)
      "=> FAIL TO REJECT. The Poisson model FITS -- a large p-value is the GOOD outcome here.\n"
    else
      "=> REJECT. The Poisson model does not fit these counts.\n")

# ── 6. POISSON APPROXIMATES THE BINOMIAL ────────────────────────────────────
cat("\n===== 6. POISSON AS A BINOMIAL APPROXIMATION =====\n")
cat("n = 2000, p = 0.0015  ->  lambda = np = 3\n\n")
cmp <- data.frame(x = 0:8,
                  binomial = round(dbinom(0:8, 2000, 0.0015), 6),
                  poisson  = round(dpois(0:8, 3), 6))
cmp$difference <- round(cmp$binomial - cmp$poisson, 6)
print(cmp, row.names = FALSE)
cat(sprintf("\nLargest absolute difference: %.6f -- negligible.\n",
            max(abs(cmp$difference))))

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

barplot(dbinom(0:8, 20, 0.06), names.arg = 0:8, col = "#5B2A86", border = NA,
        main = "Binomial(20, 0.06)", xlab = "Defective", ylab = "P(x)")

barplot(rbind(as.integer(obs_b)[1:9], exp_b[1:9]), beside = TRUE,
        names.arg = 0:8, col = c("#5B2A86", "#0FA3A3"), border = NA,
        main = "Defects: observed vs binomial", xlab = "Defective", ylab = "Batches")
legend("topright", c("observed", "expected"),
       fill = c("#5B2A86", "#0FA3A3"), bty = "n")

barplot(dpois(0:12, lam_hat), names.arg = 0:12, col = "#0B7A7A", border = NA,
        main = sprintf("Poisson(%.2f)", lam_hat), xlab = "Calls", ylab = "P(x)")

obs_full <- as.integer(table(factor(calls$calls, levels = 0:kmax)))
barplot(rbind(obs_full, dpois(0:kmax, lam_hat) * N), beside = TRUE,
        names.arg = 0:kmax, col = c("#5B2A86", "#0FA3A3"), border = NA,
        main = "Calls: observed vs Poisson", xlab = "Calls per hour", ylab = "Hours")
legend("topright", c("observed", "expected"),
       fill = c("#5B2A86", "#0FA3A3"), bty = "n")

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