suppressPackageStartupMessages({
  library(dplyr)
})

# clone_trial_arms() creates a cloned dataset for target trial emulation.
#
# At time zero (ICU admission), each eligible patient is duplicated into two
# hypothetical treatment arms:
#   - "early"    = assigned to start vasopressors within 2 hours
#   - "no_early" = assigned to not start vasopressors within 2 hours
#
# Each clone is then followed over the patient's actual observed trajectory.
# A clone is censored if the patient deviated from the assigned strategy
# (protocol deviation) or died before the 2-hour treatment window closed.
#
# Arguments:
# - longitudinal_data: output from simulate_icu_cohort_longitudinal().
#
# Returns:
# - A named list with two pieces:
#   - clone_baseline: one row per clone at time 0. This is the main analysis
#     dataset for weighting and effect estimation.
#   - clone_longitudinal: the full longitudinal data for each clone. Useful
#     for teaching trajectories and visualising when censoring occurs.
#
# Example REPL use:
#
#   source("R/simulate_icu_cohort_longitudinal.R")
#   source("R/clone_trial_arms.R")
#   long_data <- simulate_icu_cohort_longitudinal(1000, 20260531)
#   clones <- clone_trial_arms(long_data)
#   clones$clone_baseline |> count(clone_strategy)
#
clone_trial_arms <- function(longitudinal_data) {
  # Keep only patients who meet the target trial eligibility criteria.
  eligible_ids <- longitudinal_data |>
    filter(time_hours == 0, eligible == 1) |>
    distinct(patient_id) |>
    pull(patient_id)

  eligible_longitudinal <- longitudinal_data |>
    filter(patient_id %in% eligible_ids)

  # Extract patient-level baseline information (constant across time).
  # We use only the time 0 row to avoid duplicate rows per patient.
  patient_summary <- eligible_longitudinal |>
    filter(time_hours == 0) |>
    select(
      patient_id,
      age,
      sex,
      sofa_score,
      lactate,
      map,
      early_vasopressor,
      time_to_vasopressor_hours,
      death_before_2h,
      death_28d
    )

  # Step 1: Create two clones per eligible patient at baseline.
  # Each clone represents one arm of the target trial.
  clone_baseline <- patient_summary |>
    crossing(clone_strategy = c("early", "no_early")) |>
    mutate(
      # Censoring rules for the per-protocol analysis:
      #
      # - "early" clone: censored if the patient did NOT actually start
      #   vasopressors within 2 hours (protocol deviation).
      #
      # - "no_early" clone: censored if the patient DID actually start
      #   vasopressors within 2 hours (protocol deviation).
      #
      # - Both clones: censored if the patient died before the 2-hour
      #   treatment window closed (death before treatment).
      clone_censored = case_when(
        clone_strategy == "early" & early_vasopressor == 0 ~ 1,
        clone_strategy == "no_early" & early_vasopressor == 1 ~ 1,
        death_before_2h == 1 ~ 1,
        TRUE ~ 0
      ),
      clone_censored_reason = case_when(
        clone_strategy == "early" & early_vasopressor == 0 ~ "protocol_deviation",
        clone_strategy == "no_early" & early_vasopressor == 1 ~ "protocol_deviation",
        death_before_2h == 1 ~ "death_before_treatment",
        TRUE ~ "none"
      )
    )

  # Step 2: Expand each clone to all time points from the original data.
  # A clone is "at risk" only up to the point of censoring.
  # After censoring, the clone drops out and does not contribute to
  # later time points.
  clone_longitudinal <- eligible_longitudinal |>
    left_join(
      clone_baseline |>
        select(patient_id, clone_strategy, clone_censored, clone_censored_reason),
      by = "patient_id"
    ) |>
    mutate(
      clone_at_risk = case_when(
        clone_censored == 1 & time_hours > 2 ~ 0,
        TRUE ~ 1
      )
    )

  list(
    clone_baseline = clone_baseline,
    clone_longitudinal = clone_longitudinal
  )
}