pdf(NULL) # suppress Rplots.pdf in batch mode
# ============================================================================
# FIGURE 1: ITERATION PLOTS (PNAS)
# Stratified ATE and iteration-level performance for cosine sim and dreamsim
# ============================================================================

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

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

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

source(here("shared", "utils.R"))

# === DATA IMPORT ===
user_iter_file <- here("similarity_exp_replication", "inputs",
                       "main_df_promptLvl_2026-03-22_PROD_HASHED_TRIMMED.csv")

result <- generate_data_from_user_iter_file(user_iter_file)
user_iter_df <- result$user_iter_df
user_df <- result$user_df

# === DATA PROCESSING ===
main_df <- user_iter_df %>%
  rename(image_id = originalID)

working_df <- main_df %>%
  select(image_id, userId, referenceImageName, cosine_sim_mean, dream_sim_mean, cosine_sim_observed,
         dream_sim_observed, iteration, dalle_version, junk_flag, dalle3verbatimReplay_mean_cosine_score,
         dalle2Replay_mean_cosine_score, dalle3verbatimReplay_mean_dreamsim_score, dalle2Replay_mean_dreamsim_score,
         dalle3verbatim_replayCount, dalle2Replay_replayCount, replayCompleteness) %>%
  arrange(userId, iteration) %>%
  filter(dalle_version != "3-revised-prompt") %>%
  mutate(replay_cosine_sim_mean = case_when(
    dalle_version == "2" ~ dalle2Replay_mean_cosine_score,
    dalle_version == "3-verbatim-prompt" ~ dalle3verbatimReplay_mean_cosine_score,
    TRUE ~ NA_real_
  )) %>%
  mutate(replay_dream_sim_mean = case_when(
    dalle_version == "2" ~ dalle2Replay_mean_dreamsim_score,
    dalle_version == "3-verbatim-prompt" ~ dalle3verbatimReplay_mean_dreamsim_score,
    TRUE ~ NA_real_
  )) %>%
  filter(junk_flag == FALSE) %>%
  filter(!is.na(dalle2Replay_mean_cosine_score) & !is.na(dalle3verbatimReplay_mean_cosine_score))

# Z-Scores across task
working_df <- working_df %>%
  group_by(referenceImageName) %>%
  mutate(dream_sim_zscore_taskwide        = (dream_sim_mean - mean(dream_sim_mean, na.rm = TRUE)) / sd(dream_sim_mean, na.rm = TRUE),
         cosine_sim_zscore_taskwide       = (cosine_sim_mean - mean(cosine_sim_mean, na.rm = TRUE)) / sd(cosine_sim_mean, na.rm = TRUE),
         dream_sim_zscore_taskwide_noisy  = (dream_sim_observed - mean(dream_sim_observed, na.rm = TRUE)) / sd(dream_sim_observed, na.rm = TRUE),
         cosine_sim_zscore_taskwide_noisy = (cosine_sim_observed - mean(cosine_sim_observed, na.rm = TRUE)) / sd(cosine_sim_observed, na.rm = TRUE),
         cosine_sim_zscore_taskwide_replay = (replay_cosine_sim_mean - mean(replay_cosine_sim_mean, na.rm = TRUE)) / sd(replay_cosine_sim_mean, na.rm = TRUE),
         dream_sim_zscore_taskwide_replay  = (replay_dream_sim_mean - mean(replay_dream_sim_mean, na.rm = TRUE)) / sd(replay_dream_sim_mean, na.rm = TRUE)) %>%
  ungroup()

# Frequency table
freq_table <- main_df %>%
  select(image_id, userId, referenceImageName, cosine_sim_mean, dream_sim_mean, cosine_sim_observed,
         dream_sim_observed, iteration, dalle_version, junk_flag, dalle3verbatimReplay_mean_cosine_score,
         dalle2Replay_mean_cosine_score, dalle3verbatimReplay_mean_dreamsim_score, dalle2Replay_mean_dreamsim_score,
         dalle3verbatim_replayCount, dalle2Replay_replayCount, replayCompleteness) %>%
  count(dalle_version, iteration) %>%
  spread(key = iteration, value = n, fill = 0)

