# Week 6 supplied machinery: source this file; do not memorise its internals.
# stats supplies simulation and GLMs. dplyr, tibble, ggplot2 and readr reuses the course's table
# grammar and supplies ggplot2. No specialist power package is required.
library("dplyr")
library("tibble")
library("ggplot2")
library("readr")

# Room 1 is the reference for every group, matching the Week 5 starters.
scenario <- function(group, sensitivity = FALSE, null = FALSE) {
  stopifnot(length(group) == 1, group %in% 1:4)
  s <- list(group = group,
    label = c("Gaussian gain", "Binary treatment", "Poisson events", "Gamma FCR")[group],
    beta_0 = c(695, qlogis(.30), log(3), log(2.80))[group],
    beta_diet = c(45, log(.60), log(.65), log(.92))[group],
    room_effects = list(c(0, 25, 45, 35), c(0, .35, .65, .50),
                       c(0, .20, .40, .30), c(0, -.06, -.10, -.04))[[group]],
    family = list(gaussian("identity"), binomial("logit"),
                  poisson("log"), Gamma("log"))[[group]],
    sd = 60, shape = 25, nb_size = Inf,
    scenario = if (sensitivity) "Sensitivity" else "Baseline")
  if (sensitivity) {
    if (group == 1) s$sd <- 90
    if (group == 2) s$beta_0 <- qlogis(.15)
    if (group == 3) s$beta_0 <- log(1.5)
    if (group == 4) s$shape <- 9
  }
  if (null) s$beta_diet <- 0
  s$null <- null
  s
}

make_design <- function(total_pens = 16, seed = 505) {
  stopifnot(length(total_pens) == 1, is.finite(total_pens),
            total_pens >= 16, total_pens %% 8 == 0)
  set.seed(seed)
  per_room <- total_pens / 4
  tibble(pen_id = sprintf("P%03d", seq_len(total_pens)),
         room = factor(rep(paste0("Room ", 1:4), each = per_room)),
         position = rep(seq_len(per_room), 4)) |>
    group_by(room) |>
    mutate(diet = sample(rep(c("Control", "High fibre"), each = per_room / 2))) |>
    ungroup() |>
    mutate(diet = factor(diet, levels = c("Control", "High fibre")))
}

generate_response <- function(design, s) {
  # No seed reset here: each call draws new responses from the running stream.
  d <- design |>
    mutate(eta = s$beta_0 + s$beta_diet * (diet == "High fibre") +
             s$room_effects[as.integer(room)], mu = s$family$linkinv(eta))
  d$response <- switch(s$group,
    rnorm(nrow(d), d$mu, s$sd),
    rbinom(nrow(d), 1, d$mu),
    if (is.finite(s$nb_size)) rnbinom(nrow(d), mu = d$mu, size = s$nb_size)
      else rpois(nrow(d), d$mu),
    rgamma(nrow(d), shape = s$shape, scale = d$mu / s$shape))
  d
}

