pdf(NULL) # suppress Rplots.pdf in batch mode
# ============================================================================
# FIGURE 1: LOGO EXPERIMENT ITERATION PLOTS
# Stratified ATE and iteration-level BT beta performance
# ============================================================================

rm(list = ls())
set.seed(1324)

cbPalette <- c("#E69F00", "#56B4E9", "#009E73", "#F0E442",
               "#999999", "#0072B2", "#D55E00", "#CC79A7")

# === LIBRARIES ===
library(tidyverse)
library(fixest)
library(modelsummary)
library(gt)
library(kableExtra)
library(questionr)
library(cowplot)
library(here)
library(assertthat)
library(tidylog, warn.conflicts = FALSE)

# === DATA IMPORT ===
EPSILON <- 10^-10
repl_map <- c("2"                 = "2",
              "3-verbatim-prompt" = "3",
              "3-revised-prompt"  = "3revised")

strat_vars <- c('logoTaskSelection', 'iteration')
score_var <- 'strength'

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))

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  = "_") %>%
  rename(exp_bt_beta = strength) %>%
  mutate(replayDalleVersion = str_replace_all(replayDalleVersion, repl_map),
         exp_bt_beta = pmax(exp_bt_beta, EPSILON),
         strength = log(exp_bt_beta),
         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\n')
}
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\n')
}

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()

main_df <- replay_df_full %>%
  filter(iteration < 10) %>%
  filter(!is.na(strength)) %>%
  filter(dalleVersion == replayDalleVersion) %>%
  select(userId, image_id, dalleVersion, replayDalleVersion, logoTaskSelection, iteration, strength, score, z_score, logo_task, junk_flag, everything()) %>%
  arrange(userId, dalleVersion, replayDalleVersion, logoTaskSelection, iteration)

# === DATA PROCESSING ===
working_df <- main_df %>%
  select(userId, image_id, dalleVersion, replayDalleVersion, logoTaskSelection, iteration, strength, score, z_score, logo_task, junk_flag) %>%
  rename(dalle_version = dalleVersion,
         replay_dalle_version = replayDalleVersion) %>%
  arrange(userId, dalle_version, replay_dalle_version, logoTaskSelection, iteration) %>%
  filter(dalle_version != "3revised")

freq_table <- main_df %>%
  select(userId, image_id, dalleVersion, replayDalleVersion, logoTaskSelection, iteration, strength, logo_task, junk_flag) %>%
  rename(dalle_version = dalleVersion,
         replay_dalle_version = replayDalleVersion) %>%
  count(dalle_version, iteration) %>%
  spread(key = iteration, value = n, fill = 0)

print(freq_table)
freq(working_df$junk_flag)

# ============================================================================
# OVERALL ATE (STRATIFIED BY TASK AND ITERATION)
# ============================================================================
strat_vars <- c('logoTaskSelection', 'iteration')

weights <- working_df %>%
  filter(junk_flag == FALSE) %>%
  group_by(across(all_of(strat_vars))) %>%
  summarise(cnt = n()) %>%
  ungroup() %>%
  mutate(weight = cnt / sum(cnt)) %>%
  select(-cnt)

pre_strat_df <- working_df %>%
  filter(junk_flag == FALSE) %>%
  group_by(across(c(all_of(strat_vars), dalle_version))) %>%
  summarise(
    mean_Strength = mean(strength, na.rm = TRUE),
    se_Strength = sd(strength, na.rm = TRUE) / sqrt(sum(!is.na(strength))),
    n_obs = sum(!is.na(strength))
  ) %>%
  ungroup()

stratified_summary_df <- pre_strat_df %>%
  inner_join(weights, by = strat_vars) %>%
  group_by(dalle_version) %>%
  summarise(across(mean_Strength, ~ weighted.mean(.x, weight)),
            across(se_Strength, ~ sqrt(sum((.x**2) * (weight**2)))),
            sum_weight = sum(weight)) %>%
  ungroup()

assert_that(all(stratified_summary_df$sum_weight > 0.999 & stratified_summary_df$sum_weight < 1.001))
stratified_summary_df <- stratified_summary_df %>% select(-sum_weight)