print(freq_table)
freq(main_df$junk_flag)

cat("Working df:", nrow(working_df), "observations,", n_distinct(working_df$userId), "users\n")

# === CORRELATION CHECK ===
cor.test(working_df$dream_sim_mean, working_df$cosine_sim_mean) %>%
  report()

correlation <- cor(working_df$dream_sim_mean, working_df$cosine_sim_mean)
print(correlation)

# ============================================================================
# OVERALL ATE (STRATIFIED BY TASK AND ITERATION, COSINE SIM)
# ============================================================================
strat_vars <- c('referenceImageName', '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_CosineScore = mean(replay_cosine_sim_mean, na.rm = TRUE),
    se_CosineScore = sd(replay_cosine_sim_mean, na.rm = TRUE) / sqrt(sum(!is.na(replay_cosine_sim_mean))),
    n_obs = sum(!is.na(replay_cosine_sim_mean))
  ) %>%
  ungroup()

stratified_summary_df <- pre_strat_df %>%
  inner_join(weights, by = strat_vars) %>%
  group_by(dalle_version) %>%
  summarise(across(mean_CosineScore, ~ weighted.mean(.x, weight)),
            across(se_CosineScore, ~ 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_CosineScore, se_CosineScore),
    names_glue = "{.value}{pair}"
  ) %>%
  mutate(dim_CosineScore = mean_CosineScore2 - mean_CosineScore1,
         se_dim_CosineScore = sqrt(se_CosineScore1**2 + se_CosineScore2**2),
         dim_CosineScore_lower_ci = dim_CosineScore - qnorm(0.975) * se_dim_CosineScore,
         dim_CosineScore_upper_ci = dim_CosineScore + qnorm(0.975) * se_dim_CosineScore)

dim_ate_stratified <- difference_in_means_df$dim_CosineScore
dim_CosineScore_lower_ci_constant <- difference_in_means_df$dim_CosineScore_lower_ci
dim_CosineScore_upper_ci_constant <- difference_in_means_df$dim_CosineScore_upper_ci

difference_in_means_df %>% select(dim_CosineScore, dim_CosineScore_lower_ci, dim_CosineScore_upper_ci)

# ============================================================================
# ITERATION-LEVEL PLOT (COSINE SIM, 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('referenceImageName')

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_CosineScore = mean(replay_cosine_sim_mean, na.rm = TRUE),
    se_CosineScore = sd(replay_cosine_sim_mean, na.rm = TRUE) / sqrt(sum(!is.na(replay_cosine_sim_mean))),
    n_obs = sum(!is.na(replay_cosine_sim_mean))
  ) %>%
  ungroup()

