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

# estimate_iptw_vasopressor_mortality_effect() estimates a mortality effect
# using inverse probability of treatment weighting (IPTW).
#
# IPTW creates a pseudo-population in which treatment assignment is independent
# of the measured baseline confounders.
#
# This function uses unstabilized weights for clarity in a first teaching pass.
#
# Arguments:
# - icu_data: a data frame or tibble with the columns created by
#   simulate_icu_cohort().
#
# Returns:
# - A named list with six pieces:
#   - eligible_icu_patients: the eligible patients with propensity scores and
#     unstabilized IPTW weights appended.
#   - propensity_score_model: the fitted logistic regression for treatment.
#   - weight_summary: min, max, mean, and median weight by observed treatment
#     group.
#   - baseline_balance: a tibble comparing unweighted and weighted means for
#     each baseline confounder, by treatment group.
#   - weighted_mortality_risks: weighted 28-day mortality risk under each
#     treatment strategy.
#   - iptw_mortality_effect_estimates: weighted risk difference and risk ratio.
#
# Example REPL use:
#
#   source("R/simulate_icu_cohort.R")
#   icu_data <- simulate_icu_cohort(n_patients = 1000, seed = 20260531)
#   iptw_analysis <- estimate_iptw_vasopressor_mortality_effect(icu_data)
#   iptw_analysis$iptw_mortality_effect_estimates
#
estimate_iptw_vasopressor_mortality_effect <- function(icu_data) {
  # Keep only patients who meet the target trial eligibility criteria.
  eligible_icu_patients <- icu_data |>
    filter(eligible == 1)

  # Fit a logistic regression model for the probability of receiving early
  # vasopressors given baseline confounders.
  # This is the propensity score model.
  propensity_score_model <- glm(
    early_vasopressor ~ age + sex + sofa_score + lactate + map,
    data = eligible_icu_patients,
    family = binomial()
  )

  # predict(..., type = "response") gives the predicted probability of
  # early_vasopressor == 1 for each patient.
  propensity_score <- predict(
    propensity_score_model,
    newdata = eligible_icu_patients,
    type = "response"
  )

  # Unstabilized IPTW weights:
  # If treated, weight = 1 / P(treated | confounders).
  # If untreated, weight = 1 / P(untreated | confounders).
  iptw_weight <- ifelse(
    eligible_icu_patients$early_vasopressor == 1,
    1 / propensity_score,
    1 / (1 - propensity_score)
  )

  # Add the propensity score and weight back into the data for diagnostics
  # and later use.
  eligible_icu_patients <- eligible_icu_patients |>
    mutate(
      propensity_score = propensity_score,
      iptw_weight = iptw_weight
    )

  # Summarize the weight distribution by observed treatment group.
  # Extreme weights (very large values) can indicate positivity problems.
  weight_summary <- eligible_icu_patients |>
    group_by(early_vasopressor) |>
    summarize(
      patient_count = n(),
      min_weight = min(iptw_weight),
      max_weight = max(iptw_weight),
      mean_weight = mean(iptw_weight),
      median_weight = median(iptw_weight),
      .groups = "drop"
    ) |>
    mutate(
      treatment_group = factor(
        early_vasopressor,
        levels = c(0, 1),
        labels = c("No early vasopressor", "Early vasopressor")
      )
    ) |>
    select(treatment_group, patient_count, min_weight, max_weight, mean_weight, median_weight)

  # Build a baseline balance table that shows unweighted and weighted means
  # side by side. This lets learners see whether IPTW is rebalancing the
  # confounders across treatment groups.
  #
  # We compute weighted means manually with sum(x * w) / sum(w) so the code
  # is transparent and does not depend on an external weighted-summary package.

  early_idx <- eligible_icu_patients$early_vasopressor == 1
  no_early_idx <- eligible_icu_patients$early_vasopressor == 0

  w_early <- eligible_icu_patients$iptw_weight[early_idx]
  w_no_early <- eligible_icu_patients$iptw_weight[no_early_idx]

  weighted_mean <- function(x, w) {
    sum(x * w) / sum(w)
  }

  baseline_balance <- tibble(
    characteristic = c(
      "Age, years",
      "Male sex, %",
      "SOFA score",
      "Lactate, mmol/L",
      "MAP, mmHg"
    ),
    no_early_unweighted = c(
      mean(eligible_icu_patients$age[no_early_idx]),
      100 * mean(eligible_icu_patients$sex[no_early_idx] == "male"),
      mean(eligible_icu_patients$sofa_score[no_early_idx]),
      mean(eligible_icu_patients$lactate[no_early_idx]),
      mean(eligible_icu_patients$map[no_early_idx])
    ),
    no_early_weighted = c(
      weighted_mean(eligible_icu_patients$age[no_early_idx], w_no_early),
      100 * weighted_mean(eligible_icu_patients$sex[no_early_idx] == "male", w_no_early),
      weighted_mean(eligible_icu_patients$sofa_score[no_early_idx], w_no_early),
      weighted_mean(eligible_icu_patients$lactate[no_early_idx], w_no_early),
      weighted_mean(eligible_icu_patients$map[no_early_idx], w_no_early)
    ),
    early_unweighted = c(
      mean(eligible_icu_patients$age[early_idx]),
      100 * mean(eligible_icu_patients$sex[early_idx] == "male"),
      mean(eligible_icu_patients$sofa_score[early_idx]),
      mean(eligible_icu_patients$lactate[early_idx]),
      mean(eligible_icu_patients$map[early_idx])
    ),
    early_weighted = c(
      weighted_mean(eligible_icu_patients$age[early_idx], w_early),
      100 * weighted_mean(eligible_icu_patients$sex[early_idx] == "male", w_early),
      weighted_mean(eligible_icu_patients$sofa_score[early_idx], w_early),
      weighted_mean(eligible_icu_patients$lactate[early_idx], w_early),
      weighted_mean(eligible_icu_patients$map[early_idx], w_early)
    )
  )

  # Compute weighted 28-day mortality risk in each treatment group.
  # weighted.mean() from base R computes a weighted average.
  weighted_mortality_risk_early <- weighted.mean(
    eligible_icu_patients$death_28d[early_idx],
    w_early
  )

  weighted_mortality_risk_no_early <- weighted.mean(
    eligible_icu_patients$death_28d[no_early_idx],
    w_no_early
  )

  weighted_mortality_risks <- tibble(
    treatment_strategy = c("Early vasopressor", "No early vasopressor"),
    patient_count = nrow(eligible_icu_patients),
    weighted_mortality_risk_28d = c(
      weighted_mortality_risk_early,
      weighted_mortality_risk_no_early
    )
  )

  # Compute the IPTW effect estimates.
  iptw_mortality_effect_estimates <- tibble(
    estimate = c(
      "IPTW 28-day mortality risk, early vasopressor",
      "IPTW 28-day mortality risk, no early vasopressor",
      "IPTW risk difference",
      "IPTW risk ratio"
    ),
    value = c(
      weighted_mortality_risk_early,
      weighted_mortality_risk_no_early,
      weighted_mortality_risk_early - weighted_mortality_risk_no_early,
      weighted_mortality_risk_early / weighted_mortality_risk_no_early
    )
  )

  list(
    eligible_icu_patients = eligible_icu_patients,
    propensity_score_model = propensity_score_model,
    weight_summary = weight_summary,
    baseline_balance = baseline_balance,
    weighted_mortality_risks = weighted_mortality_risks,
    iptw_mortality_effect_estimates = iptw_mortality_effect_estimates
  )
}