pdf(NULL) # suppress Rplots.pdf in batch mode
library(tidyverse)
library(broom)
library(ggthemes)
library(sandwich)
library(lmtest)
library(kableExtra)
library(here)
set.seed(1324)

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_CosineScore, CosineScore), ~ mean(.x), .names="mean_{.col}"),
                across(c(z_CosineScore, CosineScore), ~ sd(.x)/sqrt(n()), .names="se_{.col}")
                ) %>%
      ungroup() %>%
      inner_join(weights, by=strat_vars) %>%
      summarise(across(c(mean_z_CosineScore, mean_CosineScore), ~ weighted.mean(.x, weight)),
                across(c(se_z_CosineScore, se_CosineScore), ~ sqrt(sum((.x**2)*(weight**2))))) %>%
      ungroup()
    )
  } else {
    summary_df <- suppressMessages(data %>%
      summarise(across(c(z_CosineScore, CosineScore), ~ mean(.x), .names="mean_{.col}"),
                across(c(z_CosineScore, CosineScore), ~ 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
    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_CosineScore=sd(mean_z_CosineScore),
              se_clustered_CosineScore=sd(mean_CosineScore))
}

# 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
# replay_file <- '~/research/prompteng_analysis/data/finalized_datasets_for_analysis/main_df_promptLvl_2024-07-29_notWinsorized_PROD.csv'

replay_file <- here("similarity_exp_replication", "inputs",
                       "main_df_promptLvl_2026-03-22_PROD_HASHED_TRIMMED.csv")

replay_df_full <- read_csv(replay_file) %>%
  rename(dalleVersion=dalle_version) %>%
  mutate(dalleVersion=gsub('3-verbatim-prompt', '3', dalleVersion),
         dalleVersion=gsub('3-revised-prompt', '3revised', dalleVersion)) %>%
  filter(iteration < 10)

strat_vars <- c('referenceImageName', 'iteration')
replay_df_full <- replay_df_full %>%
  select(userId, dalleVersion, junk_flag, all_of(strat_vars), matches('.*Replay_mean_.*_score$')) %>%
  pivot_longer(cols=matches('.*Replay_mean_cosine_score$|.*Replay_mean_dreamsim_score$'),
               names_pattern = 'dalle(.*)Replay_mean_(.*)_score',
               names_to=c('replayDalleVersion', 'score')) %>%
  pivot_wider(names_from = score, values_from = value) %>%
  filter(replayDalleVersion %in% c('2', '3verbatim', '3revised')) %>% 
  mutate(replayDalleVersion=gsub('3verbatim', '3', replayDalleVersion)) %>%
  rename(CosineScore=cosine, DreamScore=dreamsim) %>%
  filter(!is.na(CosineScore)) 

# remove prompts with errors
replay_df <- replay_df_full %>% filter(junk_flag==FALSE)

# subset the data to the two main models: 2 and 3verbatim (referred to as 3)
replay_df <- replay_df %>%
  filter(dalleVersion %in% c('2', '3'),
         replayDalleVersion %in% c('2', '3'))
# calculate z-score for comparison between these two models
replay_df <- replay_df %>%
  group_by(across(all_of(strat_vars))) %>%
  mutate(z_CosineScore = scale(CosineScore)[,1],
         z_DreamScore = scale(DreamScore)[,1]) %>%
  ungroup()

# chi-square test of balance across all 45 image/dalle conditons
tmp <- replay_df_full %>%
  group_by(userId) %>%
  summarise(dalleVersion=first(dalleVersion),
            referenceImageName=first(referenceImageName)) %>%
  group_by(dalleVersion, referenceImageName) %>%
  summarise(cnt=n())
chisq.test(tmp$cnt)

################################################################################
######## Mediation analysis through replay (direct/indirect/ATE decomposition in a plot)
# 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

# Plot the progression from base model to forward direct and indirect effect and then back
# cosine 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_CosineScore = mean_z_CosineScore - base_row$mean_z_CosineScore,
         se_z_CosineScore = se_z_CosineScore,
         se_clustered_z_CosineScore = se_clustered_z_CosineScore,
         mean_CosineScore = mean_CosineScore - base_row$mean_CosineScore,
         se_CosineScore = se_CosineScore,
         se_clustered_CosineScore = se_clustered_CosineScore,
         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_CosineScore, default=0), NA),
         arrow_y_start=ifelse(label!='DALL-E 2\nDALL-E 2', lag(mean_CosineScore, default=0), NA),
         arrow_z_y_end=ifelse(label!='DALL-E 2\nDALL-E 2', mean_z_CosineScore, NA),
         arrow_y_end=ifelse(label!='DALL-E 2\nDALL-E 2', mean_CosineScore, 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_CosineScore, mean_CosineScore) %>%
  mutate(pct_change_z_CosineScore = (mean_z_CosineScore/base_row$mean_z_CosineScore - 1)*100,
         pct_change_CosineScore = (mean_CosineScore/base_row$mean_CosineScore - 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


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_CosineScore,
                      ymin = mean_z_CosineScore - 2 * se_clustered_z_CosineScore,
                      ymax = mean_z_CosineScore + 2 * se_clustered_z_CosineScore),
                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 ~ 'Cosine Similarity 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 cosine 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_CosineScore,
                      ymin = mean_CosineScore - 2 * se_clustered_CosineScore,
                      ymax = mean_CosineScore + 2 * se_clustered_CosineScore),
                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 ~ 'Cosine Similarity'), 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(paste0('~/research/prompteng_analysis/results/replay_forward_backward_progression_base_model', base_model, '_raw_score.pdf'), p, width=8, height=6, units='in')