stratified_summary_df <- pre_strat_df %>%
  inner_join(weights, by = strat_vars) %>%
  group_by(iteration, dalle_version) %>%
  summarise(across(mean_CosineScore, ~ weighted.mean(.x, weight)),
            across(se_CosineScore, ~ 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 %>%
  group_by(iteration) %>%
  mutate(pair = row_number()) %>%
  ungroup() %>%
  pivot_wider(
    names_from = pair,
    values_from = c(dalle_version, mean_CosineScore, se_CosineScore),
    names_glue = "{.value}{pair}"
  ) %>%
  mutate(dim_CosineScore = mean_CosineScore2 - mean_CosineScore1,
         se_dim_CosineScore = sqrt(se_CosineScore1**2 + se_CosineScore2**2),
         dim_CosineScore_lower_ci = dim_CosineScore - qnorm(0.975) * se_dim_CosineScore,
         dim_CosineScore_upper_ci = dim_CosineScore + qnorm(0.975) * se_dim_CosineScore) %>%
  mutate(estimate = "stratified")

difference_in_means_df <- difference_in_means_df %>%
  mutate(dim_CosineScore_lower_ci_constant = dim_CosineScore_lower_ci_constant,
         dim_CosineScore_upper_ci_constant = dim_CosineScore_upper_ci_constant)

ate_plot <- ggplot(data = difference_in_means_df,
       aes(x = as.factor(iteration), y = dim_CosineScore)) +
  geom_rect(aes(xmin = -Inf, xmax = Inf,
                ymin = dim_CosineScore_lower_ci_constant[1], ymax = dim_CosineScore_upper_ci_constant[1]),
            fill = "lightblue", alpha = 0.05, inherit.aes = FALSE) +
  geom_point(size = point_size) +
  geom_errorbar(aes(ymin = dim_CosineScore - qnorm(0.975) * se_dim_CosineScore,
                    ymax = dim_CosineScore + qnorm(0.975) * se_dim_CosineScore), width = error_bar_size) +
  geom_hline(yintercept = dim_ate_stratified, linetype = "solid", color = "blue") +
  geom_hline(yintercept = 0, linetype = "dashed", color = "black") +
  annotate("text", x = 1, y = dim_ate_stratified, label = "ATE", vjust = -0.5, hjust = 1.1, color = "black") +
  labs(x = "Attempt Number",
       y = expression(paste(Delta, " Cosine Similarity"))) +
  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")
  )

annotations_df <- data.frame(
  iteration = c(0, 5),
  mean_CosineScore = c(0.835, 0.795),
  dalle_version = c("DALLE-3 (Superior)", "DALLE-2 (Inferior)")
)

main_plot <- ggplot(data = stratified_summary_df,
                    aes(x = as.factor(iteration), y = mean_CosineScore, color = dalle_version)) +
  geom_point(aes(shape = dalle_version), size = point_size, position = position_dodge(width = dodge_amt)) +
  geom_errorbar(aes(ymin = mean_CosineScore - qnorm(0.975) * se_CosineScore,
                    ymax = mean_CosineScore + qnorm(0.975) * se_CosineScore), 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_CosineScore, label = dalle_version),
            hjust = -0.2, size = 6) +
  labs(x = "Attempt Number", y = "Cosine Similarity",
       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()
  )

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("similarity_exp_replication", "outputs", "figures"), recursive = TRUE, showWarnings = FALSE)
ggsave(here("similarity_exp_replication", "outputs", "figures", "fig1_iterPlot_stratified.pdf"),
       plot = fig, device = "pdf", width = 8, height = 5)

# ============================================================================
# REGRESSION VIEW (COSINE SIM)
# ============================================================================
working_df <- working_df %>%
  mutate(referenceImageName = factor(referenceImageName),
         userId = factor(userId))

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

fe_ols <- feols(replay_cosine_sim_mean ~ iteration * dalle_version + dalle_version | referenceImageName, pdat)
fe_ols_zscored <- feols(cosine_sim_zscore_taskwide_replay ~ iteration * dalle_version + dalle_version, pdat)

est_clu <- summary(fe_ols, vcov = ~userId)
est_clu_zscored <- summary(fe_ols_zscored, vcov = ~userId)
modelsummary(list("Cosine Sim" = est_clu, "Z-Score CS" = est_clu_zscored),
             statistic = "{std.error} ({p.value})", fmt = 4, output = "kableExtra", stars = TRUE,
             notes = list('Clustered SEs (p-values)'))

# ============================================================================
# OVERALL ATE (STRATIFIED, DREAMSIM)
# ============================================================================
strat_vars <- c('referenceImageName', '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_dreamScore = mean(replay_dream_sim_mean, na.rm = TRUE),
    se_dreamScore = sd(replay_dream_sim_mean, na.rm = TRUE) / sqrt(sum(!is.na(replay_dream_sim_mean))),
    n_obs = sum(!is.na(replay_dream_sim_mean))
  ) %>%
  ungroup()

