pdf(NULL) # suppress Rplots.pdf in batch mode
library(here)
library(tidyverse)
library(broom)
library(ggthemes)
library(sandwich)
library(lmtest)
library(kableExtra)
library(furrr)
plan(multicore, workers = min(16, parallelly::availableCores()))
set.seed(1324)
EPSILON <- 10^-10

##############################################
stratified_mean <- function(data, strat_vars, weights) {
  if(length(strat_vars) > 0) {
    summary_df <- suppressMessages(data %>%
      group_by(across(c(all_of(strat_vars)))) %>%
      summarise(across(c(z_score, score), ~ mean(.x), .names="mean_{.col}"),
                across(c(z_score, score), ~ sd(.x)/sqrt(n()), .names="se_{.col}")
                ) %>%
      ungroup() %>%
      inner_join(weights, by=strat_vars) %>%
      summarise(across(c(mean_z_score, mean_score), ~ weighted.mean(.x, weight)),
                across(c(se_z_score, se_score), ~ sqrt(sum((.x**2)*(weight**2))))) %>%
      ungroup()
    )
  } else {
    summary_df <- suppressMessages(data %>%
      summarise(across(c(z_score, score), ~ mean(.x), .names="mean_{.col}"),
                across(c(z_score, score), ~ sd(.x)/sqrt(n()), .names="se_{.col}")))
  }
  return(summary_df)
}

boot_fn_clustered <- function(data, strat_vars, cluster_var, weights, B = 300) {
  unique_clusters <- unique(data[[cluster_var]])
  
  # Perform B bootstrap iterations
  1:B %>%
    # For each iteration, sample clusters with replacement
    map(~ sample(unique_clusters, replace = TRUE)) %>%
    # For each bootstrap sample, extract the rows corresponding to the sampled clusters
    future_map(~ {
      sampled_indices <- which(data[[cluster_var]] %in% .x)
      stratified_mean(data[sampled_indices,], strat_vars, weights)
    }) %>%
    # put all outputs from stratified_mean into a single df
    list_rbind() %>%
    # Calculate the standard error
    summarise(se_clustered_z_score=sd(mean_z_score),
              se_clustered_score=sd(mean_score))
}

# function that returns a tibble with stratified means and bootstrapped ses
stratified_mean_boot_cluster_se <- function(data, strat_vars, cluster_var, weights, B = 300) {
  strat_means <- stratified_mean(data, strat_vars, weights)
  boot_ses <- boot_fn_clustered(data, strat_vars, cluster_var, weights, B)
  combined_res <- bind_cols(strat_means, boot_ses)
  return(combined_res)
}

################################################################################
######## Prepare data
repl_map <- c("2"                 = "2",
              "3-verbatim-prompt" = "3",
              "3-revised-prompt"  = "3revised")
image_file <- here("logo_exp_replication", "inputs", "main_df_logoExp_promptLvl_2026-03-22_PROD_HASHED_TRIMMED.csv")
image_df <- read_csv(image_file) %>%
  mutate(dalleVersion = str_replace_all(dalleVersion, repl_map))
table(image_df$dalleVersion)


repl_map <- c("dalle3verbatim" = "3",
              "dalle3revised"  = "3revised",
              "dalle2"         = "2") 
replay_file <- here("logo_exp_replication", "inputs", "logo_experiment_mislabel_excluded_image_rankings_20_perc_TRIMMED.csv")
replay_df_full <- read_csv(replay_file) %>%
  separate(col  = image_id, into = c("image_id", "replayDalleVersion"), sep  = "_") %>%
  mutate(replayDalleVersion = str_replace_all(replayDalleVersion, repl_map),
         bt_score=pmax(strength, EPSILON),
         strength=log(bt_score),
         fraction_win=num_wins/(num_wins+num_losses))

orig_size <- nrow(replay_df_full)
orig_images <- n_distinct(image_df$image_id)
replay_df_full <- replay_df_full %>%
  inner_join(image_df, by='image_id')
if(nrow(replay_df_full) != orig_size) {
  cat('Lost', orig_size - nrow(replay_df_full), 'replay images after merging with image df')
}
if(n_distinct(replay_df_full$image_id) != orig_images) {
  cat('Lost', orig_images - n_distinct(replay_df_full$image_id), 'unique images after merging with image df')
}

replay_df_full <- replay_df_full %>%
  filter(iteration < 10) %>%
  filter(!is.na(strength))

# remove prompts that were not successfully replayed across all models
orig_size <- nrow(replay_df_full)
orig_images <- n_distinct(image_df$image_id)
orig_users <- n_distinct(replay_df_full$userId)
bad_images <- replay_df_full %>% group_by(image_id) %>% summarise(valid_replays=n()) %>%
  filter(valid_replays < 3) %>% pull(image_id)
replay_df_full <- replay_df_full %>% filter(!image_id %in% bad_images)
if(nrow(replay_df_full) != orig_size) {
  cat('Lost', orig_size - nrow(replay_df_full), 'replay images after removing those without all replay models')
}
if(n_distinct(replay_df_full$image_id) != orig_images) {
  cat('Lost', orig_images - n_distinct(replay_df_full$image_id), 'unique images after removing those without all replay models')
}
if(n_distinct(replay_df_full$userId) != orig_users) {
  cat('Lost', orig_users - n_distinct(replay_df_full$userId), 'unique users after removing images without all replay models')
}

strat_vars <- c('logo_task', 'iteration')
score_var <- 'strength'
# calculate z-score for comparison between these two models
replay_df_full <- replay_df_full %>%
  mutate(score=.data[[score_var]]) %>%
  group_by(across(all_of(strat_vars))) %>%
  mutate(z_score = scale(score)[,1]) %>%
  ungroup()

# chi-square test of balance across all tasks/dalle conditions
tmp <- replay_df_full %>%
  group_by(userId) %>%
  summarise(dalleVersion=first(dalleVersion),
            logo_task=first(logo_task)) %>%
  group_by(dalleVersion, logo_task) %>%
  summarise(cnt=n())
chisq.test(tmp$cnt)
ggplot(replay_df_full, aes(score, fill=replayDalleVersion)) +
  geom_histogram() +
  scale_x_log10()

################################################################################
######## Mediation analysis through replay (direct/indirect/ATE decomposition in a plot)
strat_vars <- c('logo_task', 'iteration')
# subset the data to the two main models
#valid_models <- c('2', '3', '3revised')
#valid_models <- c('2', '3revised')
valid_models <- c('2', '3')
replay_df <- replay_df_full %>%
  filter(dalleVersion %in% valid_models &
         replayDalleVersion %in% valid_models)

cat("🔄 Running bootstrap (1/3) — please be patient, this may take ~3 min...\n")
# stratified averages of replay scores across all iterations
weights <- replay_df %>%
  group_by(across(all_of(strat_vars))) %>%
  summarise(cnt=n()) %>%
  ungroup() %>%
  mutate(weight=cnt/sum(cnt)) %>%
  select(-cnt)

