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

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")
full_df <- read_csv(image_file) %>%
  rename(dalle_version=dalleVersion) %>%
  mutate(dalle_version= str_replace_all(dalle_version, repl_map),
         prompt_length=map_int(prompt, ~ length(tokenize_words(.x)[[1]])),
         revised_prompt_length=map_int(revisedPrompt, ~ length(tokenize_words(.x)[[1]])),
         log_prompt_length=log(prompt_length),
         dalle_txt=paste('DALL-E', dalle_version),
         iteration=iteration+1,
         StartDate = as.POSIXct(StartDate, format = '%Y-%m-%d %H:%M:%S', tz = 'UTC'),
         EndDate   = as.POSIXct(EndDate, format = '%Y-%m-%d %H:%M:%S', tz = 'UTC'),
         duration  = as.numeric(difftime(EndDate, StartDate, units = 'secs'))
         )

# get stats on how long it took for participants
full_df %>% group_by(dalle_version) %>%
  summarise(across(duration, list(mean = mean, median = median, max = max, min = min),
                   .names = '{fn}_time'))
full_df %>% summarise(across(duration, list(mean = mean, median = median, max = max, min = min),
                        .names = '{fn}_time'))

########################
# compute cosine similarity of the prompt to its previous iteration
# 1. Build a global TF–IDF matrix for all prompts
it <- itoken(full_df$prompt, 
             tokenizer = word_tokenizer, 
             progressbar = FALSE)
vocab      <- create_vocabulary(it)
vectorizer <- vocab_vectorizer(vocab)
dtm        <- create_dtm(it, vectorizer)
tfidf_mat  <- TfIdf$new()$fit_transform(dtm)

# helper to compute cosine between two rows of a dgCMatrix
cosine_sim <- function(mat, i, j) {
  ui <- mat[i, ]
  uj <- mat[j, ]
  num <- sum(ui * uj)
  den <- sqrt(sum(ui * ui) * sum(uj * uj))
  if (den == 0) NA_real_ else num/den
}
# helper: only compute cosine when we have a true “prev iteration = iter − 1” in the same group
safe_cosine <- function(curr_row, prev_row,
                        user,   prev_user,
                        ver,    prev_ver,
                        iter,   prev_iter) {
  if (is.na(prev_row) ||
      user    != prev_user ||
      ver     != prev_ver  ||
      (iter - prev_iter) != 1) {
    return(NA_real_)
  }
  cosine_sim(tfidf_mat, curr_row, prev_row)
}
safe_contains <- function(curr_prompt, prev_prompt,
                          user,   prev_user,
                          ver,    prev_ver,
                          iter,   prev_iter) {
  if (is.na(prev_prompt) ||
      user    != prev_user ||
      ver     != prev_ver  ||
      (iter - prev_iter) != 1) {
    return(NA)  # logical NA
  }
  str_detect(curr_prompt, fixed(prev_prompt))
}

# 2. Join TF–IDF back to df via row indices, then compute lagged cosine
# row_id must be assigned BEFORE sorting so it stays aligned with tfidf_mat
cat("🔄 Computing pairwise cosine similarities (this may take several minutes)...\n")
full_df <- full_df %>%
  mutate(row_id=row_number()) %>%
  arrange(userId, dalle_version, iteration) %>%
  mutate(prev_row=lag(row_id),
         prev_iter=lag(iteration),
         prev_prompt= lag(prompt),
         prev_user=lag(userId),
         prev_ver=lag(dalle_version),
         cosine_sim_prev_prompt = pmap_dbl(list(row_id,        prev_row,
                                                userId,        prev_user,
                                                dalle_version, prev_ver,
                                                iteration,     prev_iter),
                                           safe_cosine),
         contains_prev_prompt = pmap_lgl(list(prompt,          prev_prompt,
                                              userId,          prev_user,
                                              dalle_version,   prev_ver,
                                              iteration,       prev_iter),
                                         safe_contains)) %>%
  select(-row_id, -prev_row, -prev_iter, -prev_prompt, -prev_user, -prev_ver)
########################

# subset the data to the two main models
#valid_models <- c('2', '3', '3revised')
#valid_models <- c('3', '3revised')
#valid_models <- c('2', '3revised')
valid_models <- c('2', '3')
df <- full_df %>%
  filter(dalle_version %in% valid_models)

# take average per userid,logo,iteration
# 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, logoTaskSelection) %>%
  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()

################################################################################
### cosine sim and prompt length by model iteration
strat_vars <- c('logoTaskSelection')
dep_vars <- c('cosine_sim_prev_prompt', 'prompt_length', 'log_prompt_length', 'contains_prev_prompt',
              'longer_prompt', 'longer_contains_prev_prompt')
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',
    iteration == 1 & dalle_txt == 'DALL-E 3revised' ~ 'DALL-E3revised_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' & !str_detect(dalle_txt, "_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(str_detect(dalle_txt, '_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.4, 0.45, 0.5)) +
  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("logo_exp_replication", "outputs", "figures"), recursive = TRUE, showWarnings = FALSE)
ggsave(here("logo_exp_replication", "outputs", "figures", "prompt_length_similarity_by_model_task2.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 ~ dalle_version + logoTaskSelection + 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 + logoTaskSelection, 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 + logoTaskSelection + 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 + logoTaskSelection + 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
cat("🔄 Loading udpipe model (will download on first run)...\n")
library(udpipe)
model_file <- udpipe_download_model(language = 'english')$file_model
ud_model <- udpipe_load_model(model_file)
cat("✅ udpipe model loaded\n")

cat("🔄 Annotating", nrow(df), "prompts with POS tags...\n")
batch_size <- 1000
n_batches <- ceiling(nrow(df) / batch_size)
anno_list <- vector("list", n_batches)
for (i in seq_len(n_batches)) {
  idx <- ((i - 1) * batch_size + 1):min(i * batch_size, nrow(df))
  anno_list[[i]] <- as.data.frame(udpipe_annotate(ud_model,
                                                   x = df$prompt[idx],
                                                   doc_id = df$image_id[idx]))
  cat(sprintf("\r  Progress: %d/%d batches (%d%%)  ", i, n_batches, round(100 * i / n_batches)))
}
cat("\n")
anno_df <- bind_rows(anno_list)
cat("✅ POS annotation complete:", nrow(anno_df), "tokens\n")

pos_df <- anno_df %>%
  # keep word tokens only (drop punctuation etc.)
  filter(!is.na(upos), upos != 'PUNCT', str_detect(token, '[A-Za-z0-9]')) %>%
  group_by(doc_id) %>%
  summarise(prompt_length_words = n(),
            pos_ADJ_count = sum(upos == 'ADJ'),
            pos_ADV_count = sum(upos == 'ADV'),
            pos_VERB_count = sum(upos == 'VERB'),
            pos_NOUN_count = sum(upos == 'NOUN')) %>%
  ungroup() %>%
  mutate(pos_ADJ_count_norm  = pos_ADJ_count  / prompt_length_words,
         pos_ADV_count_norm  = pos_ADV_count  / prompt_length_words,
         pos_VERB_count_norm = pos_VERB_count / prompt_length_words,
         pos_NOUN_count_norm = pos_NOUN_count / prompt_length_words,
         pos_ADJ_NOUN_norm   = pos_ADJ_count_norm + pos_NOUN_count_norm) %>%
  rename(image_id = doc_id)
pos_df <- pos_df %>% inner_join(df %>% select(userId, image_id, dalle_version), by='image_id') %>%
  mutate(dalle_version=gsub('3-verbatim-prompt', '3', dalle_version))

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

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