difference_in_means_df <- stratified_summary_df %>%
  mutate(pair = row_number()) %>%
  ungroup() %>%
  pivot_wider(
    names_from = pair,
    values_from = c(dalle_version, mean_Strength, se_Strength),
    names_glue = "{.value}{pair}"
  ) %>%
  mutate(dim_Strength = mean_Strength2 - mean_Strength1,
         se_dim_Strength = sqrt(se_Strength1**2 + se_Strength2**2),
         dim_Strength_lower_ci = dim_Strength - qnorm(0.975) * se_dim_Strength,
         dim_Strength_upper_ci = dim_Strength + qnorm(0.975) * se_dim_Strength)

dim_Strength_ATE_stratified <- difference_in_means_df$dim_Strength
dim_Strength_ATE_SE <- difference_in_means_df$se_dim_Strength
dim_Strength_ATE_lower_ci_constant <- difference_in_means_df$dim_Strength_lower_ci
dim_Strength_ATE_upper_ci_constant <- difference_in_means_df$dim_Strength_upper_ci

difference_in_means_df %>% select(dim_Strength, dim_Strength_lower_ci, dim_Strength_upper_ci)

# ============================================================================
# ITERATION-LEVEL PLOT (STRATIFIED BY TASK)
# ============================================================================
error_bar_size <- 0
point_size <- 4
axis_text_size <- 16
axis_title_size <- 18
legend_text_size <- 18
dodge_amt <- 0.2

strat_vars <- c('logoTaskSelection')

stopifnot(length(strat_vars) > 0)

weights <- working_df %>%
  filter(junk_flag == FALSE) %>%
  group_by(across(all_of(strat_vars))) %>%
  summarise(cnt = n()) %>%
  ungroup() %>%
  mutate(weight = cnt / sum(cnt)) %>%
  select(-cnt)

pre_strat_df <- working_df %>%
  filter(junk_flag == FALSE) %>%
  group_by(across(c(all_of(strat_vars), iteration, dalle_version))) %>%
  summarise(
    mean_Strength = mean(strength, na.rm = TRUE),
    se_Strength = sd(strength, na.rm = TRUE) / sqrt(sum(!is.na(strength))),
    n_obs = sum(!is.na(strength))
  ) %>%
  ungroup()

stratified_summary_df <- pre_strat_df %>%
  inner_join(weights, by = strat_vars) %>%
  group_by(iteration, dalle_version) %>%
  summarise(across(mean_Strength, ~ weighted.mean(.x, weight)),
            across(se_Strength, ~ sqrt(sum((.x**2) * (weight**2)))),
            sum_weight = sum(weight)) %>%
  ungroup()

stopifnot(all(stratified_summary_df$sum_weight > 0.999 & stratified_summary_df$sum_weight < 1.001))
stratified_summary_df <- stratified_summary_df %>% select(-sum_weight)

difference_in_means_df <- stratified_summary_df %>%
  group_by(iteration) %>%
  mutate(pair = row_number()) %>%
  ungroup() %>%
  pivot_wider(
    names_from = pair,
    values_from = c(dalle_version, mean_Strength, se_Strength),
    names_glue = "{.value}{pair}"
  ) %>%
  mutate(dim_Strength = mean_Strength2 - mean_Strength1,
         se_dim_Strength = sqrt(se_Strength1**2 + se_Strength2**2),
         dim_Strength_lower_ci = dim_Strength - qnorm(0.975) * se_dim_Strength,
         dim_Strength_upper_ci = dim_Strength + qnorm(0.975) * se_dim_Strength) %>%
  mutate(estimate = "stratified")

difference_in_means_df <- difference_in_means_df %>%
  mutate(dim_Strength_ATE_lower_ci_constant = dim_Strength_ATE_lower_ci_constant,
         dim_Strength_ATE_upper_ci_constant = dim_Strength_ATE_upper_ci_constant)

