# Load necessary libraries
library(dplyr)
library(haven)

# Load Titanic dataset
titanic <- read_dta("https://github.com/scunning1975/mixtape/raw/master/titanic.dta") %>%
  mutate(
    d = ifelse(class == 1, 1, 0),  # Treatment indicator
    s = case_when(  # Stratify by sex and age
      sex == 0 & age == 1 ~ 1,  # Adult female
      sex == 0 & age == 0 ~ 2,  # Child female
      sex == 1 & age == 1 ~ 3,  # Adult male
      sex == 1 & age == 0 ~ 4,  # Child male
      TRUE ~ 0
    )
  )

# Step 1: Calculate differences in mean survival rates for each stratum
calculate_differences <- function(data, strata) {
  data %>%
    filter(s == strata) %>%
    summarize(
      ey11 = mean(survived[d == 1], na.rm = TRUE),  # Treated
      ey10 = mean(survived[d == 0], na.rm = TRUE),  # Control
      diff = ey11 - ey10  # Difference
    ) %>%
    mutate(stratum = strata)
}

differences <- bind_rows(lapply(1:4, function(i) calculate_differences(titanic, i)))

# Step 2: Calculate total observations for ATE, ATT, and ATU
n_total <- nrow(titanic)
n_treated <- nrow(titanic %>% filter(d == 1))
n_control <- nrow(titanic %>% filter(d == 0))

# Calculate stratum-specific counts
stratum_counts <- titanic %>%
  group_by(s, d) %>%
  summarize(count = n(), .groups = "drop") %>%
  pivot_wider(names_from = d, values_from = count, names_prefix = "d_") %>%
  mutate(
    total = d_1 + d_0,
    wt_ate = total / n_total,          # ATE weights
    wt_att = d_1 / n_treated,          # ATT weights
    wt_atu = d_0 / n_control           # ATU weights
  )

# Step 3: Merge weights with differences
weights_and_differences <- differences %>%
  left_join(stratum_counts, by = c("stratum" = "s"))

# Step 4: Calculate aggregate ATE, ATT, and ATU
ate <- sum(weights_and_differences$diff * weights_and_differences$wt_ate, na.rm = TRUE)
att <- sum(weights_and_differences$diff * weights_and_differences$wt_att, na.rm = TRUE)
atu <- sum(weights_and_differences$diff * weights_and_differences$wt_atu, na.rm = TRUE)

# Output results
list(
  ATE = ate,
  ATT = att,
  ATU = atu
)