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

# cosine_sim_file <- '~/research/prompteng_analysis/data/finalized_datasets_for_analysis/main_df_promptLvl_2024-07-29_notWinsorized_PROD.csv'

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

df <- read_csv(cosine_sim_file) %>%
  filter(dalle_version %in% c('2', '3-verbatim-prompt')) %>%
  mutate(dalle_version=gsub('3-verbatim-prompt', '3', dalle_version),
         prompt_length=prompt_length_words,
         log_prompt_length=log(prompt_length),
         dalle_txt=paste('DALL-E', dalle_version),
         iteration=iteration+1)
df <- df %>% filter(junk_flag == FALSE)
 
# take average per userid,referenceimage,iteration
df <- df %>% group_by(referenceImageName, iteration) %>%
  mutate(z_cosine_sim_mean = scale(cosine_sim_mean)[,1]) %>%
  ungroup()
# average prompt length per model-iteration
df <- df %>%
  group_by(dalle_version, iteration) %>%
  mutate(mean_cosine_sim_prev_prompt=mean(cosine_sim_prev_prompt)) %>%
  ungroup()
# difference in prompt length
df <- df %>%
  group_by(userId, referenceImageName) %>%
  arrange(iteration) %>%
  mutate(prev_prompt_length = ifelse(lag(iteration)==iteration-1, lag(prompt_length), NA),
         prompt_length_diff = prompt_length - prev_prompt_length,
         longer_prompt=(prompt_length_diff > 0),
         longer_contains_prev_prompt=(longer_prompt & contains_prev_prompt)) %>%
  ungroup()

cor.test(df$prompt_length_words, df$prompt_length_chars)

################################################################################
### cosine sim and prompt length by model iteration
strat_vars <- c('referenceImageName')
dep_vars <- c('cosine_sim_prev_prompt', 'prompt_length', 'log_prompt_length', 'contains_prev_prompt',
              'longer_prompt', 'longer_contains_prev_prompt',
              'dist_to_user_centroid_junkFiltered')
weights <- df %>%
  group_by(across(all_of(strat_vars))) %>%
  summarise(cnt=n()) %>%
  ungroup() %>%
  mutate(weight=cnt/sum(cnt)) %>%
  select(-cnt)
summary_df <- df %>%
  group_by(across(c(all_of(strat_vars), dalle_txt, iteration))) %>%
  summarise(across(all_of(dep_vars), ~ mean(.x, na.rm=TRUE), .names="mean_{.col}"),
            across(all_of(dep_vars), ~ sd(.x, na.rm=TRUE)/sqrt(sum(!is.na(.x))), .names="se_{.col}")) %>%
  ungroup() %>%
  inner_join(weights, by=strat_vars) %>%
  group_by(dalle_txt, iteration) %>%
  summarise(across(all_of(paste0('mean_', dep_vars)), ~ weighted.mean(.x, weight)),
            across(all_of(paste0('se_', dep_vars)), ~ sqrt(sum((.x**2)*(weight**2)))),
            sum_weight=sum(weight)) %>% 
  ungroup()

summary_df <- summary_df %>%
  mutate(dalle_txt = case_when(
    iteration == 1 & dalle_txt == 'DALL-E 2' ~ 'DALL-E2_init',
    iteration == 1 & dalle_txt == 'DALL-E 3' ~ 'DALL-E3_init',
    TRUE ~ dalle_txt
  )) %>%
  mutate(iteration_numeric = as.numeric(iteration))
summary_df$iteration <- as.factor(summary_df$iteration)

error_bar_size <- 1
point_size <- 5
axis_text_size <- 21
axis_title_size <- 21
legend_text_size <- 19
dodge <- position_dodge(width = 0.4)

p <- ggplot() +
  geom_errorbar(data = summary_df,
                aes(x = iteration, y = mean_prompt_length,
                    ymin = mean_prompt_length - 2 * se_prompt_length,
                    ymax = mean_prompt_length + 2 * se_prompt_length,
                    color = mean_cosine_sim_prev_prompt, group = dalle_txt),
                width = 0, size = error_bar_size, position = dodge) +
  geom_point(data = summary_df %>% filter(iteration != '0' & dalle_txt != 'DALL-E2_init' & dalle_txt != 'DALL-E3_init'),
             aes(x = iteration, y = mean_prompt_length, color = mean_cosine_sim_prev_prompt, shape = dalle_txt),
             size = point_size, position = dodge) +
  geom_point(data = summary_df %>% filter(dalle_txt == 'DALL-E2_init' | dalle_txt == 'DALL-E3_init'),
             aes(x = iteration, y = mean_prompt_length, shape = dalle_txt),
             color = "#969696", size = point_size, stroke = 1, position = dodge,
             show.legend = FALSE) +
  scale_color_viridis(option = 'D', breaks = c(0.75, 0.8, 0.85)) +
  scale_shape_manual(values = c('DALL-E 2' = 16, 'DALL-E 3' = 17, 'DALL-E2_init' = 21, 'DALL-E3_init' = 24),
                     breaks = c('DALL-E 2', 'DALL-E 3')) +    # don't show the category of first iterations
  labs(x = 'Attempt Number', y = 'Prompt Length in Words', color = 'Similarity to\nPrevious Prompt') +
  guides(color = guide_colorbar(order = 1),
         shape = guide_legend(order = 2, title = NULL)) +
  theme_clean() +
  theme(axis.text = element_text(colour = "black", size = axis_text_size),
        axis.title = element_text(colour = "black", size = axis_title_size),
        panel.grid.major.y = element_blank(), 
        panel.grid.minor.y = element_blank(), 
        legend.title = element_text(size = legend_text_size),
        legend.text = element_text(size = legend_text_size),
        legend.position = "inside",
        legend.position.inside = c(0.26, 0.75),
        legend.box = "horizontal",
        legend.background = element_blank(),
        plot.background = element_blank())
