# Class 8 follow-up: plan a future experiment in which whole pens receive diets.
# Run this script from the supplied Class 8 project folder, one section at a
# time. It compares four plans: 16 or 32 pens, with pen SD 40 or 60 g/day.
# Every pen contains six animals. Half the pens receive each diet.
# This simpler future experiment has no room blocks or bedding treatments.

# 1. Load the packages used below -------------------------------------------
# simglm generates trials, fits models and calculates power.
# lme4 fits the daily-gain model; performance checks whether fitting succeeded.
# purrr repeats steps for different plans and trials.
# readr, dplyr and tidyr read, prepare and save the tables.
library("simglm")
library("lme4", warn.conflicts = FALSE)
library("performance")
library("purrr")
library("readr")
library("dplyr", warn.conflicts = FALSE)
library("tidyr")

# 2. Describe one possible experiment ---------------------------------------
# These numbers describe the future study we want to simulate.
# Changing them changes the question answered by the power calculation.
settings <- list(
  # Fit a diet difference and a shared difference for each pen.
  formula = gain ~ diet + (1 | pen),
  fixed = list(
    diet = list(
      var_type = "factor",
      levels = c("Control", "Supplement"),
      # Level 2 means that researchers assign diet to whole pens.
      var_level = 2,
      # Give half the pens each diet; all animals in a pen share its diet.
      force_equal = TRUE
    )
  ),
  # Level 1 counts animals per pen; level 2 counts pens in the experiment.
  sample_size = list(level1 = 6, level2 = 16),
  # Control mean = 700 g/day; Supplement adds an assumed 35 g/day.
  reg_weights = c(700, 35),
  # Shared pen SD = 40 g/day; simglm needs variance, which is SD squared.
  randomeffect = list(int_pen = list(variance = 40^2, var_level = 2)),
  # Animals also differ within their pens, with SD = 35 g/day.
  error = list(variance = 35^2),
  # Store the fitting function without calling it yet. Later, simglm calls
  # lmer() on each simulated dataset, using the formula above and REML.
  model_fit = list(model_function = lme4::lmer, REML = TRUE),
  # Count detections using a two-sided approximate normal test at 0.05.
  power = list(alpha = 0.05, dist = "qnorm", direction = "two-tailed"),
  # Try 100 new experiments for each plan. This is a rough demonstration.
  replications = 100
)

# 3. Inspect one invented trial before repeating anything --------------------
# A fixed seed lets us reproduce the same simulated examples when rerunning.
# Set it once: resetting it inside the repetitions would repeat the same trial.
set.seed(821029)
one_trial <- simglm(settings)
one_trial |>
  select(pen, diet, gain) |>
  slice_head(n = 3)

# Keep one row per pen before counting assignments. Counting animal rows
# instead would count measurements, rather than independently assigned pens.
one_trial |>
  distinct(pen, diet) |>
  count(diet)

# 4. List the four plans ----------------------------------------------------
# expand_grid() creates every combination of the two pen counts and two SDs.
# The order is 16 then 32 pens at SD 40, followed by 16 then 32 at SD 60.
# Keeping this order also keeps the random draws attached to the same plans.
plans <- expand_grid(pen_sd = c(40, 60), pens = c(16, 32)) |>
  select(pens, pen_sd) |>
  arrange(pen_sd, pens)
plans

# 5. Give each plan its own simulation settings ------------------------------
# map2() takes the first pen count and first SD together, then the second
# pair, and so on. function(pens, pen_sd) names the two values for one plan.
# The braces contain the steps to repeat for each pair.
# Copy settings first so changing one plan does not change the others.
plan_settings <- map2(plans$pens, plans$pen_sd, function(pens, pen_sd) {
  one_plan <- settings
  one_plan$sample_size$level2 <- pens
  one_plan$randomeffect$int_pen$variance <- pen_sd^2
  # Request one trial per call, so a fitting error does not stop all 100.
  one_plan$replications <- 1
  one_plan
})
# map2() returns a list: plan_settings[[1]] contains the first plan's settings.

# 6. Simulate and fit 100 trials for each plan --------------------------------
# The outer map() takes one set of plan settings at a time.
# The inner map() repeats the simulation for trial numbers 1 to 100.
# replicate_simulation() generates new animal gains AND fits lmer().
# It returns a list of models; [[1]] takes the one model requested above.
fit_attempts <- map(plan_settings, function(one_plan) {
  map(seq_len(settings$replications), function(trial_number) {
    tryCatch(
      replicate_simulation(one_plan)[[1]],
      # Save an error in the list instead of replacing the failed trial.
      # identity returns the error unchanged, so we can inspect it later.
      error = identity
    )
  })
})
# fit_attempts[[2]][[8]] is the fitted model, or error, for plan 2, trial 8.
# Inspect that entry if you need to investigate a failed trial.