summary_df <- replay_df %>%
  group_by(dalleVersion, replayDalleVersion) %>%
  group_modify(~ stratified_mean_boot_cluster_se(.x, strat_vars, 'userId', weights)) %>%
  ungroup()
summary_df

end <- summary_df[summary_df$dalleVersion==3 & summary_df$replayDalleVersion==3,]$mean_score
mid <- summary_df[summary_df$dalleVersion==2 & summary_df$replayDalleVersion==3,]$mean_score
start <- summary_df[summary_df$dalleVersion==2 & summary_df$replayDalleVersion==2,]$mean_score
(mid - start) / (end - start)

ggplot(summary_df, aes(replayDalleVersion, mean_z_score, color=dalleVersion)) +
  geom_point(position=position_dodge(width=0.3), size=2) + 
  geom_errorbar(aes(ymin=mean_z_score - 2*se_z_score,
                    ymax=mean_z_score + 2*se_z_score),
                linewidth=1, width=0.2, position = position_dodge(width = 0.3)) +
  labs(y=paste('mean z socre of', score_var))

# Plot the progression from base model to forward direct and indirect effect and then back
# z-score
base_model <- '2'
base_row <- summary_df[summary_df$dalleVersion == base_model & summary_df$replayDalleVersion == base_model, ]
normalized_summary_df <- summary_df %>%
  mutate(mean_z_score = mean_z_score - base_row$mean_z_score,
         se_z_score = se_z_score,
         se_clustered_z_score = se_clustered_z_score,
         mean_score = mean_score - base_row$mean_score,
         se_score = se_score,
         se_clustered_score = se_clustered_score,
         label=case_when(
           dalleVersion == base_model & replayDalleVersion > base_model ~ 'DALL-E 2\nDALL-E 3',
           dalleVersion > base_model & replayDalleVersion > base_model ~ 'DALL-E 3\nDALL-E 3',
           dalleVersion > base_model & replayDalleVersion == base_model ~ 'DALL-E 3\nDALL-E 2',
           dalleVersion == base_model & replayDalleVersion == base_model ~ 'DALL-E 2\nDALL-E 2',
           TRUE ~ "Something is wrong")) %>%
  #filter(dalleVersion != base_model | replayDalleVersion != base_model) %>%
  mutate(label=factor(label,
                      levels = c('DALL-E 2\nDALL-E 2', 'DALL-E 2\nDALL-E 3', 'DALL-E 3\nDALL-E 3', 'DALL-E 3\nDALL-E 2'),
                      ordered = TRUE)) %>%
  arrange(label) %>%
  mutate(arrow_z_y_start=ifelse(label!='DALL-E 2\nDALL-E 2', lag(mean_z_score, default=0), NA),
         arrow_y_start=ifelse(label!='DALL-E 2\nDALL-E 2', lag(mean_score, default=0), NA),
         arrow_z_y_end=ifelse(label!='DALL-E 2\nDALL-E 2', mean_z_score, NA),
         arrow_y_end=ifelse(label!='DALL-E 2\nDALL-E 2', mean_score, NA),
         arrow_x=order(label) - 0.25,
         next_arrow_x=ifelse(label!='DALL-E 2\nDALL-E 2', lead(arrow_x, default = max(order(label))), NA),
         effect=case_when(
           dalleVersion == base_model & replayDalleVersion > base_model ~ 'Model Effect',
           dalleVersion > base_model & replayDalleVersion == base_model ~ 'Model Effect',
           dalleVersion > base_model & replayDalleVersion > base_model ~ 'Prompting Effect',
           TRUE ~ "Something is wrong"),
         effect_type=case_when(
           grepl('^Model', effect) ~ 'Direct',
           grepl('^Prompt', effect) ~ 'Indirect',
           TRUE ~ effect))
cat(paste('Percentage changes with respect to base model:', base_model))
summary_df %>%
  select(dalleVersion, replayDalleVersion, mean_z_score, mean_score) %>%
  mutate(pct_change_z_score = (mean_z_score/base_row$mean_z_score - 1)*100,
         pct_change_score = (mean_score/base_row$mean_score - 1)*100)

full_effect_arrow <- data.frame(arrow_x=max(normalized_summary_df$arrow_x) + 0.6,
                                next_arrow_x=normalized_summary_df[which.max(normalized_summary_df$arrow_y_end),]$next_arrow_x,
                                arrow_z_y_start=min(normalized_summary_df$arrow_z_y_start, na.rm=TRUE),
                                arrow_y_start=min(normalized_summary_df$arrow_y_start, na.rm=TRUE),
                                arrow_z_y_end=max(normalized_summary_df$arrow_z_y_end, na.rm=TRUE),
                                arrow_y_end=max(normalized_summary_df$arrow_y_end, na.rm=TRUE),
                                effect_type='Full',
                                effect='ATE')
normalized_summary_df <- bind_rows(normalized_summary_df, full_effect_arrow)
normalized_summary_df
dir.create(here("logo_exp_replication", "outputs", "data"), recursive = TRUE, showWarnings = FALSE)
dir.create(here("logo_exp_replication", "outputs", "figures"), recursive = TRUE, showWarnings = FALSE)
dir.create(here("logo_exp_replication", "outputs", "tables"), recursive = TRUE, showWarnings = FALSE)
write_csv(normalized_summary_df,
          here("logo_exp_replication", "outputs", "data",
               paste0('writing_task_model_prompt_effect_dalle_', paste(valid_models, collapse='_'), '.csv')))

error_bar_size <- 1
point_size <- 1.1
axis_text_size <- 17
axis_title_size <- 17 
legend_text_size <- 17
p <- ggplot(normalized_summary_df) +
  geom_hline(yintercept = 0, size = 2, color = "white") +
  geom_pointrange(data=normalized_summary_df %>% filter(!is.na(label)),
                  aes(x=label, y=mean_z_score,
                      ymin = mean_z_score - 2 * se_clustered_z_score,
                      ymax = mean_z_score + 2 * se_clustered_z_score),
                size=point_size, linewidth=error_bar_size, na.rm=TRUE) +
  geom_hline(yintercept = 0, linetype='dotted', linewidth=1) +
  geom_segment(aes(x = arrow_x, xend = arrow_x, y = arrow_z_y_start, yend = arrow_z_y_end, color=effect_type),
               arrow = arrow(length = unit(0.3, "cm")), linewidth = error_bar_size) +
  geom_text(aes(x = arrow_x, y = (arrow_z_y_start + arrow_z_y_end) / 2, label = effect, color=effect_type),
            nudge_x=-0.13, nudge_y = 0.00, angle = 90, hjust = 0.5, size = axis_text_size-3, size.unit = 'pt') +
  geom_segment(aes(x = arrow_x, xend = next_arrow_x, y = arrow_z_y_end, yend = arrow_z_y_end),
               linetype='dotted', linewidth=1) +
  scale_color_manual(values=c('Direct'='indianred', 'Indirect'='dodgerblue', 'Full'='black'), guide='none') +
  coord_cartesian(x=c(1.45,3.9)) +
  theme_clean() +
  labs(x='Replay Model', y=expression(Delta ~ 'Z-Score'), color='Prompted Model') +
  theme(axis.text.y=element_text(colour="black", size=axis_text_size),
        axis.text.x=element_text(colour="black", size=axis_text_size, angle = 0),
        axis.title.y=element_text(colour="black", size=axis_title_size),
        axis.title.x=element_blank(),
        legend.title=element_text(size=axis_title_size),
        legend.text=element_text(size=legend_text_size),
        plot.background = element_blank())