p
dir.create(here("similarity_exp_replication", "outputs", "figures"), recursive = TRUE, showWarnings = FALSE)
ggsave(here("similarity_exp_replication", "outputs", "figures", "prompt_length_similarity_by_model.pdf"), p, width=9, height=4.5, units='in')

################################################################################
# compare prompt characteristics (e.g. length) from both models across all iterations

# linear model with cluster robust standard errors: does the prompt length increase by model?
tmp_df <- df %>% mutate(iteration_factor=as.factor(iteration))
model_fit <- lm(prompt_length_words ~ dalle_version + referenceImageName + iteration_factor, tmp_df)
model_fit_robust <- coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")
cat(paste('Regression results of prompt length on dalle version with cluster robust standard errors:'))
model_fit_robust
cat(paste('Percent change in prompt_length due to dalle_version:'))
model_fit_robust["dalle_version3", "Estimate"]/mean(tmp_df[tmp_df$dalle_version=='2',]$prompt_length)*100

# does the prompt length difference increase through iterations? (interaction term)
model_fit <- lm(prompt_length ~ dalle_version*iteration + referenceImageName, tmp_df)
cat(paste('Regression results of prompt length on dalle version across iterations with cluster robust standard errors:'))
coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")

# does the prompt similarity increase by model?
model_fit <- lm(cosine_sim_prev_prompt ~ dalle_version + referenceImageName + iteration_factor, tmp_df %>% filter(iteration>1))
cat(paste('Regression results of prompt similarity on dalle version across iterations with cluster robust standard errors:'))
coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")

# prompt convergence: regression analysis of distance to task centroid by model 
model_fit <- lm(dist_to_user_centroid_junkFiltered ~ dalle_version + referenceImageName + iteration_factor, tmp_df)
cat(paste('Regression results of prompt distance to user centroid on dalle version with cluster robust standard errors:'))
coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")


# robustness check: stratification without cluster robust errors
iter_weights <- df %>%
  mutate(iteration=as.factor(iteration)) %>%
  filter(iteration != '0') %>%
  group_by(iteration) %>%
  summarise(cnt=n()) %>%
  ungroup() %>%
  mutate(weight=cnt/sum(cnt)) %>%
  select(-cnt)
res <- summary_df %>%
  inner_join(iter_weights, by='iteration') %>%
  group_by(dalle_txt) %>%
  summarise(across(all_of(paste0('mean_', dep_vars)), ~ weighted.mean(.x, weight)),
            across(all_of(paste0('se_', dep_vars)), ~ sqrt(sum((.x**2)*(weight**2)))),
            sum_weight=sum(weight)) %>% 
  ungroup() %>%
  summarise(across(matches("^mean_"), ~ .[dalle_txt == 'DALL-E 2'], .names = 'base_{col}'),
            across(matches("^mean_"), ~ .[dalle_txt == 'DALL-E 3'] - .[dalle_txt == 'DALL-E 2'], .names = 'ate_{col}'),
            across(matches("^se_"), ~ sqrt(sum(.^2)), .names = 'se_ate_{col}')) %>%
  pivot_longer(cols = everything(), 
               names_to = c(".value", "column"), 
               names_pattern = "(^base_mean|^ate_mean|^se_ate_se)_(.*)") %>%
  rename(ate=ate_mean, se_ate=se_ate_se) %>%
  mutate(pct_change=ate*100/base_mean,
         z=ate/se_ate,
         p_ate = 2*(1-pnorm(abs(ate), sd=se_ate)))
res

################################################################################
#### test part of speech frequency by model
# pos_file <- '~/research/prompteng_analysis/data/finalized_datasets_for_analysis/part_of_speech.csv'
# pos_file = here("similarity_exp_replication", "inputs", "part_of_speech.csv")
pos_file = here("similarity_exp_replication", "inputs", "part_of_speech_HASHED_TRIMMED.csv")

pos_df <- read_csv(pos_file) %>%
  select(c(image_id, dalleVersion, prompt_length_words, pos_ADJ_count_norm,
           pos_ADV_count_norm, pos_VERB_count_norm, pos_NOUN_count_norm)) %>%
  mutate(dalleVersion=gsub('3-verbatim-prompt', '3', dalleVersion),
         pos_ADJ_NOUN_norm=pos_ADJ_count_norm + pos_NOUN_count_norm) %>%
  inner_join(df %>% select(userId, originalID), by=join_by(image_id == originalID))

# simple avg of adj/noun per model
pos_df %>% group_by(dalleVersion) %>% summarise(mean(pos_ADJ_NOUN_norm))

# linear model
model_fit <- lm(pos_ADJ_NOUN_norm ~ dalleVersion, pos_df)
coeftest(model_fit, vcov=vcovCL, cluster=~userId, type = "HC1")