# 7. Check the fits before using their estimates -----------------------------
# A fit is usable here if fitting converged and its estimates and standard
# errors are finite. A standard error describes uncertainty in an estimate.
# NULL marks a trial we cannot use. return(NULL) ends the checks for that
# trial only; it does not stop the other trials or the script.

# Define the checks for ONE trial first. fit is the model, or saved error,
# supplied to this function. The function returns its usable coefficient
# table, or NULL if that trial cannot contribute to the power calculation.
check_trial_fit <- function(fit) {
  # A failed fitting attempt contains an error instead of a model.
  # Display its explanation, then stop checking this particular trial.
  if (inherits(fit, "error")) {
    message("Fitting error: ", conditionMessage(fit))
    return(NULL)
  }

  # Check that the returned object is an lmer model and that its numerical
  # fitting converged before trying to use its estimates.
  if (!inherits(fit, "lmerMod")) {
    return(NULL)
  }
  if (!isTRUE(check_convergence(fit))) {
    return(NULL)
  }

  # Read the fitted coefficients, including the Supplement diet change.
  # If extraction fails, report the error and leave this trial out.
  coefficients <- tryCatch(extract_coefficients(fit), error = identity)
  if (inherits(coefficients, "error")) {
    message("Coefficient error: ", conditionMessage(coefficients))
    return(NULL)
  }

  # The next checks need these four columns. Stop if any are missing.
  required <- c("term", "estimate", "std.error", "statistic")
  if (!all(required %in% names(coefficients))) {
    return(NULL)
  }

  # Require exactly one diet-change row, finite estimates and test statistics,
  # and positive, finite standard errors. Missing or infinite values cannot
  # give a usable test of the assumed diet change.
  usable <- sum(coefficients$term == "dietSupplement") == 1L &&
    all(is.finite(coefficients$estimate)) &&
    all(is.finite(coefficients$std.error) & coefficients$std.error > 0) &&
    all(is.finite(coefficients$statistic))
  if (!usable) {
    return(NULL)
  }

  # Return the coefficient table for a usable fit. Give simglm the test
  # statistic so it applies the two-sided threshold itself.
  # In simglm 1.0.0, giving two-sided p-values instead halves alpha again.
  coefficients |> select(-any_of("p.value"))
}

# Now apply those checks. The outer map() handles one plan at a time.
# The inner map() passes each fitted trial to check_trial_fit by name.
usable_estimates <- map(fit_attempts, function(plan_fits) {
  checked <- map(plan_fits, check_trial_fit)
  # compact() removes the NULL entries. All attempts remain in fit_attempts,
  # and we count the removed entries as failures in the final table.
  compact(checked)
})

# 8. Calculate power separately for each plan --------------------------------
# map_dbl() repeats this calculation for each plan and returns one number
# per plan. Power is the fraction of usable trials detecting the assumed
# 35 g/day diet change with the test specified in settings$power.
power_values <- map_dbl(usable_estimates, function(estimates) {
  # If every trial failed, power cannot be estimated: keep it missing (NA).
  if (length(estimates) == 0) {
    return(NA_real_)
  }
  power <- compute_statistics(
    estimates,
    settings,
    power = TRUE,
    type_1_error = FALSE,
    precision = FALSE
  )
  # Keep the diet-change result, rather than the intercept result.
  power |>
    filter(term == "dietSupplement") |>
    pull(power)
})

# 9. Bring the results together and keep failures visible --------------------
# map_int(..., length) counts the usable fits for each plan.
# Multiplying power by usable recovers how many trials detected the change.
# If all fits failed, detections stay zero rather than becoming missing.
plans <- plans |>
  mutate(
    attempted = settings$replications,
    usable = map_int(usable_estimates, length),
    detected = if_else(usable == 0, 0, round(power_values * usable)),
    power_among_usable = power_values,
    failed = attempted - usable,
    # Simulation SE describes the noise from running only 100 trials.
    simulation_se = sqrt(
      power_among_usable * (1 - power_among_usable) / usable
    ),
    # A failed fit might have missed or detected the change. These two
    # bounds count every failure first as a miss, then as a detection.
    lower_if_failures_miss = detected / attempted,
    upper_if_failures_detect = (detected + failed) / attempted
  )
# Show every column when the table is displayed interactively.
options(tibble.width = Inf)
plans
# When all fits fail, the bounds are 0 to 1; power and simulation SE are NA.
# Investigate failed fits; do not replace them with new simulated trials.
# Near power 0.5, 100 trials give simulation SE about 0.05. A real planning
# decision needs more trials and justified assumptions about gains and spread.

# Save the table so it can be opened in R or shared with your study group.
dir.create("outputs", showWarnings = FALSE)
write_csv(plans, "outputs/power-plans.csv")