analyse_one <- function(d, s, alpha = .05) {
  # A declared conservative screening procedure, not a formal separation test.
  # Any warning/invalid or extreme binary fit is unresolved and not a detection.
  warnings <- character()
  result <- tryCatch(withCallingHandlers({
    full <- glm(response ~ diet + room, family = s$family, data = d)
    reduced <- glm(response ~ room, family = s$family, data = d)
    tab <- coef(summary(full))
    beta <- unname(coef(full)["dietHigh fibre"])
    se <- tab["dietHigh fibre", "Std. Error"]
    test <- if (s$group %in% c(1, 4)) "F" else "Chisq"
    comparison <- anova(reduced, full, test = test)
    p <- comparison[[if (test == "F") "Pr(>F)" else "Pr(>Chi)"]][2]
    unstable <- any(vapply(list(full, reduced), function(f) {
      !isTRUE(f$converged) || isTRUE(f$boundary) ||
        any(!is.finite(coef(f))) || f$rank < length(coef(f)) ||
        (s$group == 2 && any(fitted(f) < 1e-6 | fitted(f) > 1 - 1e-6))
    }, logical(1)))
    unstable <- unstable || any(!is.finite(c(beta, se, p))) || se <= 0 ||
      (s$group == 2 && (abs(beta) > 10 || se > 10))
    if (length(p) != 1 || p < 0 || p > 1) unstable <- TRUE
    critical <- if (s$group %in% c(1, 4)) qt(.975, df.residual(full)) else qnorm(.975)
    tibble(beta = beta, se = se, p = p,
           ci_low = beta - critical * se, ci_high = beta + critical * se,
           converged = isTRUE(full$converged) && isTRUE(reduced$converged),
           boundary = isTRUE(full$boundary) || isTRUE(reduced$boundary),
           unresolved = unstable, error = "")
  }, warning = function(w) {
    warnings <<- c(warnings, conditionMessage(w))
    invokeRestart("muffleWarning")
  }), error = function(e) tibble(beta = NA_real_, se = NA_real_, p = NA_real_,
      ci_low = NA_real_, ci_high = NA_real_, converged = FALSE, boundary = NA,
      unresolved = TRUE, error = conditionMessage(e)))
  result$warnings <- paste(unique(warnings), collapse = " / ")
  result$unresolved <- result$unresolved | length(warnings) > 0
  result$detected <- !result$unresolved & !is.na(result$p) & result$p < alpha
  result
}

run_simulations <- function(s, total_pens = 16, B = 200, seed = 6001) {
  stopifnot(length(B) == 1, is.finite(B), B >= 1, B == as.integer(B))
  design <- make_design(total_pens)
  set.seed(seed) # Once per batch, never inside generate_response().
  bind_rows(lapply(seq_len(B), function(i) {
    analyse_one(generate_response(design, s), s) |> mutate(iteration = i)
  })) |>
    mutate(group = s$group, label = s$label, scenario = s$scenario,
           null = s$null, total_pens = total_pens, pens_per_diet = total_pens / 2,
           seed = seed)
}

summarise_runs <- function(runs) {
  # Wilson Monte Carlo intervals remain meaningful at observed 0% or 100%.
  runs |>
    group_by(group, label, scenario, null, total_pens, pens_per_diet) |>
    summarise(attempted = n(), detections = sum(detected),
              unresolved = sum(unresolved), .groups = "drop") |>
    mutate(detection = detections / attempted, unresolved_rate = unresolved / attempted,
      mc_se = sqrt(detection * (1 - detection) / attempted),
      centre = (detection + qnorm(.975)^2 / (2 * attempted)) /
        (1 + qnorm(.975)^2 / attempted),
      half = qnorm(.975) * sqrt(detection * (1 - detection) / attempted +
        qnorm(.975)^2 / (4 * attempted^2)) / (1 + qnorm(.975)^2 / attempted),
      lower = pmax(0, centre - half), upper = pmin(1, centre + half)) |>
    select(-centre, -half)
}

run_curve <- function(s, sizes = c(16, 32, 48, 64), B = 200, seed = 6001) {
  bind_rows(lapply(sizes, function(n) run_simulations(s, n, B, seed + n)))
}

plot_curve <- function(summary) {
  ggplot(summary, aes(total_pens, detection, colour = scenario, group = scenario)) +
    geom_hline(yintercept = .8, linetype = 2, colour = "grey50") +
    geom_line() + geom_point(size = 2) +
    geom_errorbar(aes(ymin = lower, ymax = upper), width = 1.5) +
    scale_colour_manual(values = c(Baseline = "#003d73", Sensitivity = "#cc6600")) +
    scale_x_continuous(breaks = c(16, 32, 48, 64)) +
    scale_y_continuous(limits = c(0, 1), labels = scales::label_percent()) +
    labs(x = "Total pens (half on each diet)",
         y = "Operational detection probability", colour = NULL,
         caption = "95% Monte Carlo intervals; unresolved analyses count as no detection.\nFour fixed rooms; fixed balanced allocation at each size.") +
    theme_minimal(base_size = 13) + theme(legend.position = "bottom")
}