suppressPackageStartupMessages({
  library(dplyr)
  library(tidyr)
})

# simulate_icu_cohort_longitudinal() generates a sparse longitudinal dataset
# for the ICU septic shock teaching example.
#
# It reuses simulate_icu_cohort() for the baseline data-generating process so
# that cross-sectional and longitudinal results are directly comparable.
# Then it adds time-varying covariates, treatment timing, and censoring events
# at sparse time points.
#
# A "longitudinal" dataset has multiple rows per patient: one row for each
# time at which a variable was measured or an event could occur.
#
# Arguments:
# - n_patients: how many ICU patients to simulate.
# - seed: reproducibility seed.
#
# Returns:
# - A tibble with one row per patient per time point.
# - Time points: 0h, 2h, 6h, 12h, 24h, day7, day14, day21, day28.
# - Columns include baseline variables, time-varying covariates,
#   treatment indicators, and censoring/outcome variables.
#
# Example REPL use:
#
#   source("R/simulate_icu_cohort.R")
#   source("R/simulate_icu_cohort_longitudinal.R")
#   long_data <- simulate_icu_cohort_longitudinal(n_patients = 5, seed = 1)
#   long_data |> filter(patient_id == 1)
#
simulate_icu_cohort_longitudinal <- function(n_patients = 1000, seed = 20260531) {
  # Step 1: Generate baseline cross-sectional data.
  # simulate_icu_cohort() handles set.seed() and the full baseline DGP.
  # By calling it here, we guarantee that the baseline columns are identical
  # to those produced by the cross-sectional notebook workflows.
  patients <- simulate_icu_cohort(n_patients = n_patients, seed = seed)

  # Step 2: Add the longitudinal-specific variable.
  # Some 28-day deaths occur before the 2-hour treatment window closes.
  # This creates censoring for the per-protocol analysis.
  # Only patients who die by day 28 can die early; the probability is 15%.
  patients <- patients |>
    mutate(
      death_before_2h = ifelse(death_28d == 1, rbinom(n(), 1, 0.15), 0)
    )

  # Step 3: Define sparse time points for longitudinal follow-up.
  # We measure at clinically meaningful intervals during the first day,
  # then weekly until day 28.
  time_grid <- tibble(
    time_hours = c(0, 2, 6, 12, 24, 168, 336, 504, 672),
    time_label = c("0h", "2h", "6h", "12h", "24h", "day7", "day14", "day21", "day28")
  )

  # Step 4: Expand each patient to all time points.
  # crossing() creates the Cartesian product: every patient paired with every
  # time point. This is a standard tidyverse way to build a longitudinal grid.
  longitudinal <- patients |>
    crossing(time_grid) |>
    arrange(patient_id, time_hours)

  # Step 5: Add time-varying indicators.
  # These depend on the patient's actual treatment timing and baseline values.
  longitudinal <- longitudinal |>
    mutate(
      # Has vasopressor therapy started by this time point?
      vasopressor_started = as.integer(time_hours >= time_to_vasopressor_hours),

      # Time-varying MAP: improves modestly after vasopressors start.
      # this is a huge assumption/limitation to the simulation dataset, since some
      # patients will get better, some stay same, and some worse
      map_current = case_when(
        vasopressor_started == 1 ~ pmin(map + 5 + rnorm(n(), mean = 0, sd = 3), 95),
        TRUE ~ pmax(map - 1 + rnorm(n(), mean = 0, sd = 3), 35)
      ) |>
        round(1),

      # Time-varying lactate: decreases modestly after vasopressors start.
      lactate_current = case_when(
        vasopressor_started == 1 ~ pmax(lactate - 0.5 + rnorm(n(), mean = 0, sd = 0.3), 0.5),
        TRUE ~ pmin(lactate + 0.2 + rnorm(n(), mean = 0, sd = 0.3), 15)
      ) |>
        round(1),

      # Alive indicator: 0 if the patient has died by this time point.
      # Simplified: death_before_2h occurs at 2h; all other deaths at day 28.
      alive = case_when(
        death_before_2h == 1 & time_hours > 2 ~ 0,
        death_28d == 1 & time_hours > 672 ~ 0,
        TRUE ~ 1
      )
    )

  # Step 6: Select columns for the analytic dataset.
  # We drop intermediate helper columns (severity_score, prob_early_vasopressor,
  # mortality_linear_predictor, prob_death_28d) so learners focus on the
  # observable variables.
  longitudinal |>
    select(
      patient_id,
      time_hours,
      time_label,
      age,
      sex,
      sofa_score,
      lactate,
      map,
      lactate_current,
      map_current,
      suspected_sepsis,
      hypotension_at_baseline,
      elevated_lactate_at_baseline,
      eligible,
      early_vasopressor,
      time_to_vasopressor_hours,
      vasopressor_started,
      death_before_2h,
      death_28d,
      alive
    )
}