suppressPackageStartupMessages({
  library(dplyr)
  library(tibble)
})

# estimate_standardized_vasopressor_mortality_effect() estimates a simple
# adjusted mortality effect using outcome regression and standardization.
#
# Standardization asks a target-trial-style question:
# "Among the same eligible ICU patients, what would the average 28-day mortality
# risk be if everyone followed the early vasopressor strategy versus if everyone
# followed the no early vasopressor strategy?"
#
# This is also called g-computation in this simple baseline-treatment setting.
# It is our first adjusted analysis, so it is intentionally simple and explicit.
#
# Arguments:
# - icu_data: a data frame or tibble with the columns created by
#   simulate_icu_cohort().
#
# Returns:
# - A named list with four pieces:
#   - eligible_icu_patients: only the target-trial eligible patients.
#   - mortality_outcome_model: the fitted logistic regression model.
#   - standardized_mortality_risks: average predicted 28-day mortality risk
#     under each treatment strategy.
#   - standardized_mortality_effect_estimates: standardized risk difference and
#     risk ratio.
#
# Example REPL use:
#
#   source("R/simulate_icu_cohort.R")
#   icu_data <- simulate_icu_cohort(n_patients = 1000, seed = 20260531)
#   standardized_analysis <- estimate_standardized_vasopressor_mortality_effect(icu_data)
#   standardized_analysis$standardized_mortality_effect_estimates
#
estimate_standardized_vasopressor_mortality_effect <- function(icu_data) {
  eligible_icu_patients <- icu_data |>
    filter(eligible == 1)

  # glm() fits a generalized linear model.
  # family = binomial() makes this a logistic regression for a 0/1 outcome.
  # The model adjusts for baseline severity variables that affect both treatment
  # decisions and mortality risk.
  mortality_outcome_model <- glm(
    death_28d ~ early_vasopressor + age + sex + sofa_score + lactate + map,
    data = eligible_icu_patients,
    family = binomial()
  )

  # Make two copies of the same eligible patients.
  # The only thing we change is the treatment strategy column.
  # This creates two counterfactual prediction datasets.
  eligible_patients_if_early_vasopressor <- eligible_icu_patients |>
    mutate(early_vasopressor = 1)

  eligible_patients_if_no_early_vasopressor <- eligible_icu_patients |>
    mutate(early_vasopressor = 0)

  # predict(..., type = "response") returns predicted probabilities from the
  # logistic regression model, not log-odds.
  predicted_mortality_if_early_vasopressor <- predict(
    mortality_outcome_model,
    newdata = eligible_patients_if_early_vasopressor,
    type = "response"
  )

  predicted_mortality_if_no_early_vasopressor <- predict(
    mortality_outcome_model,
    newdata = eligible_patients_if_no_early_vasopressor,
    type = "response"
  )

  standardized_mortality_risk_early_vasopressor <- mean(
    predicted_mortality_if_early_vasopressor
  )

  standardized_mortality_risk_no_early_vasopressor <- mean(
    predicted_mortality_if_no_early_vasopressor
  )

  standardized_mortality_risks <- tibble(
    treatment_strategy = c("Early vasopressor", "No early vasopressor"),
    patient_count = nrow(eligible_icu_patients),
    standardized_mortality_risk_28d = c(
      standardized_mortality_risk_early_vasopressor,
      standardized_mortality_risk_no_early_vasopressor
    )
  )

  standardized_mortality_effect_estimates <- tibble(
    estimate = c(
      "Standardized 28-day mortality risk, early vasopressor",
      "Standardized 28-day mortality risk, no early vasopressor",
      "Standardized risk difference",
      "Standardized risk ratio"
    ),
    value = c(
      standardized_mortality_risk_early_vasopressor,
      standardized_mortality_risk_no_early_vasopressor,
      standardized_mortality_risk_early_vasopressor - standardized_mortality_risk_no_early_vasopressor,
      standardized_mortality_risk_early_vasopressor / standardized_mortality_risk_no_early_vasopressor
    )
  )

  list(
    eligible_icu_patients = eligible_icu_patients,
    mortality_outcome_model = mortality_outcome_model,
    standardized_mortality_risks = standardized_mortality_risks,
    standardized_mortality_effect_estimates = standardized_mortality_effect_estimates
  )
}