# Create directory if it doesn't exist
dir.create(
  here("similarity_exp_replication", "outputs", "figures"),
  recursive = TRUE,
  showWarnings = FALSE
)

ggsave(
  here("similarity_exp_replication", "outputs", "figures",
       paste0("replay_decomposition_base_model_", base_model, ".pdf")),
  p,
  width = 8,
  height = 6,
  units = 'in'
)

################################################################################
##### 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, referenceImageName, iteration, effect, CosineScore, z_CosineScore)

# z-score results
model_fit <- lm(z_CosineScore ~ iteration + referenceImageName + effect, tmp_data)
model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat(paste('Regression results of z_CosineScore for', effect_type, 'effect with cluster robust standard errors:'))
model_summary_robust
cat('95% confidence interval for z_CosineScore 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(CosineScore ~ iteration + referenceImageName + effect, tmp_data)
model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat(paste('Regression results of CosineScore for', effect_type, 'effect with cluster robust standard errors:'))
model_summary_robust
cat('95% confidence interval for CosineScore 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_CosineScore = mean_z_CosineScore - mean_z_CosineScore1,
         se_clustered_ate_z_CosineScore = sqrt(se_clustered_z_CosineScore**2 + se_clustered_z_CosineScore1**2),
         p_ate_z_CosineScore = 2*(1-pnorm(abs(ate_z_CosineScore), sd=se_clustered_ate_z_CosineScore)),
         ate_z_CosineScore_lower_ci = ate_z_CosineScore - qnorm(0.975)*se_clustered_ate_z_CosineScore,
         ate_z_CosineScore_upper_ci = ate_z_CosineScore + qnorm(0.975)*se_clustered_ate_z_CosineScore,
         ate_CosineScore = mean_CosineScore - mean_CosineScore1,
         se_clustered_ate_CosineScore = sqrt(se_clustered_CosineScore**2 + se_clustered_CosineScore1**2),
         p_ate_CosineScore = 2*(1-pnorm(abs(ate_CosineScore), sd=se_clustered_ate_CosineScore)),
         ate_CosineScore_lower_ci = ate_CosineScore - qnorm(0.975)*se_clustered_ate_CosineScore,
         ate_CosineScore_upper_ci = ate_CosineScore + qnorm(0.975)*se_clustered_ate_CosineScore,
         ) %>%
  select(-matches('^mean_z_CosineScore*|^se_z_CosineScore*|^mean_CosineScore*|^se_CosineScore*|se_clustered_z_CosineScore*|se_clustered_CosineScore*')) %>%
  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_CosineScore, se_clustered_ate_CosineScore, ate_CosineScore_lower_ci, ate_CosineScore_upper_ci, p_ate_CosineScore))