ate_plot <- ggplot(data = difference_in_means_df,
       aes(x = as.factor(iteration), y = dim_Strength)) +
  geom_rect(aes(xmin = -Inf, xmax = Inf,
                ymin = dim_Strength_ATE_lower_ci_constant[1], ymax = dim_Strength_ATE_upper_ci_constant[1]),
            fill = "lightblue", alpha = 0.05, inherit.aes = FALSE) +
  geom_point(size = point_size) +
  geom_errorbar(aes(ymin = dim_Strength - qnorm(0.975) * se_dim_Strength,
                    ymax = dim_Strength + qnorm(0.975) * se_dim_Strength), width = error_bar_size) +
  geom_hline(yintercept = dim_Strength_ATE_stratified, linetype = "solid", color = "blue") +
  geom_hline(yintercept = 0, linetype = "dashed", color = "black") +
  annotate("text", x = 2.3, y = dim_Strength_ATE_stratified, label = "ATE", vjust = -0.5, hjust = 1.1, color = "black") +
  labs(x = "Attempt Number",
       y = expression(paste(Delta, " BT Betas"))) +
  scale_x_discrete(labels = 1:10) +
  theme(
    axis.text = element_text(colour = "black", size = axis_text_size),
    axis.title = element_text(colour = "black", size = axis_title_size),
    legend.position = "inside",
    legend.position.inside = c(0.18, 0.9),
    legend.title = element_blank(),
    legend.text = element_text(size = legend_text_size),
    panel.grid.major = element_blank(),
    panel.grid.minor = element_blank(),
    panel.background = element_rect(fill = "white", colour = NA),
    plot.background = element_rect(fill = "white", colour = NA),
    plot.subtitle = element_text(size = 16),
    panel.border = element_blank(),
    axis.line = element_line(colour = "black")
  )

dalle3_mean <- stratified_summary_df %>% filter(dalle_version == "3") %>% pull(mean_Strength) %>% mean()
dalle2_mean <- stratified_summary_df %>% filter(dalle_version == "2") %>% pull(mean_Strength) %>% mean()
iter_positions <- unique(sort(stratified_summary_df$iteration))
iter_pos_1 <- iter_positions[1]
iter_pos_mid <- iter_positions[ceiling(length(iter_positions) / 2)]

annotations_df <- data.frame(
  iteration = c(iter_pos_1, iter_pos_mid),
  mean_Strength = c(dalle3_mean, dalle2_mean),
  dalle_version = c("DALLE-3 (Superior)", "DALLE-2 (Inferior)")
)

main_plot <- ggplot(data = stratified_summary_df,
                    aes(x = as.factor(iteration), y = mean_Strength, color = dalle_version)) +
  geom_point(aes(shape = dalle_version), size = point_size, position = position_dodge(width = dodge_amt)) +
  geom_errorbar(aes(ymin = mean_Strength - qnorm(0.975) * se_Strength,
                    ymax = mean_Strength + qnorm(0.975) * se_Strength), width = error_bar_size,
                position = position_dodge(width = dodge_amt)) +
  geom_line(aes(group = dalle_version), alpha = 0.5, linetype = "dashed") +
  geom_text(data = annotations_df, aes(x = as.factor(iteration), y = mean_Strength + 0.2, label = dalle_version),
            hjust = -0.2, size = 6) +
  labs(x = "Attempt Number", y = "BT Betas",
       color = "Estimate Type", shape = "Estimate Type") +
  scale_color_manual(values = c("#f1a340", "#998ec3", "#f1a340", "#998ec3")) +
  scale_shape_manual(values = c(16, 16)) +
  scale_x_discrete(labels = 1:10) +
  theme(
    axis.text = element_text(colour = "black", size = axis_text_size),
    axis.title = element_text(colour = "black", size = axis_title_size),
    panel.grid.major = element_blank(),
    panel.grid.minor = element_blank(),
    panel.background = element_rect(fill = "white", colour = NA),
    plot.background = element_rect(fill = "white", colour = NA),
    legend.position = "none",
    legend.title = element_blank(),
    legend.text = element_text(size = legend_text_size),
    plot.subtitle = element_text(size = 16),
    panel.border = element_blank(),
    axis.line = element_line(colour = "black"),
    axis.title.x = element_blank(),
    axis.text.x = element_blank(),
    axis.ticks.x = element_blank()
  )

