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

# Load the dataset directly from GitHub
ri_data_url <- "https://github.com/scunning1975/mixtape/raw/master/ri.dta"
ri_data <- read_dta(url(ri_data_url))
ri_data <- ri_data %>% mutate(id = row_number())  # Add unique ID column if not already present

# Generate combinations of IDs (k = 4)
combinations <- as.data.frame(t(combn(ri_data$id, 4)))
colnames(combinations) <- paste0("treated", 1:4)
combinations <- combinations %>% mutate(permutation = row_number())

# Cross join ri_data with combinations
cross_joined <- expand.grid(id = ri_data$id, permutation = combinations$permutation) %>%
  left_join(combinations, by = "permutation")

# Create the treatment indicator
cross_joined <- cross_joined %>%
  mutate(d = ifelse(id == treated1 | id == treated2 | id == treated3 | id == treated4, 1, 0))

# Ensure the treatment indicator (d) is correctly created in the cross_joined data frame
if ("d" %in% colnames(ri_data)) {
  ri_data <- ri_data %>% rename(d_ri = d)
}

# Merge cross_joined with original data and calculate grouped means
grouped <- cross_joined %>%
  left_join(ri_data, by = "id") %>%
  group_by(permutation, d) %>%
  summarize(mean_y = mean(y, na.rm = TRUE), .groups = "drop")

# Calculate average treatment effect (ATE) for each permutation
ate_data <- grouped %>%
  pivot_wider(names_from = d, values_from = mean_y, names_prefix = "d_") %>%
  mutate(ate = d_1 - d_0) %>%
  arrange(ate) %>%
  mutate(rank = row_number())

# Calculate p-value for observed permutation (permutation == 1)
observed_rank <- ate_data %>%
  filter(permutation == 1) %>%
  pull(rank)

p_value <- observed_rank / nrow(ate_data)

# Print the observed rank and p-value
print(paste("Observed rank:", observed_rank))
print(paste("Observed p-value:", p_value))


# Ensure ATE data is sorted correctly
ate_data <- grouped %>%
  pivot_wider(names_from = d, values_from = mean_y, names_prefix = "d_") %>%
  mutate(ate = d_1 - d_0) %>%
  arrange(ate) %>%  # Break ties by permutation ID
  mutate(rank = row_number())

# Validate observed permutation rank
observed_rank <- ate_data %>%
  filter(permutation == 1) %>%
  pull(rank)

# Compute p-value based on observed rank
p_value <- observed_rank / nrow(ate_data)

# Print results
print(paste("Observed rank:", observed_rank))
print(paste("Observed p-value:", p_value))