# 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_CosineScore - direct_ate$ate_CosineScore ,' ',
      'and P-value of comparing (backward or forward) total and direct effect:',
      2*(1- pnorm(abs(total_ate$ate_CosineScore - direct_ate$ate_CosineScore),
                  sd=sqrt(total_ate$se_clustered_ate_CosineScore**2 + direct_ate$se_clustered_ate_CosineScore**2))))

################################################################################
##### Direct/Indirect/Total effect by quantile/percentile in distribution using a linear model with Cluster robust SE
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(CosineScore_pct_rank=percent_rank(CosineScore)*100,
         CosineScore_pct_rank_group = sapply(CosineScore_pct_rank, find_closest)) %>%
  ungroup() %>%
  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, referenceImageName, iteration, effect, CosineScore, z_CosineScore,
         CosineScore_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_CosineScore ~ iteration + referenceImageName + effect +
                  as.factor(CosineScore_pct_rank_group) + effect:CosineScore_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_CosineScore for', effect_type, '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(CosineScore ~ iteration + referenceImageName + effect +
                         as.factor(CosineScore_pct_rank_group) + effect:CosineScore_pct_rank_group, tmp_direct)
model_fit_indirect <- lm(CosineScore ~ iteration + referenceImageName + effect +
                         as.factor(CosineScore_pct_rank_group) + effect:CosineScore_pct_rank_group, tmp_indirect)
model_fit_total <- lm(CosineScore ~ iteration + referenceImageName + effect +
                         as.factor(CosineScore_pct_rank_group) + effect:CosineScore_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 CosineScore 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("CosineScore_pct_rank_group", "Cosine Score Rank Group", coefname))
