library(tidyverse)
library(haven)

# Load the HIV dataset directly
hiv <- read_dta("https://raw.github.com/scunning1975/mixtape/master/thornton_hiv.dta")

# Function to calculate ATE with fixed 2,222 treated units
permuteHIV <- function(df, random = TRUE) {
  tb <- df
  
  if (random == TRUE) {
    # Randomly select exactly 2,222 respondents for treatment
    treatment_indices <- sample(nrow(tb), 2222)
    tb <- tb %>%
      mutate(any = if_else(row_number() %in% treatment_indices, 1, 0))
  }
  
  # Calculate mean outcome for treatment and control groups
  te1 <- tb %>%
    filter(any == 1) %>%
    pull(got) %>%
    mean(na.rm = TRUE)
  
  te0 <- tb %>%
    filter(any == 0) %>%
    pull(got) %>%
    mean(na.rm = TRUE)
  
  # Calculate ATE
  ate <- te1 - te0
  return(ate)
}

# Calculate observed ATE without randomization
observed_ate <- permuteHIV(hiv, random = FALSE)

# Set the number of permutations
iterations <- 1000

# Generate randomization distribution of ATEs
permutation <- tibble(
  iteration = 1:iterations,
  ate = c(observed_ate, map_dbl(2:iterations, ~permuteHIV(hiv, random = TRUE)))
)

# Calculate the p-value by comparing observed ATE with random ATEs
p_value <- mean(permutation$ate >= observed_ate)

# Display results
cat("Observed ATE:", observed_ate, "\n")
cat("P-value:", p_value, "\n")