library(tidyverse)
library(haven)
library(broom) # For tidying model outputs

read_data <- function(df) {
  full_path <- paste("https://raw.github.com/scunning1975/mixtape/master/", df, sep = "")
  df <- read_dta(full_path)
  return(df)
}

titanic <- read_data("titanic.dta") %>%
  mutate(d = if_else(class == 1, 1, 0))

ey1 <- titanic %>%
  filter(d == 1) %>%
  pull(survived) %>%
  mean()

ey0 <- titanic %>%
  filter(d == 0) %>%
  pull(survived) %>%
  mean()

sdo <- ey1 - ey0

# Run a regression of survived onto d
regression_result <- lm(survived ~ d, data = titanic)
tidy(regression_result)

# Generate dummies for the stratified variables according to the s variable
titanic <- titanic %>%
  mutate(
    female = as.numeric(sex == 0),
    male = as.numeric(sex == 1),
    s = case_when(
      female == 1 & age == 1 ~ 1,  # adult female
      female == 1 & age == 0 ~ 2,  # child female
      female == 0 & age == 1 ~ 3,  # adult male
      female == 0 & age == 0 ~ 4   # child male
    ),
    s1 = as.numeric(s == 1),
    s2 = as.numeric(s == 2),
    s3 = as.numeric(s == 3),
    s4 = as.numeric(s == 4)
  )

# Summarize counts and means by treatment status (d)
summary_by_d <- titanic %>%
  group_by(d) %>%
  summarize(
    count_s1 = sum(s1),
    mean_s1 = mean(s1),
    count_s2 = sum(s2),
    mean_s2 = mean(s2),
    count_s3 = sum(s3),
    mean_s3 = mean(s3),
    count_s4 = sum(s4),
    mean_s4 = mean(s4),
    .groups = 'drop'
  )

print(summary_by_d)