p

# raw score
p <- ggplot(normalized_summary_df) +
  geom_hline(yintercept = 0, size = 2, color = "white") +
  geom_pointrange(data=normalized_summary_df %>% filter(!is.na(label)),
                  aes(x=label, y=mean_score,
                      ymin = mean_score - 2 * se_clustered_score,
                      ymax = mean_score + 2 * se_clustered_score),
                size=point_size, linewidth=error_bar_size, na.rm=TRUE) +
  geom_hline(yintercept = 0, linetype='dotted', linewidth=1) +
  geom_segment(aes(x = arrow_x, xend = arrow_x, y = arrow_y_start, yend = arrow_y_end, color=effect_type),
               arrow = arrow(length = unit(0.3, "cm")), linewidth = error_bar_size) +
  geom_text(aes(x = arrow_x, y = (arrow_y_start + arrow_y_end) / 2, label = effect, color=effect_type),
            nudge_x=-0.11, nudge_y = 0.0002, angle = 90, hjust = 0.5, size = axis_text_size-3, size.unit = 'pt') +
  geom_segment(aes(x = arrow_x, xend = next_arrow_x, y = arrow_y_end, yend = arrow_y_end),
               linetype='dotted', linewidth=1) +
  scale_color_manual(values=c('Direct'='indianred', 'Indirect'='dodgerblue', 'Full'='black'), guide='none') +
  #coord_cartesian(x=c(1.45,3.9), y=c(-0.003,0.0195)) +
  theme_clean() +
  #scale_y_continuous(minor_breaks = seq(0.005 , 0.02, 0.005)) +
  labs(x='Replay Model', y=expression(Delta ~ 'Score'), color='Prompted Model') +
  theme(axis.text.y=element_text(colour="black", size=axis_text_size),
        axis.text.x=element_text(colour="black", size=axis_text_size, angle = 0),
        axis.title.y=element_text(colour="black", size=axis_title_size),
        axis.title.x=element_blank(),
        legend.title=element_text(size=axis_title_size),
        legend.text=element_text(size=legend_text_size),
        plot.background = element_blank())
p
ggsave(here("logo_exp_replication", "outputs", "figures",
            paste0('replay_decomposition_base_model_', base_model, '_task2.pdf')),
       p, width=8, height=6, units='in')

ate_df <- summary_df %>% cross_join(summary_df, suffix = c("", "1")) %>%
  mutate(ate_z_score = mean_z_score - mean_z_score1,
         se_clustered_ate_z_score = sqrt(se_clustered_z_score**2 + se_clustered_z_score1**2),
         p_ate_z_score = 2*(1-pnorm(abs(ate_z_score), sd=se_clustered_ate_z_score)),
         ate_z_score_lower_ci = ate_z_score - qnorm(0.975)*se_clustered_ate_z_score,
         ate_z_score_upper_ci = ate_z_score + qnorm(0.975)*se_clustered_ate_z_score,
         ate_score = mean_score - mean_score1,
         se_clustered_ate_score = sqrt(se_clustered_score**2 + se_clustered_score1**2),
         p_ate_score = 2*(1-pnorm(abs(ate_score), sd=se_clustered_ate_score)),
         ate_score_lower_ci = ate_score - qnorm(0.975)*se_clustered_ate_score,
         ate_score_upper_ci = ate_score + qnorm(0.975)*se_clustered_ate_score,
         ) %>%
  select(-matches('^mean_z_score*|^se_z_score*|^mean_score*|^se_score*|se_clustered_z_score*|se_clustered_score*')) %>%
  filter((dalleVersion1 == base_model    & replayDalleVersion1 == base_model &
          dalleVersion  == dalleVersion1 & replayDalleVersion  != replayDalleVersion1) |     # direct effect
         (dalleVersion1 == base_model    & replayDalleVersion != base_model &
          dalleVersion  != dalleVersion1 & replayDalleVersion == replayDalleVersion1) |      # indirect effect
         (dalleVersion1 == base_model    & replayDalleVersion1 == base_model &
          dalleVersion  != dalleVersion1 & replayDalleVersion  != replayDalleVersion1)) %>%  # full effect
  mutate(label=case_when(
           (dalleVersion != dalleVersion1 & replayDalleVersion != replayDalleVersion1) ~ 'Total Effect\n(Model & Prompting)', 
           dalleVersion > base_model ~ 'Indirect Effect\n(Prompting)', 
           dalleVersion < base_model ~ 'Indirect Effect\n(Prompting)',
           replayDalleVersion < base_model ~ 'Direct Effect\n(Model)',
           replayDalleVersion > base_model ~ 'Direct Effect\n(Model)',
           TRUE ~ "Something is wrong"))
ate_df %>% select(c(label, ate_score, se_clustered_ate_score, ate_score_lower_ci, ate_score_upper_ci, p_ate_score))


################################################################################
##### Forward/Backward ATE/Direct/Indirect estimation using a linear model with Cluster robust SE
# reverse these two base/other models to get backward (change base to 2 and other to 1)
base_model <- sort(unique(replay_df$dalleVersion))[1]
other_model <- sort(unique(replay_df$dalleVersion))[2]
#base_model <- sort(unique(replay_df$dalleVersion))[2]
#other_model <- sort(unique(replay_df$dalleVersion))[1]
cat('Using base model of', base_model, 'and other model of', other_model)

# which effect type?
#effect_type <- 'direct'
#effect_type <- 'indirect'
effect_type <- 'total'
tmp_data <- replay_df %>%
  filter((effect_type=='direct' & dalleVersion==base_model) |
         (effect_type=='indirect' & replayDalleVersion==other_model) |
         (effect_type=='total' & replayDalleVersion==dalleVersion)) %>%
  mutate(
    userId=as.factor(userId),
    iteration=as.factor(iteration),
    effect=case_when(
      effect_type=='direct' ~ ifelse(replayDalleVersion==base_model, 0, 1),
      effect_type=='indirect' ~ ifelse(dalleVersion==base_model, 0, 1),
      effect_type=='total' ~ ifelse(dalleVersion==base_model, 0, 1),
      TRUE ~ -1)) %>%
  select(userId, dalleVersion, replayDalleVersion, logo_task, iteration, effect, score, z_score, bt_score)

