suppressPackageStartupMessages({
  library(dplyr)
})

# estimate_iptw_and_ipcw_effect() computes combined IPTW and IPCW weights
# for a cloned target trial emulation dataset.
#
# IPTW (inverse probability of treatment weighting) balances baseline
# confounders across the two cloned treatment arms.
#
# IPCW (inverse probability of censoring weighting) accounts for clones
# that are censored due to protocol deviation or death before treatment.
#
# This function uses baseline confounders only for IPCW, keeping the
# teaching example simple and transparent.
#
# Arguments:
# - clone_baseline: output from clone_trial_arms()$clone_baseline.
#   One row per clone with baseline confounders, clone strategy,
#   censoring indicators, and 28-day outcome.
#
# Returns:
# - A named list with five pieces:
#   - propensity_score_model: logistic regression for treatment assignment.
#   - censoring_model: logistic regression for being censored at 2 hours.
#   - clone_weights: clone_baseline with propensity scores, IPTW weights,
#     IPCW weights, and combined weights appended.
#   - weight_diagnostics: summary of combined weights by clone strategy.
#   - weighted_effect_estimates: per-protocol mortality risks, RD, and RR.
#
# Example REPL use:
#
#   source("R/estimate_iptw_and_ipcw_effect.R")
#   effects <- estimate_iptw_and_ipcw_effect(clones$clone_baseline)
#   effects$weighted_effect_estimates
#
estimate_iptw_and_ipcw_effect <- function(clone_baseline) {
  # Step 1: Propensity score model.
  # We predict the probability of receiving early vasopressors in the
  # OBSERVED data using baseline confounders.
  # Because clone_baseline has two rows per patient, we deduplicate
  # by patient_id so the model is not fit on duplicate observations.
  unique_patients <- clone_baseline |>
    distinct(patient_id, .keep_all = TRUE)

  propensity_model <- glm(
    early_vasopressor ~ age + sex + sofa_score + lactate + map,
    data = unique_patients,
    family = binomial()
  )

  # Merge the predicted propensity score back onto the clone-level data.
  patient_propensity <- unique_patients |>
    mutate(
      propensity_score = predict(propensity_model, type = "response")
    ) |>
    select(patient_id, propensity_score)

  clone_weights <- clone_baseline |>
    left_join(patient_propensity, by = "patient_id")

  # Step 2: IPTW weights (unstabilized).
  # Early clone weight = 1 / P(early | confounders)
  # No_early clone weight = 1 / P(no_early | confounders)
  clone_weights <- clone_weights |>
    mutate(
      iptw_weight = case_when(
        clone_strategy == "early" ~ 1 / propensity_score,
        clone_strategy == "no_early" ~ 1 / (1 - propensity_score)
      )
    )

  # Step 3: Censoring model (baseline confounders only).
  # We predict the probability that a clone is censored at the 2-hour
  # mark, given its assigned strategy and baseline characteristics.
  censoring_model <- glm(
    clone_censored ~ clone_strategy + age + sex + sofa_score + lactate + map,
    data = clone_weights,
    family = binomial()
  )

  # Predicted probability of being censored, and therefore the weight
  # needed to remain uncensored.
  clone_weights <- clone_weights |>
    mutate(
      prob_censored = predict(censoring_model, type = "response"),
      prob_not_censored = 1 - prob_censored,
      ipcw_weight = 1 / prob_not_censored
    )

  # Step 4: Combined weights.
  clone_weights <- clone_weights |>
    mutate(
      combined_weight = iptw_weight * ipcw_weight
    )

  # Step 5: Weight diagnostics.
  # Extreme combined weights can indicate positivity problems in either
  # the treatment model or the censoring model.
  weight_diagnostics <- clone_weights |>
    group_by(clone_strategy) |>
    summarize(
      n_clones = n(),
      min_weight = min(combined_weight),
      max_weight = max(combined_weight),
      mean_weight = mean(combined_weight),
      median_weight = median(combined_weight),
      .groups = "drop"
    )

  # Step 6: Weighted effect estimates.
  # We compute the weighted 28-day mortality risk in each clone arm,
  # using only the clones that were NOT censored.
  # Each uncensored clone contributes according to its combined weight,
  # which accounts for both confounding and informative censoring.
  uncensored_clones <- clone_weights |>
    filter(clone_censored == 0)

  early_clones <- uncensored_clones |>
    filter(clone_strategy == "early")

  no_early_clones <- uncensored_clones |>
    filter(clone_strategy == "no_early")

  weighted_risk_early <- weighted.mean(
    early_clones$death_28d,
    early_clones$combined_weight
  )

  weighted_risk_no_early <- weighted.mean(
    no_early_clones$death_28d,
    no_early_clones$combined_weight
  )

  weighted_effect_estimates <- tibble(
    estimate = c(
      "IPTW + IPCW 28-day mortality risk, early vasopressor",
      "IPTW + IPCW 28-day mortality risk, no early vasopressor",
      "IPTW + IPCW risk difference",
      "IPTW + IPCW risk ratio"
    ),
    value = c(
      weighted_risk_early,
      weighted_risk_no_early,
      weighted_risk_early - weighted_risk_no_early,
      weighted_risk_early / weighted_risk_no_early
    )
  )

  list(
    propensity_score_model = propensity_model,
    censoring_model = censoring_model,
    clone_weights = clone_weights,
    weight_diagnostics = weight_diagnostics,
    weighted_effect_estimates = weighted_effect_estimates
  )
}