stratified_summary_df <- pre_strat_df %>%
  inner_join(weights, by = strat_vars) %>%
  group_by(dalle_version) %>%
  summarise(across(mean_dreamScore, ~ weighted.mean(.x, weight)),
            across(se_dreamScore, ~ 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_dreamScore, se_dreamScore),
    names_glue = "{.value}{pair}"
  ) %>%
  mutate(dim_dreamScore = mean_dreamScore2 - mean_dreamScore1,
         se_dim_dreamScore = sqrt(se_dreamScore1**2 + se_dreamScore2**2),
         dim_dreamScore_lower_ci = dim_dreamScore - qnorm(0.975) * se_dim_dreamScore,
         dim_dreamScore_upper_ci = dim_dreamScore + qnorm(0.975) * se_dim_dreamScore,
         z_score_dim_dreamScore = dim_dreamScore / se_dim_dreamScore,
         p_value_dim_dreamScore = 2 * (1 - pnorm(abs(z_score_dim_dreamScore))))

dim_ate_stratified <- difference_in_means_df$dim_dreamScore
dim_dreamScore_lower_ci_constant <- difference_in_means_df$dim_dreamScore_lower_ci
dim_dreamScore_upper_ci_constant <- difference_in_means_df$dim_dreamScore_upper_ci

difference_in_means_df %>% select(dim_dreamScore, dim_dreamScore_lower_ci, dim_dreamScore_upper_ci,
                                  z_score_dim_dreamScore, p_value_dim_dreamScore)

# ============================================================================
# OVERALL ATE (Z-SCORED DREAMSIM, UNSTRATIFIED)
# ============================================================================
unstrat_test_df <- working_df %>%
  filter(junk_flag == FALSE) %>%
  group_by(dalle_version) %>%
  summarise(
    mean_Dream_ZScore = mean(dream_sim_zscore_taskwide_replay, na.rm = TRUE),
    se_Dream_ZScore = sd(dream_sim_zscore_taskwide_replay, na.rm = TRUE) / sqrt(sum(!is.na(dream_sim_zscore_taskwide_replay))),
    n_obs = sum(!is.na(dream_sim_zscore_taskwide_replay))
  )

unstrat_dim_df <- unstrat_test_df %>%
  mutate(pair = row_number()) %>%
  select(-n_obs) %>%
  ungroup() %>%
  pivot_wider(
    names_from = pair,
    values_from = c(dalle_version, mean_Dream_ZScore, se_Dream_ZScore),
    names_glue = "{.value}{pair}"
  ) %>%
  mutate(dim_Dream_ZScore = mean_Dream_ZScore2 - mean_Dream_ZScore1,
         se_dim_Dream_ZScore = sqrt(se_Dream_ZScore1**2 + se_Dream_ZScore2**2),
         p_dim_Dream_ZScore = 2 * (1 - pnorm(dim_Dream_ZScore, sd = se_dim_Dream_ZScore)),
         dim_Dream_ZScore_lower_ci = dim_Dream_ZScore - qnorm(0.975) * se_dim_Dream_ZScore,
         dim_Dream_ZScore_upper_ci = dim_Dream_ZScore + qnorm(0.975) * se_dim_Dream_ZScore)

unstrat_dim_df %>% select(dalle_version1, dalle_version2, dim_Dream_ZScore,
                          p_dim_Dream_ZScore, dim_Dream_ZScore_lower_ci, dim_Dream_ZScore_upper_ci)

# ============================================================================
# REGRESSION VIEW (DREAMSIM)
# ============================================================================
fe_ols <- feols(replay_dream_sim_mean ~ iteration * dalle_version + dalle_version | referenceImageName, pdat)
fe_ols_zscored <- feols(dream_sim_zscore_taskwide_replay ~ iteration * dalle_version + dalle_version, pdat)

est_clu <- summary(fe_ols, vcov = ~userId)
est_clu_zscored <- summary(fe_ols_zscored, vcov = ~userId)
modelsummary(list("Dream Sim" = est_clu, "Z-Score DS" = est_clu_zscored),
             statistic = "{std.error} ({p.value})", fmt = 4, output = "kableExtra", stars = TRUE,
             notes = list('Clustered SEs (p-values)'))