# z-score results
model_fit <- lm(z_score ~ iteration + logo_task + effect, tmp_data)
model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat(paste('Regression results of z_score for', effect_type, 'effect with cluster robust standard errors:'))
model_summary_robust
cat('95% confidence interval for z_score with cluster robust standard errors:')
cat("(",
    model_summary_robust["effect", "Estimate"] - 1.96*model_summary_robust["effect", "Std. Error"],
    " , ",
    model_summary_robust["effect", "Estimate"] + 1.96*model_summary_robust["effect", "Std. Error"], ")")

# raw score results
model_fit <- lm(score ~ iteration + logo_task + effect, tmp_data)
model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat(paste('Regression results of score for', effect_type, 'effect with cluster robust standard errors:'))
model_summary_robust
cat('95% confidence interval for score with cluster robust standard errors:')
cat("(",
    model_summary_robust["effect", "Estimate"] - 1.96*model_summary_robust["effect", "Std. Error"],
    " , ",
    model_summary_robust["effect", "Estimate"] + 1.96*model_summary_robust["effect", "Std. Error"], ")")

### robustness check using stratification: main results in the paper are based on regression above
ate_df <- summary_df %>% cross_join(summary_df, suffix = c("", "1")) %>%
  mutate(ate_z_score = mean_z_score - mean_z_score1,
         se_clustered_ate_z_score = sqrt(se_clustered_z_score**2 + se_clustered_z_score1**2),
         p_ate_z_score = 2*(1-pnorm(abs(ate_z_score), sd=se_clustered_ate_z_score)),
         ate_z_score_lower_ci = ate_z_score - qnorm(0.975)*se_clustered_ate_z_score,
         ate_z_score_upper_ci = ate_z_score + qnorm(0.975)*se_clustered_ate_z_score,
         ate_score = mean_score - mean_score1,
         se_clustered_ate_score = sqrt(se_clustered_score**2 + se_clustered_score1**2),
         p_ate_score = 2*(1-pnorm(abs(ate_score), sd=se_clustered_ate_score)),
         ate_score_lower_ci = ate_score - qnorm(0.975)*se_clustered_ate_score,
         ate_score_upper_ci = ate_score + qnorm(0.975)*se_clustered_ate_score,
         ) %>%
  select(-matches('^mean_z_score*|^se_z_score*|^mean_score*|^se_score*|se_clustered_z_score*|se_clustered_score*')) %>%
  filter((dalleVersion1 == base_model    & replayDalleVersion1 == base_model &
          dalleVersion  == dalleVersion1 & replayDalleVersion  != replayDalleVersion1) |     # direct effect
         (dalleVersion1 == base_model    & replayDalleVersion != base_model &
          dalleVersion  != dalleVersion1 & replayDalleVersion == replayDalleVersion1) |      # indirect effect
         (dalleVersion1 == base_model    & replayDalleVersion1 == base_model &
          dalleVersion  != dalleVersion1 & replayDalleVersion  != replayDalleVersion1)) %>%  # full effect
  mutate(label=case_when(
           (dalleVersion != dalleVersion1 & replayDalleVersion != replayDalleVersion1) ~ 'Total Effect\n(Model & Prompting)', 
           dalleVersion > base_model ~ 'Indirect Effect\n(Prompting)', 
           dalleVersion < base_model ~ 'Indirect Effect\n(Prompting)',
           replayDalleVersion < base_model ~ 'Direct Effect\n(Model)',
           replayDalleVersion > base_model ~ 'Direct Effect\n(Model)',
           TRUE ~ "Something is wrong"))

# results based on raw-score
ate_df %>% select(c(label, ate_score, se_clustered_ate_score, ate_score_lower_ci, ate_score_upper_ci, p_ate_score))
# pvalue of comparison between total effect and backward direct effect
total_ate <- ate_df %>% filter(grepl("Total", label))
direct_ate <- ate_df %>% filter(grepl("Direct", label))
paste('Difference:',
      total_ate$ate_score - direct_ate$ate_score ,' ',
      'and P-value of comparing (backward or forward) total and direct effect:',
      2*(1- pnorm(abs(total_ate$ate_score - direct_ate$ate_score),
                  sd=sqrt(total_ate$se_clustered_ate_score**2 + direct_ate$se_clustered_ate_score**2))))

################################################################################
##### Direct/Indirect/Total effect by quantile/percentile in distribution using a linear model with Cluster robust SE
strat_vars <- c('logo_task', 'iteration')
base_model <- sort(unique(replay_df$dalleVersion))[1]
other_model <- sort(unique(replay_df$dalleVersion))[2]
#base_model <- sort(unique(replay_df$dalleVersion))[2]
#other_model <- sort(unique(replay_df$dalleVersion))[1]
cat('Using base model of', base_model, 'and other model of', other_model)

# add the performance rank of each user within its own reference image and iteration
pct_groups <- seq(1,100,2)
find_closest <- function(x) {
  pct_groups[which.min(abs(pct_groups - x))]
}
tmp_data <- replay_df %>%
  group_by(across(c(all_of(strat_vars), dalleVersion, replayDalleVersion))) %>%
  mutate(score_pct_rank=percent_rank(score)*100,
         score_pct_rank_group = sapply(score_pct_rank, find_closest)) %>%
  ungroup() %>%
  mutate(
    userId=as.factor(userId),
    iteration=as.factor(iteration)) %>%
  select(userId, dalleVersion, replayDalleVersion, logo_task, iteration, score, z_score,
         score_pct_rank_group)

# data sets for each effect type
tmp_direct <- tmp_data %>% filter(dalleVersion==base_model) %>%
  mutate(effect=ifelse(replayDalleVersion==base_model, 0, 1))
tmp_indirect <- tmp_data %>% filter(replayDalleVersion==other_model) %>%
  mutate(effect=ifelse(dalleVersion==base_model, 0, 1))
tmp_total <- tmp_data %>% filter(replayDalleVersion==dalleVersion) %>%
  mutate(effect=ifelse(dalleVersion==base_model, 0, 1))

# z-score results
model_fit <- lm(z_score ~ iteration + logo_task + effect +
                  as.factor(score_pct_rank_group) + effect:score_pct_rank_group, tmp_direct)