title <- ggdraw() +
  draw_label("Logo Experiment Main Data", x = 0, hjust = 0)

plot_row <- plot_grid(main_plot, ate_plot, ncol = 1, align = "v", axis = "lr")
fig <- plot_grid(plot_row, ncol = 1, rel_heights = c(0.1, 1))
print(fig)

dir.create(here("logo_exp_replication", "outputs", "figures"), recursive = TRUE, showWarnings = FALSE)
ggsave(here("logo_exp_replication", "outputs", "figures", "fig1_iterPlot_logoExp_stratified.pdf"),
       plot = fig, device = "pdf", width = 8, height = 5)

# ============================================================================
# REGRESSION VIEW (STRENGTH, BT BETAS)
# ============================================================================
working_df <- working_df %>%
  mutate(logoTaskSelection = factor(logoTaskSelection),
         userId = factor(userId))

pdat <- panel(working_df, panel.id = c('userId', 'iteration'))

fe_ols <- feols(strength ~ iteration * dalle_version + dalle_version | logoTaskSelection, pdat)
fe_ols_zscore <- feols(z_score ~ iteration * dalle_version + dalle_version | logoTaskSelection, pdat)

est_clu <- summary(fe_ols, vcov = ~userId)
est_clu_zscore <- summary(fe_ols_zscore, vcov = ~userId)
modelsummary(list("BT Betas" = est_clu, "Z-Score BT Betas" = est_clu_zscore),
             statistic = "({std.error})", fmt = 4, output = "kableExtra", stars = TRUE,
             notes = list('Clustered SEs'))

# ============================================================================
# STRATIFIED SUMMARY TABLE (BY ITERATION)
# ============================================================================
iteration_summary_table <- difference_in_means_df %>%
  mutate(
    `DALLE-2 BT Betas` = paste0(format(round(mean_Strength1, 4), nsmall = 4),
                                " (", format(round(se_Strength1, 4), nsmall = 4), ")"),
    `DALLE-3 BT Betas` = paste0(format(round(mean_Strength2, 4), nsmall = 4),
                                " (", format(round(se_Strength2, 4), nsmall = 4), ")"),
    `Difference-in-Means BT Betas` = paste0(format(round(dim_Strength, 4), nsmall = 4),
                                            " (", format(round(se_dim_Strength, 4), nsmall = 4), ")")
  ) %>%
  select(iteration, `DALLE-2 BT Betas`, `DALLE-3 BT Betas`, `Difference-in-Means BT Betas`) %>%
  rename(Iteration = iteration) %>%
  gt() %>%
  tab_header(title = "Stratified Mean BT Betas by Iteration",
             subtitle = "Estimates with Standard Errors in Parentheses") %>%
  cols_align(align = "center") %>%
  tab_options(table.font.size = "small",
              heading.title.font.size = "medium",
              heading.subtitle.font.size = "small")

iteration_summary_table

# ============================================================================
# OVERALL ATE SUMMARY
# ============================================================================
overall_ATE_summary_table <- data.frame(
  Label = "ATE (DALLE-3 vs DALLE-2)",
  Estimate_SE = paste0(format(round(dim_Strength_ATE_stratified, 4), nsmall = 4),
                       " (", format(round(dim_Strength_ATE_SE, 4), nsmall = 4), ")"),
  CI_95 = paste0("[", format(round(dim_Strength_ATE_lower_ci_constant, 4), nsmall = 4),
                 ", ", format(round(dim_Strength_ATE_upper_ci_constant, 4), nsmall = 4), "]")
) %>%
  gt() %>%
  cols_label(Estimate_SE = "Estimate (SE)", CI_95 = "95% CI") %>%
  tab_header(title = "Overall Average Treatment Effect (Stratified)",
             subtitle = "Estimate with Clustered Standard Errors and 95% Confidence Interval") %>%
  cols_align(align = "center") %>%
  tab_options(table.font.size = "small",
              heading.title.font.size = "medium",
              heading.subtitle.font.size = "small")

overall_ATE_summary_table