direct_coefs
cat(paste('Regression results of CosineScore 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("CosineScore_pct_rank_group", "Cosine Score Rank Group", coefname))
indirect_coefs
cat(paste('Regression results of CosineScore 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("CosineScore_pct_rank_group", "Cosine 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("similarity_exp_replication", "outputs", "tables", "qte_table.tex"))

################################################################################
##### Performance and ATE variation by iteration
strat_vars <- c('referenceImageName')
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_CosineScore, CosineScore),
                   ~ mean(.x), .names="mean_{.col}"),
            across(c(z_CosineScore, CosineScore),
                   ~ 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_CosineScore, mean_CosineScore), ~ weighted.mean(.x, weight)),
            across(c(se_z_CosineScore, se_CosineScore), ~ 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_CosineScore, se_z_CosineScore, mean_CosineScore, se_CosineScore),
              names_glue='{.value}_dalle{dalleVersion}') %>%
  mutate(ate_z_CosineScore = mean_z_CosineScore_dalle3 - mean_z_CosineScore_dalle2,
         se_ate_CosineScore = sqrt(se_z_CosineScore_dalle3**2 + se_z_CosineScore_dalle2**2),
         ate_CosineScore = mean_CosineScore_dalle3 - mean_CosineScore_dalle2,
         se_ate_CosineScore = sqrt(se_CosineScore_dalle3**2 + se_CosineScore_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_CosineScore)) +
  geom_point(position=position_dodge(width=0.3), size=point_size, color='black') + 
  geom_errorbar(aes(ymin=ate_CosineScore - 1*se_ate_CosineScore,
                    ymax=ate_CosineScore + 1*se_ate_CosineScore),
                linewidth=error_bar_size, width=0.2, position = position_dodge(width = 0.3)) +
  theme_clean() +
  labs(x='Iteration', y='ATE Cosine Similarity 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())
p

p <- ggplot(summary_df, aes(as.factor(iteration), mean_CosineScore, shape=dalleVersion)) +
  geom_point(position=position_dodge(width=0.3), size=point_size, color='black') + 
  geom_pointrange(aes(ymin=mean_CosineScore - 2*se_CosineScore,
                      ymax=mean_CosineScore + 2*se_CosineScore),
                  position=position_dodge(width=0.3)) +
  labs(x='Iteration', y='Cosine Similarity 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, referenceImageName, effect, CosineScore) %>%
  group_by(referenceImageName) %>%
  mutate(z_CosineScore = scale(CosineScore)[,1]) %>%
  ungroup()

# z-score
model_fit <- lm(z_CosineScore ~ effect*iteration + referenceImageName + 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_CosineScore with cluster robust standard errors:'))
model_summary_robust

# raw score
model_fit <- lm(CosineScore ~ effect*iteration + referenceImageName + 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 CosineScore with cluster robust standard errors:'))
model_summary_robust


################################################################################
##### Revised Dalle3 against a base model (either Dalle2 or Dalle3 revised)
strat_vars <- c('referenceImageName', '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_CosineScore = scale(CosineScore)[,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, referenceImageName, iteration, effect, CosineScore, z_CosineScore)

dep_var <- 'CosineScore'
#dep_var <- 'z_CosineScore'
model_fit <- lm(as.formula(paste(dep_var, ' ~ iteration + referenceImageName + 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

### 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_CosineScore = mean_z_CosineScore - mean_z_CosineScore1,
         se_clustered_ate_z_CosineScore = sqrt(se_clustered_z_CosineScore**2 + se_clustered_z_CosineScore1**2),
         p_ate_z_CosineScore = 2*(1-pnorm(abs(ate_z_CosineScore), sd=se_clustered_ate_z_CosineScore)),
         ate_z_CosineScore_lower_ci = ate_z_CosineScore - qnorm(0.975)*se_clustered_ate_z_CosineScore,
         ate_z_CosineScore_upper_ci = ate_z_CosineScore + qnorm(0.975)*se_clustered_ate_z_CosineScore,
         ate_CosineScore = mean_CosineScore - mean_CosineScore1,
         se_clustered_ate_CosineScore = sqrt(se_clustered_CosineScore**2 + se_clustered_CosineScore1**2),
         p_ate_CosineScore = 2*(1-pnorm(abs(ate_CosineScore), sd=se_clustered_ate_CosineScore)),
         ate_CosineScore_lower_ci = ate_CosineScore - qnorm(0.975)*se_clustered_ate_CosineScore,
         ate_CosineScore_upper_ci = ate_CosineScore + qnorm(0.975)*se_clustered_ate_CosineScore,
         ) %>%
  select(-matches('^mean_z_CosineScore*|^se_z_CosineScore*|^mean_CosineScore*|^se_CosineScore*|se_clustered_z_CosineScore*|se_clustered_CosineScore*')) %>%
  filter(dalleVersion!=dalleVersion1 & dalleVersion1==base_model)

# results based on z-score
ate_df %>% select(c(ate_z_CosineScore, se_clustered_ate_z_CosineScore, ate_z_CosineScore_lower_ci, ate_z_CosineScore_upper_ci, p_ate_z_CosineScore))
# results based on raw-score
ate_df %>% select(c(ate_CosineScore, se_clustered_ate_CosineScore, ate_CosineScore_lower_ci, ate_CosineScore_upper_ci, p_ate_CosineScore))


################################################################################
##### 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, referenceImageName, iteration, CosineScore)
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 <- 'CosineScore'
  # dalle 3 coef/se
  model_fit <- lm(as.formula(paste(dep_var, ' ~ iteration + referenceImageName + 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 + referenceImageName + 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 %>%
    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, ")")

################################################################################
##### Results based on Dreasim and Regression
# reverse these two base/other models to get backward
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)

strat_vars <- c('referenceImageName', 'iteration')
replay_df <- replay_df_full %>% filter(junk_flag==FALSE)
replay_df <- replay_df %>%
  filter(dalleVersion %in% c(base_model, '3'),
         replayDalleVersion %in% c(base_model, '3'))
# calculate z-score for comparison between these two models
replay_df <- replay_df %>%
  group_by(across(all_of(strat_vars))) %>%
  mutate(z_DreamScore = scale(DreamScore)[,1]) %>%
  ungroup()

##########
# 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, referenceImageName, iteration, effect, DreamScore, z_DreamScore)

# z-score dream sim
model_fit <- lm(z_DreamScore ~ iteration + referenceImageName + effect, tmp_data)
model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat(paste('Regression results for', effect_type, 'effect on z_DreamScore with cluster robust standard errors:'))
model_summary_robust
cat('95% confidence interval 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 dream sim
model_fit <- lm(DreamScore ~ iteration + referenceImageName + effect, tmp_data)
model_summary_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat(paste('Regression results for', effect_type, 'effect on DreamScore with cluster robust standard errors:'))
model_summary_robust
cat('95% confidence interval 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"], ")")