model_summary_robust <- as.data.frame(coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")[,])
cat(paste('Regression results of z_score for direct effect with cluster robust standard errors:'))
model_summary_robust %>% rownames_to_column(var = "coefname") %>% filter(grepl("effect", coefname))

# raw score results
model_fit_direct <- lm(score ~ iteration + logo_task + effect +
                         as.factor(score_pct_rank_group) + effect:score_pct_rank_group, tmp_direct)
model_fit_indirect <- lm(score ~ iteration + logo_task + effect +
                         as.factor(score_pct_rank_group) + effect:score_pct_rank_group, tmp_indirect)
model_fit_total <- lm(score ~ iteration + logo_task + effect +
                         as.factor(score_pct_rank_group) + effect:score_pct_rank_group, tmp_total)
model_summary_direct <- as.data.frame(coeftest(model_fit_direct, vcov=vcovCL, cluster=~userId, type = "HC1")[,])
model_summary_indirect <- as.data.frame(coeftest(model_fit_indirect, vcov=vcovCL, cluster=~userId, type = "HC1")[,])
model_summary_total <- as.data.frame(coeftest(model_fit_total, vcov=vcovCL, cluster=~userId, type = "HC1")[,])

cat(paste('Regression results of score for direct effect with cluster robust standard errors:'))
direct_coefs <- model_summary_direct %>% rownames_to_column(var = "coefname") %>%
  filter(grepl("effect", coefname)) %>% mutate(coefname=gsub("effect", "Direct", coefname),
                                               coefname=gsub("score_pct_rank_group", "Score Rank Group", coefname))
direct_coefs
cat(paste('Regression results of score for indirect effect with cluster robust standard errors:'))
indirect_coefs <- model_summary_indirect %>% rownames_to_column(var = "coefname") %>%
  filter(grepl("effect", coefname)) %>% mutate(coefname=gsub("effect", "Indirect", coefname),
                                               coefname=gsub("score_pct_rank_group", "Score Rank Group", coefname))
indirect_coefs
cat(paste('Regression results of score for total effect with cluster robust standard errors:'))
total_coefs <- model_summary_total %>% rownames_to_column(var = "coefname") %>%
  filter(grepl("effect", coefname)) %>% mutate(coefname=gsub("effect", "Total", coefname),
                                               coefname=gsub("score_pct_rank_group", "Score Rank Group", coefname))
total_coefs

# save the effect coefficients along with their interaction with performance group in a single compact table 
format_estimate <- function(estimate, std_error) {
  sprintf("%.4f (%.4f)", estimate, std_error)
}
extract_results <- function(model, label) {
  model %>%  mutate(Estimate = mapply(format_estimate, Estimate, `Std. Error`),
                    `P-value` = round(`Pr(>|t|)`, 4)) %>%
    select(coefname, Estimate, `P-value`)
}
direct_results <- extract_results(direct_coefs, "Direct")
indirect_results <- extract_results(indirect_coefs, "Indirect")
total_results <- extract_results(total_coefs, "Total")
all_models <- bind_rows(direct_results, indirect_results, total_results)
all_models %>%
  #kable("html", align = 'c', col.names = c("", "Estimate (SE)", "P-value")) %>%
  kable("latex", booktabs = TRUE, align = 'c', col.names = c("", "Estimate (SE)", "P-value")) %>%
  kable_styling(latex_options = c("striped", "hold_position"), full_width = FALSE) %>%
  #add_header_above(c(" " = 1, "Model Results" = 2)) %>%
  row_spec(nrow(direct_results), hline_after = TRUE, extra_css = "border-bottom: double;") %>%
  row_spec(nrow(direct_results) + nrow(indirect_results), hline_after = TRUE, extra_css = "border-bottom: double;") %>%
  save_kable(file = here("logo_exp_replication", "outputs", "tables", "qte_table_task2.tex"))

################################################################################
##### Performance and ATE variation by iteration
strat_vars <- c('logo_task')
weights <- replay_df %>%
  filter(dalleVersion == replayDalleVersion) %>%
  group_by(across(all_of(strat_vars))) %>%
  summarise(cnt=n()) %>%
  ungroup() %>%
  mutate(weight=cnt/sum(cnt)) %>%
  select(-cnt)
summary_df <- replay_df %>%
  filter(dalleVersion == replayDalleVersion) %>%
  group_by(across(c(all_of(strat_vars), dalleVersion, iteration))) %>%
  summarise(across(c(z_score, score),
                   ~ mean(.x), .names="mean_{.col}"),
            across(c(z_score, score),
                   ~ sd(.x)/sqrt(n()), .names="se_{.col}")
            ) %>%
  ungroup()
if(length(strat_vars) > 0) {
  summary_df <- summary_df %>%
  inner_join(weights, by=strat_vars) %>%
  group_by(dalleVersion, iteration) %>%
  summarise(across(c(mean_z_score, mean_score), ~ weighted.mean(.x, weight)),
            across(c(se_z_score, se_score), ~ sqrt(sum((.x**2)*(weight**2)))),
            sum_weight=sum(weight)) %>% 
  ungroup()
  assertthat::assert_that(all(summary_df$sum_weight > 0.999 & summary_df$sum_weight < 1.001))
  summary_df <- summary_df %>% select(-sum_weight)
}
summary_df_wide <- summary_df %>%
  pivot_wider(names_from=dalleVersion,
              values_from = c(mean_z_score, se_z_score, mean_score, se_score),
              names_glue='{.value}_dalle{dalleVersion}') %>%
  mutate(ate_z_score = mean_z_score_dalle3 - mean_z_score_dalle2,
         se_ate_score = sqrt(se_z_score_dalle3**2 + se_z_score_dalle2**2),
         ate_score = mean_score_dalle3 - mean_score_dalle2,
         se_ate_score = sqrt(se_score_dalle3**2 + se_score_dalle2**2))
point_size <- 2.5
error_bar_size <- 1
axis_text_size <- 17
axis_title_size <- 17 
legend_text_size <- 16
p <- ggplot(summary_df_wide, aes(as.factor(iteration), ate_score)) +
  geom_point(position=position_dodge(width=0.3), size=point_size, color='black') + 
  geom_errorbar(aes(ymin=ate_score - 1*se_ate_score,
                    ymax=ate_score + 1*se_ate_score),
                linewidth=error_bar_size, width=0.2, position = position_dodge(width = 0.3), color='black') +
  theme_clean() +
  labs(x='Iteration', y='ATE Log BT Score') +
  theme(axis.text=element_text(color='black', size=axis_text_size),
        axis.title=element_text(color='black', size=axis_title_size),
        legend.title=element_blank(),
        #legend.position = "inside",
        #legend.position.inside = c(0.8, 0.87),
        legend.text=element_text(size=legend_text_size),
        legend.key.spacing.y = unit(0.2, 'cm'),
        legend.background = element_blank(),
        plot.background = element_blank())
p

p <- ggplot(summary_df, aes(as.factor(iteration), mean_score, shape=dalleVersion)) +
  geom_point(position=position_dodge(width=0.3), size=point_size, color='black') + 
  geom_pointrange(aes(ymin=mean_score - 2*se_score,
                      ymax=mean_score + 2*se_score),
                  position=position_dodge(width=0.3)) +
  labs(x='Iteration', y='Log BT Score', shape='Dall-E Model') +
  theme_minimal() +
  #coord_cartesian(y=c(0.785, 0.835)) +
  theme(axis.text=element_text(colour="black", size=axis_text_size),
        axis.title=element_text(colour="black", size=axis_title_size),
        legend.title=element_text(size=legend_text_size),
        legend.text=element_text(size=legend_text_size),
        legend.justification = "top",
        legend.key.height = unit(0.4, 'cm'),
        legend.key.spacing.y = unit(0, 'cm'),
        legend.box.margin = margin(-8,0,0,-10),
        panel.grid.major = element_blank(),  # Remove major grid lines
        panel.grid.minor = element_blank(), 
        axis.line = element_line(color = "black"),  # Add axis lines
        legend.background = element_blank(),
        plot.background = element_blank())
p

## ATE variation and estimation by iteration using a linear model
base_model <- sort(unique(replay_df$dalleVersion))[1]
other_model <- sort(unique(replay_df$dalleVersion))[2]
#base_model <- sort(unique(replay_df$dalleVersion))[2]
#other_model <- sort(unique(replay_df$dalleVersion))[1]
cat('Using base model of', base_model, 'and other model of', other_model)

# which effect type?
#effect_type <- 'direct'
#effect_type <- 'indirect'
effect_type <- 'total'
tmp_data <- replay_df %>%
  filter((effect_type=='direct' & dalleVersion==base_model) |
         (effect_type=='indirect' & replayDalleVersion==other_model) |
         (effect_type=='total' & replayDalleVersion==dalleVersion)) %>%
  mutate(
    userId=as.factor(userId),
    effect=case_when(
      effect_type=='direct' ~ ifelse(replayDalleVersion==base_model, 0, 1),
      effect_type=='indirect' ~ ifelse(dalleVersion==base_model, 0, 1),
      effect_type=='total' ~ ifelse(dalleVersion==base_model, 0, 1),
      TRUE ~ -1)) %>%
  select(userId, dalleVersion, replayDalleVersion, iteration, logo_task, effect, score) %>%
  group_by(logo_task) %>%
  mutate(z_score = scale(score)[,1]) %>%
  ungroup()

# z-score
model_fit <- lm(z_score ~ effect*iteration + logo_task + effect, tmp_data)
model_summary <- summary(model_fit)
model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat(paste('Regression results for', effect_type, 'effect on z_score with cluster robust standard errors:'))
model_summary_robust

# raw score
model_fit <- lm(score ~ effect*iteration + logo_task + effect, tmp_data)
model_summary <- summary(model_fit)
model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat(paste('Regression results for', effect_type, 'effect on score with cluster robust standard errors:'))
model_summary_robust


################################################################################
##### Revised Dalle3 against a base model (either Dalle2 or Dalle3 revised)
strat_vars <- c('logo_task', 'iteration')
base_model <- '2'
replay_df <- replay_df_full %>% filter(junk_flag==FALSE)
replay_df <- replay_df %>%
  filter(dalleVersion %in% c(base_model, '3revised'),
         replayDalleVersion %in% c(base_model, '3revised'))
# calculate z-score for comparison between these two models
replay_df <- replay_df %>%
  group_by(across(all_of(strat_vars))) %>%
  mutate(z_score = scale(score)[,1]) %>%
  ungroup()

# Results using regression (comparing revised dalle3 vs dalle2)
cat('Using base model of', base_model)

tmp_data <- replay_df %>%
  filter(replayDalleVersion==dalleVersion) %>%
  mutate(
    userId=as.factor(userId),
    iteration=as.factor(iteration),
    effect=ifelse(dalleVersion==base_model, 0, 1)) %>%
  select(userId, dalleVersion, replayDalleVersion, logo_task, iteration, effect, score, z_score)

dep_var <- 'score'
#dep_var <- 'z_score'
model_fit <- lm(as.formula(paste(dep_var, ' ~ iteration + logo_task + effect')), tmp_data)
model_summary <- summary(model_fit)
model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat('Regression results for', dep_var, 'with cluster robust standard errors:')
model_summary_robust

cat("✅ Bootstrap 1/3 complete\n")
cat("🔄 Running bootstrap (2/3) — please be patient, this may take ~3 min...\n")
### Robustness check: ate using stratification. The results in the table are based on regression above
weights <- replay_df %>%
  group_by(across(all_of(strat_vars))) %>%
  summarise(cnt=n()) %>%
  ungroup() %>%
  mutate(weight=cnt/sum(cnt)) %>%
  select(-cnt)

summary_df <- replay_df %>%
  group_by(dalleVersion, replayDalleVersion) %>%
  group_modify(~ stratified_mean_boot_cluster_se(.x, strat_vars, 'userId', weights)) %>%
  ungroup()

base_model <- unique(summary_df$dalleVersion)[1]
cat('Using base model of', base_model)
ate_df <- summary_df %>% cross_join(summary_df, suffix = c("", "1")) %>%
  mutate(ate_z_score = mean_z_score - mean_z_score1,
         se_clustered_ate_z_score = sqrt(se_clustered_z_score**2 + se_clustered_z_score1**2),
         p_ate_z_score = 2*(1-pnorm(abs(ate_z_score), sd=se_clustered_ate_z_score)),
         ate_z_score_lower_ci = ate_z_score - qnorm(0.975)*se_clustered_ate_z_score,
         ate_z_score_upper_ci = ate_z_score + qnorm(0.975)*se_clustered_ate_z_score,
         ate_score = mean_score - mean_score1,
         se_clustered_ate_score = sqrt(se_clustered_score**2 + se_clustered_score1**2),
         p_ate_score = 2*(1-pnorm(abs(ate_score), sd=se_clustered_ate_score)),
         ate_score_lower_ci = ate_score - qnorm(0.975)*se_clustered_ate_score,
         ate_score_upper_ci = ate_score + qnorm(0.975)*se_clustered_ate_score,
         ) %>%
  select(-matches('^mean_z_score*|^se_z_score*|^mean_score*|^se_score*|se_clustered_z_score*|se_clustered_score*')) %>%
  filter(dalleVersion!=dalleVersion1 & dalleVersion1==base_model)

# results based on z-score
ate_df %>% select(c(ate_z_score, se_clustered_ate_z_score, ate_z_score_lower_ci, ate_z_score_upper_ci, p_ate_z_score))
# results based on raw-score
ate_df %>% select(c(ate_score, se_clustered_ate_score, ate_score_lower_ci, ate_score_upper_ci, p_ate_score))


################################################################################
##### Revised Dalle3 ATE vs Verbatim Dalle3 ATE
# how much does revision reduce benefit of Dalle-3? ratio estimation. SE through bootstrap
base_model <- '2'
replay_df <- replay_df_full %>% filter(junk_flag==FALSE)
tmp_data <- replay_df %>%
  filter(replayDalleVersion==dalleVersion) %>%
  mutate(
    userId=as.factor(userId),
    iteration=as.factor(iteration)) %>%
  select(userId, dalleVersion, replayDalleVersion, logo_task, iteration, score)
dalle3_df <- tmp_data %>% filter(replayDalleVersion %in% c('2', '3'))
dalle3r_df <- tmp_data %>% filter(replayDalleVersion %in% c('2', '3revised'))

dalle3_3revised_ratio <- function(d3_df, d3r_df) {
  dep_var <- 'score'
  # dalle 3 coef/se
  model_fit <- lm(as.formula(paste(dep_var, ' ~ iteration + logo_task + dalleVersion')),
                  d3_df)
  model_summary <- summary(model_fit)
  model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
  dalle3_coef <- model_summary_robust["dalleVersion3", "Estimate"]
  dalle3_se <- model_summary_robust["dalleVersion3", "Std. Error"]
  # dalle 3 revised coef/se
  model_fit <- lm(as.formula(paste(dep_var, ' ~ iteration + logo_task + dalleVersion')),
                  d3r_df)
  model_summary <- summary(model_fit)
  model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
  dalle3r_coef <- model_summary_robust["dalleVersion3revised", "Estimate"]
  dalle3r_se <- model_summary_robust["dalleVersion3revised", "Std. Error"]
  
  ratio_est <- 1 - (dalle3r_coef/dalle3_coef)
  return(ratio_est)
}
  
boot_ratio_fn = function(dalle3_df, dalle3r_df, B = 300) {
  1:B %>%
    future_map(~ {
      # For each iteration, generate a sample of each df with replacement
      sample_dalle3_df <- dalle3_df[sample(nrow(dalle3_df), replace = TRUE), ]
      sample_dalle3r_df <- dalle3r_df[sample(nrow(dalle3r_df), replace = TRUE), ]
      # Obtain the fn estimate for each pair of bootstrap sample
      dalle3_3revised_ratio(sample_dalle3_df, sample_dalle3r_df)
    }) %>%
    unlist() %>%
    sd()
}
ratio_est <- dalle3_3revised_ratio(dalle3_df, dalle3r_df)
ratio_se_boot <- boot_ratio_fn(dalle3_df, dalle3r_df, 600)
cat('Reduction in ATE of DALLE3 vs DALLE3-revised: ', ratio_est)
cat('95% confidence interval with cluster robust standard errors and bootstrap of the ratio:')
cat("(", ratio_est - 1.96*ratio_se_boot, " , ", ratio_est + 1.96*ratio_se_boot, ")")


################################################################################
# NOTE: DreamScore section removed — DreamScore is a main experiment metric
# and does not exist in the logo experiment dataset.
################################################################################



##### Heterogeneity of the direct and indirect effects by performance
# the performance tile of each user within its own task and iteration
strat_vars <- c('logo_task', 'iteration')
pct_groups <- c(0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100)
#pct_groups <- c(0, 20, 40, 60, 80, 100)
find_closest <- function(x) {
  pct_groups[which.min(abs(pct_groups - x))]
}
tmp_data <- replay_df %>%
  group_by(across(c(all_of(strat_vars), dalleVersion, replayDalleVersion))) %>%
  mutate(score_pct_rank=percent_rank(score)*100,
         score_pct_rank_group = sapply(score_pct_rank, find_closest)) %>%
  ungroup()

cat("✅ Bootstrap 2/3 complete\n")
cat("🔄 Running bootstrap (3/3) — please be patient, this may take ~3 min...\n")
# using stratification
cat(paste('Stratification vars:', paste(strat_vars, collapse=','), '\n'))
base_model <- sort(unique(replay_df$dalleVersion))[1]
other_model <- sort(unique(replay_df$dalleVersion))[2]
cat('Using base model of', base_model, 'and other model of', other_model, '\n')

summary_df <- tmp_data %>%
  group_by(dalleVersion, replayDalleVersion, score_pct_rank_group) %>%
  group_modify(~ stratified_mean_boot_cluster_se(.x, strat_vars, 'userId', weights)) %>%
  ungroup()
summary_df

cat('Using base model of', base_model, '\n')
ate_df <- summary_df %>% inner_join(summary_df, by="score_pct_rank_group",
                                    relationship="many-to-many", suffix = c("", "1")) %>%
  mutate(ate_z_score = mean_z_score - mean_z_score1,
         se_ate_z_score = sqrt(se_z_score**2 + se_z_score1**2),
         se_clustered_ate_z_score = sqrt(se_clustered_z_score**2 + se_clustered_z_score1**2),
         ate_score = mean_score - mean_score1,
         se_ate_score = sqrt(se_score**2 + se_score1**2),
         se_clustered_ate_score = sqrt(se_clustered_score**2 + se_clustered_score1**2)) %>%
  select(-matches('^mean_z_score*|^se_z_score*|^mean_score*|^se_score*')) %>%
  filter((dalleVersion1 == base_model    & replayDalleVersion1 == base_model &
          dalleVersion  == dalleVersion1 & replayDalleVersion  != replayDalleVersion1) |     # direct effect
         (dalleVersion1 == base_model    & replayDalleVersion != base_model &
          dalleVersion  != dalleVersion1 & replayDalleVersion == replayDalleVersion1) |      # indirect effect
         (dalleVersion1 == base_model    & replayDalleVersion1 == base_model &
          dalleVersion  != dalleVersion1 & replayDalleVersion  != replayDalleVersion1)) %>%  # full effect
  mutate(label=case_when(
           (dalleVersion != dalleVersion1 & replayDalleVersion != replayDalleVersion1) ~ 'ATE (Model & Prompting)',
           dalleVersion > base_model ~ 'Prompting Effect', 
           dalleVersion < base_model ~ 'Prompting Effect',
           replayDalleVersion < base_model ~ 'Model Effect',
           replayDalleVersion > base_model ~ 'Model Effect',
           TRUE ~ "Something is wrong"),
         label=factor(label, levels=c('Model Effect',
                                      'ATE (Model & Prompting)',
                                      'Prompting Effect',
                                      'Something is wrong'),
                      ordered=TRUE))

# first tile has a very large se
ate_df1 <- ate_df

point_size <- 2.5
error_bar_size <- 1
axis_text_size <- 19
axis_title_size <- 19
legend_text_size <- 16
p2 <- ggplot(ate_df1, aes(score_pct_rank_group, ate_z_score, color=label)) +
  geom_point(position=position_dodge(width=2.4), size=point_size) + 
  geom_errorbar(aes(ymin=ate_z_score - 1.96*se_clustered_ate_z_score,
                    ymax=ate_z_score + 1.96*se_clustered_ate_z_score),
                linewidth=error_bar_size, width=0.2, position = position_dodge(width = 2.4)) +
  geom_hline(yintercept = 0, linetype='longdash', color='black') +
  scale_color_manual(values=c('indianred', 'black', 'dodgerblue')) +
  scale_x_continuous(breaks = seq(0, 100, by = 10)) +
  theme_clean() +
  #coord_cartesian(y=c(0.0,0.25)) +
  labs(x='Performance Percentile', y = expression(Delta ~ "Z-Log BT Score")) +
  theme(axis.text=element_text(colour="black", size=axis_text_size),
        axis.title=element_text(colour="black", size=axis_title_size),
        legend.title=element_blank(),
        legend.position = "inside",
        legend.position.inside = c(0.8, 0.87),
        legend.text=element_text(size=legend_text_size),
        legend.key.spacing.y = unit(0.2, 'cm'),
        legend.background = element_blank(),
        plot.background = element_blank())
p2

p <- ggplot(ate_df1, aes(score_pct_rank_group, ate_score, color=label)) +
  geom_point(position=position_dodge(width=3.6), size=point_size) + 
  geom_errorbar(aes(ymin=ate_score - 1.96*se_clustered_ate_score,
                    ymax=ate_score + 1.96*se_clustered_ate_score),
                linewidth=error_bar_size, width=0.2, position = position_dodge(width = 3.6)) +
  geom_hline(yintercept = 0, linetype='longdash', color='black') +
  scale_color_manual(values=c('indianred', 'black', 'dodgerblue')) +
  scale_x_continuous(breaks = seq(0, 100, by = 10)) +
  theme_clean() +
  #coord_cartesian(y=c(-0.0022, 0.027)) +
  labs(x='Performance Percentile', y = expression(Delta ~ "Log BT Score")) +
  theme(axis.text=element_text(colour="black", size=axis_text_size),
        axis.title=element_text(colour="black", size=axis_title_size),
        legend.title=element_blank(),
        legend.position = "inside",
        legend.position.inside = c(0.76, 0.89),
        legend.text=element_text(size=legend_text_size),
        legend.key.spacing.y = unit(0.2, 'cm'),
        legend.background = element_blank(),
        plot.background = element_blank())
p

#################
# using regression
regression_ate_cluster_se <- function(data, dep_var, control_vars, cluster_var) {
  model_fit <- lm(as.formula(paste(dep_var, '~', paste(control_vars, collapse=' + '), '+ effect')),
                  data)
  model_fit_cluster <- coeftest(model_fit, vcov = vcovCL, cluster = ~userId)
  effect_coef <- model_fit_cluster["effect", "Estimate"]
  effect_se <- model_fit_cluster["effect", "Std. Error"]
  effect_p <- model_fit_cluster["effect", "Pr(>|t|)"]
  return(data.frame(coef = effect_coef, se = effect_se, p_val=effect_p))
}
base_model <- sort(unique(replay_df$dalleVersion))[1]
other_model <- sort(unique(replay_df$dalleVersion))[2]
cat('Using base model of', base_model, 'and other model of', other_model)

direct_df <- tmp_data %>% filter(dalleVersion==base_model) %>%
  mutate(effect_type='direct', effect=ifelse(replayDalleVersion==other_model, 1, 0))
indirect_df <- tmp_data %>% filter(replayDalleVersion==other_model) %>%
  mutate(effect_type='indirect', effect=ifelse(dalleVersion==other_model, 1, 0))
total_df <- tmp_data %>% filter(dalleVersion==replayDalleVersion) %>%
  mutate(effect_type='total', effect=ifelse(dalleVersion==other_model, 1, 0))
reg_df <- bind_rows(direct_df, indirect_df, total_df) %>%
  select(c(userId, effect_type, effect, score, z_score, iteration, logo_task, score_pct_rank_group)) %>%
  mutate(iteration=as.factor(iteration),
         userId=as.factor(userId))

ate_df <- reg_df %>%
  group_by(effect_type, score_pct_rank_group) %>%
  group_modify(~ regression_ate_cluster_se(.x, 'score',
                                           c('logo_task', 'iteration'),
                                           'userId')) %>%
  ungroup() %>%
  mutate(label=case_when(
           (effect_type == 'total') ~ 'ATE (Model & Prompting)',
            effect_type == 'indirect' ~ 'Prompting Effect', 
            effect_type == 'direct' ~ 'Model Effect',
           TRUE ~ "Something is wrong"),
         label=factor(label, levels=c('Model Effect',
                                      'ATE (Model & Prompting)',
                                      'Prompting Effect',
                                      'Something is wrong'),
                      ordered=TRUE))

point_size <- 2.5
error_bar_size <- 1
axis_text_size <- 19
axis_title_size <- 19
legend_text_size <- 16
p <- ggplot(ate_df, aes(score_pct_rank_group, coef, color=label)) +
  geom_point(position=position_dodge(width=3.6), size=point_size) + 
  geom_errorbar(aes(ymin=coef - 1.96*se, ymax=coef + 1.96*se),
                linewidth=error_bar_size, width=0.2, position = position_dodge(width = 3.6)) +
  geom_hline(yintercept = 0, linetype='longdash', color='black') +
  scale_color_manual(values=c('indianred', 'black', 'dodgerblue')) +
  scale_x_continuous(breaks = seq(0, 100, by = 10)) +
  theme_clean() +
  #coord_cartesian(y=c(-0.001, 0.027)) +
  labs(x='Performance Decile', y = expression(Delta ~ "Transformed BT-score")) +
  theme(axis.text=element_text(colour="black", size=axis_text_size),
        axis.title=element_text(colour="black", size=axis_title_size),
        legend.title=element_blank(),
        legend.position = "inside",
        legend.position.inside = c(0.77, 0.89),
        legend.text=element_text(size=legend_text_size),
        legend.key.spacing.y = unit(0.2, 'cm'),
        legend.background = element_blank(),
        plot.background = element_blank())
p

ate_df <- reg_df %>%
  group_by(effect_type, score_pct_rank_group) %>%
  group_modify(~ regression_ate_cluster_se(.x, 'z_score',
                                           c('logo_task', 'iteration'),
                                           'userId')) %>%
  ungroup() %>%
  mutate(label=case_when(
           (effect_type == 'total') ~ 'ATE (Model & Prompting)',
            effect_type == 'indirect' ~ 'Prompting Effect', 
            effect_type == 'direct' ~ 'Model Effect',
           TRUE ~ "Something is wrong"),
         label=factor(label, levels=c('Model Effect', 'Prompting Effect', 'ATE (Model & Prompting)', 'Something is wrong'),
                      ordered=TRUE))

point_size <- 2.5
error_bar_size <- 1
axis_text_size <- 19
axis_title_size <- 19
legend_text_size <- 16
p <- ggplot(ate_df, aes(score_pct_rank_group, coef, color=label)) +
  geom_point(position=position_dodge(width=3.6), size=point_size) + 
  geom_errorbar(aes(ymin=coef - 1.96*se, ymax=coef + 1.96*se),
                linewidth=error_bar_size, width=0.2, position = position_dodge(width = 3.6)) +
  geom_hline(yintercept = 0, linetype='longdash', color='black') +
  scale_color_manual(values=c('indianred', 'dodgerblue', 'black')) +
  scale_x_continuous(breaks = seq(0, 100, by = 10)) +
  theme_clean() +
  #coord_cartesian(y=c(0.0,0.25)) +
  labs(x='Performance Decile', y = expression(Delta ~ "Z-Log BT Score")) +
  theme(axis.text=element_text(colour="black", size=axis_text_size),
        axis.title=element_text(colour="black", size=axis_title_size),
        legend.title=element_blank(),
        legend.position = "inside",
        legend.position.inside = c(0.76, 0.89),
        legend.text=element_text(size=legend_text_size),
        legend.key.spacing.y = unit(0.2, 'cm'),
        legend.background = element_blank(),
        plot.background = element_blank())
p

cat("✅ Bootstrap 3/3 complete — all bootstraps finished\n")
