suppressPackageStartupMessages({
  library(dplyr)
})

# simulate_icu_cohort() is our first reusable project primitive.
#
# A "primitive" is a small function that does one useful job for the project.
# Here, the job is to create one simulated ICU observational cohort that we can
# reuse in scripts, notebooks, and later targets pipelines.
#
# Arguments:
# - n_patients: how many ICU patients to simulate.
# - seed: a number that makes the random simulation reproducible.
#
# Returns:
# - A tibble with one row per ICU patient.
# - The columns include baseline variables, eligibility indicators, observed
#   early vasopressor treatment, treatment timing, and 28-day mortality.
#
# Example REPL use:
#
#   source("R/simulate_icu_cohort.R")
#   icu_data <- simulate_icu_cohort(n_patients = 5, seed = 1)
#   icu_data
#
# Example output shape:
#
#   # A tibble: 5 x 13
#     patient_id   age sex    sofa_score lactate   map suspected_sepsis ...
#          <int> <dbl> <chr>       <dbl>   <dbl> <dbl>            <int> ...
#   1          1    56 female          8     2.9    59                1 ...
#
simulate_icu_cohort <- function(n_patients = 1000, seed = 20260531) {
  # set.seed() fixes the random-number stream.
  # That means the same inputs produce the same simulated dataset each time.
  set.seed(seed)

  # tibble() creates a modern data frame.
  # Each argument below becomes a column.
  # Each column must have either one value or n_patients values.
  tibble(
    # seq_len(n_patients) creates patient IDs 1, 2, ..., n_patients.
    patient_id = seq_len(n_patients),

    # rnorm() draws from a normal distribution.
    # round() makes age whole-number-like.
    # pmin() and pmax() cap ages to a plausible ICU range.
    age = round(rnorm(n_patients, mean = 65, sd = 14)) |>
      pmin(95) |>
      pmax(18),

    # sample() draws categorical values.
    # replace = TRUE means each patient gets an independent draw.
    sex = sample(
      c("female", "male"),
      size = n_patients,
      replace = TRUE,
      prob = c(0.45, 0.55)
    ),

    # rpois() draws count-like SOFA scores from a Poisson distribution.
    # pmin(20) keeps the simulated score within a plausible upper range.
    sofa_score = rpois(n_patients, lambda = 7) |>
      pmin(20),

    # rlnorm() creates right-skewed lactate values.
    # Clinical lab values often have this kind of skew.
    lactate = round(rlnorm(n_patients, meanlog = log(3), sdlog = 0.5), 1) |>
      pmin(15),

    # MAP is mean arterial pressure.
    # Lower MAP means more hypotension and greater shock severity.
    map = round(rnorm(n_patients, mean = 62, sd = 10)) |>
      pmin(95) |>
      pmax(35),

    # rbinom(..., size = 1) creates 0/1 indicators.
    # Here 1 means suspected sepsis is present at ICU admission.
    suspected_sepsis = rbinom(n_patients, size = 1, prob = 0.90)
  ) |>
    # mutate() adds new columns or changes existing columns.
    # These columns depend on the baseline variables created above.
    mutate(
      # as.integer(TRUE) is 1 and as.integer(FALSE) is 0.
      hypotension_at_baseline = as.integer(map < 65),
      elevated_lactate_at_baseline = as.integer(lactate >= 2),

      # This is the simulated eligibility definition for the target trial.
      # A patient is eligible only if all three baseline criteria are true.
      eligible = as.integer(
        suspected_sepsis == 1 &
          hypotension_at_baseline == 1 &
          elevated_lactate_at_baseline == 1
      ),

      # This is not a real clinical score.
      # It is a simulation device that makes sicker patients more likely to
      # receive early vasopressors and more likely to die.
      severity_score = 0.04 * (age - 65) +
        0.18 * (sofa_score - 7) +
        0.25 * (lactate - 3) -
        0.04 * (map - 62),

      # plogis() converts any real number into a probability between 0 and 1.
      # Higher severity_score gives a higher chance of early vasopressors.
      prob_early_vasopressor = plogis(-0.2 + severity_score),

      # This is the observed treatment group in the observational data.
      # It is not randomized; it depends on severity through the probability above.
      early_vasopressor = rbinom(n(), size = 1, prob = prob_early_vasopressor),

      # ifelse() chooses one value when the condition is TRUE and another when FALSE.
      # Early-treated patients get a time between 0 and 2 hours.
      # Not-early-treated patients get a time between 2 and 24 hours.
      time_to_vasopressor_hours = ifelse(
        early_vasopressor == 1,
        runif(n(), min = 0, max = 2),
        runif(n(), min = 2, max = 24)
      ) |>
        round(2),

      # This linear predictor controls each patient's mortality risk.
      # Sicker patients have higher risk; early vasopressor has a modest
      # protective effect in the data-generating process.
      mortality_linear_predictor = -1.4 +
        0.03 * (age - 65) +
        0.20 * (sofa_score - 7) +
        0.28 * (lactate - 3) -
        0.05 * (map - 62) -
        0.25 * early_vasopressor,

      # Convert the mortality linear predictor into a probability, then draw
      # the observed 28-day death indicator.
      prob_death_28d = plogis(mortality_linear_predictor),
      death_28d = rbinom(n(), size = 1, prob = prob_death_28d)
    ) |>
    # select() keeps the columns learners should analyze.
    # We intentionally drop helper columns like severity_score and probabilities.
    select(
      patient_id,
      age,
      sex,
      sofa_score,
      lactate,
      map,
      suspected_sepsis,
      hypotension_at_baseline,
      elevated_lactate_at_baseline,
      eligible,
      early_vasopressor,
      time_to_vasopressor_hours,
      death_28d
    )
}