pdf(NULL) # suppress Rplots.pdf in batch mode
library(ggplot2)
library(stringr)
library(lmtest)
library(broom)
library(sandwich)
library(quickmatch)
library(tidyverse)
library(here)

source(here("shared", "utils.R"))
set.seed(1324)
user_iter_file <- here("similarity_exp_replication", "inputs",
                       "main_df_promptLvl_2026-03-22_PROD_HASHED_TRIMMED.csv")

# source("~/research/prompteng_analysis/replication/utils.R")
# user_iter_file <- '~/research/prompteng_analysis/data/finalized_datasets_for_analysis/main_df_promptLvl_2024-07-29_notWinsorized_PROD.csv'
#user_iter_file <- '~/research/prompteng_analysis/data/finalized_datasets_for_analysis/main_df_promptLvl_2024-07-24_winsorized.csv'
user_iter_df <- read_csv(user_iter_file) %>% filter(junk_flag==FALSE)

# Add extra iteration features now using only valid iterations and generate user-level and nlp data
result <- generate_data_from_user_iter(user_iter_df)
user_iter_df <- result$user_iter_df
user_df <- result$user_df

dep_vars <- c('mean_cosine_sim_mean', 'mean_dream_sim_mean',
              'z_mean_cosine_sim_mean', 'z_mean_dream_sim_mean',
              'pct_rank_mean_cosine_sim_mean', 'pct_rank_mean_dream_sim_mean')
iter_dep_vars <- c('cosine_sim_mean', 'dream_sim_mean',
                   'z_cosine_sim_mean', 'z_dream_sim_mean',
                   'pct_rank_cosine_sim_mean', 'pct_rank_dream_sim_mean')

nlp_iter_dep_vars <- c('tsr_prev_prompt', 'cosine_sim_prev_prompt', 'contains_prev_prompt',
                       'prompt_length_words_diff', 'prompt_length_chars_diff')
nlp_dep_vars <- c('embedding_var',
                  'n_total_topics_tsr_tasklvl_mean', 'n_total_topics_cosine_tasklvl_mean',
                  'mean_tsr_prev_prompt', 'mean_cosine_sim_prev_prompt',
                  'mean_contains_prev_prompt',
                  'mean_prompt_length_chars', 'mean_prompt_length_words')

user_dep_vars <- c(dep_vars, nlp_dep_vars)
iter_dep_vars <- c(iter_dep_vars, nlp_iter_dep_vars)


################################################################################
# user-iter-level analysis

##################
# Method 1: Post-stratification z-test in binned groups of observed z score
strat_vars <- c('dalle_version')
indep_var <- 'cosine_sim_observed_z_prev'
#indep_var <- 'dream_sim_observed_z_prev'

tmp <- user_iter_df %>% filter(iteration > 0 & !is.na(.data[[indep_var]]))
# remove observations with crazy observed z (bottom and top 0.5%)
observed_z_threshold <- 3
orig_size <- nrow(tmp)
tmp <- tmp %>% filter(abs(.data[[indep_var]]) < observed_z_threshold)
new_size <- nrow(tmp)
paste0('Removed ', orig_size - new_size, ' strange observations (', round((orig_size - new_size)*100/orig_size, 2), '%)')
indep_var <- paste0(indep_var, '_bracket')
cat('Independent varibale is', indep_var)

weights <- tmp %>%
  group_by(across(all_of(strat_vars))) %>%
  summarise(cnt=n()) %>%
  ungroup() %>%
  mutate(weight=cnt/sum(cnt))

summary_df <- tmp %>% 
  group_by(across(all_of(c(strat_vars, indep_var)))) %>%
  summarise(across(all_of(iter_dep_vars),
                   ~ mean(.x, na.rm=T), .names="mean_{.col}"),
            across(all_of(iter_dep_vars),
                   ~ sd(.x, na.rm=T)/sqrt(n()), .names="se_{.col}")
            ) %>%
  ungroup() %>%
  inner_join(weights, by=strat_vars) %>%
  group_by(across(all_of(indep_var))) %>%
  summarise(across(all_of(paste0('mean_', iter_dep_vars)), ~ weighted.mean(.x, weight)),
            across(all_of(paste0('se_', iter_dep_vars)), ~ sqrt(sum((.x**2)*(weight**2)))),
            sum_weight=sum(weight))
assertthat::assert_that(all(summary_df$sum_weight > 0.999 & summary_df$sum_weight < 1.001))
summary_df

# compute pairwise ate p-values
pairs_df <- crossing(bracket1=summary_df[[indep_var]], bracket2=summary_df[[indep_var]]) %>%
  inner_join(summary_df, by=c('bracket1' = indep_var)) %>%
  inner_join(summary_df, by=c('bracket2' = indep_var), suffix=c('_1', '_2')) %>%
  mutate(across(matches("^mean_.*_1$"), ~ . - get(sub("_1$", "_2", cur_column())), .names='{col}_ate'),
         across(matches("^se_.*_1$"), ~ sqrt(.**2 + get(sub("_1$", "_2", cur_column()))**2), .names='{col}_se'),
         bracket1_bracket2=paste0(bracket1, '-', bracket2)) %>%
  rename_with(~sub("^mean_(.*)_1_ate$", "\\1_ate", .)) %>%
  rename_with(~sub("^se_(.*)_1_se$", "\\1_se", .)) %>%
  mutate(across(matches(".*_ate"), ~ 2*(1 - pnorm(abs(./get(sub("_ate", "_se", cur_column()))))),
                .names='{col}_p_val')) %>%
  rename_with(~sub("(.*)_ate_p_val", "\\1_p_val", .)) %>%
  select(matches("(bracket1|bracket2|_ate|_se|_p_val)$")) %>%
  filter(bracket2 != bracket1)
print(n=50,
      pairs_df %>%
        pivot_longer(cols=-c(bracket1, bracket2),
                     names_to = c("metric", ".value"),
                     names_pattern = "(.*)_(ate|se|p_val)") %>%
        filter(p_val < 0.05) %>% arrange(metric))

# plot average outcomes of each bracket
idx <- 1
iter_dep_var <- iter_dep_vars[idx]
mean_dep_var <- paste0('mean_', iter_dep_var)
se_dep_var <- paste0('se_', iter_dep_var)
p_vals_title <- pairs_df %>%
  filter(as.character(bracket1) > as.character(bracket2)) %>%
  mutate(across(matches(".*_p_val"), ~ round(., 4))) %>%
  select(all_of(c('bracket1_bracket2', paste0(iter_dep_var, '_p_val')))) %>%
  unite(row_string, everything(), sep="  p value: ") %>%
  summarise(final_string = str_c(row_string, collapse = "\n"))
ggplot(summary_df, aes(x=.data[[indep_var]], y=.data[[mean_dep_var]])) +
  geom_point(position = position_dodge(width = 0.4)) +
  geom_errorbar(aes(ymin=.data[[mean_dep_var]] - 2*.data[[se_dep_var]],
                    ymax=.data[[mean_dep_var]] + 2*.data[[se_dep_var]]),
                width=0.2, position = position_dodge(width = 0.4)) +
  theme_bw() +
  ggtitle(p_vals_title) +
  theme(axis.text=element_text(colour="black", size=15),
        axis.title=element_text(colour="black", size=15),
        legend.title=element_text(size=12),
        legend.text=element_text(size=10),
        plot.title = element_text(size = 20, hjust=0.5))
pairs_df %>% select(c(matches('bracket'), matches(paste0('^',iter_dep_var))))

###################
# Method 2: OLS
# overall regression
indep_var_reg <- function(x, reg_df) {
  res <- lm(as.formula(paste(x, '~ dalle_version +', indep_var)), data=reg_df)
  res <- coeftest(res, vcov=vcovHC(res, type="HC1"))
  res <- tidy(res) %>%
    filter(term == indep_var) %>%
    rename(indep_var=term) %>%
    mutate(dep_var=x) %>%
    select(dep_var, indep_var, estimate, p.value)
  return(res)
}
indep_var <- 'cosine_sim_observed_z_prev'
res <- list_rbind(map(iter_dep_vars, \(x) indep_var_reg(x, reg_df=tmp))) %>%
  arrange(p.value)
cat('Coefs of regressing dep vars on', indep_var)
res %>% filter(p.value < 0.05)

# regression within pos and neg observed z
neg_data <- tmp %>% filter(.data[[indep_var]] <= 0)
pos_data <- tmp %>% filter(.data[[indep_var]] > 0)
neg_res <- list_rbind(map(iter_dep_vars, \(x) indep_var_reg(x, reg_df=neg_data))) %>%
  arrange(p.value) %>%
  mutate(group = "negative")
pos_res <- list_rbind(map(iter_dep_vars, \(x) indep_var_reg(x, reg_df=pos_data))) %>%
  arrange(p.value) %>%
  mutate(group = "positive")
res <- bind_rows(neg_res, pos_res)
cat('Coefs of regressing dep vars on positive/negative ranges of', indep_var)
res %>% filter(p.value < 0.05) %>% arrange(dep_var)


################################################################################
# user-level analysis

###################
# Method 1: Post-stratification z-test in binned groups of observed z score
#strat_vars <- c('dalle_version', 'referenceImageName', 'iteration')
strat_vars <- c('dalle_version')
indep_var <- 'first_cosine_sim_observed_z'
#indep_var <- 'first_dream_sim_observed_z'

tmp <- user_df %>% filter(!is.na(.data[[indep_var]]))
# remove observations with crazy observed z (bottom and top 0.5%)
observed_z_threshold <- 3
orig_size <- nrow(tmp)
tmp <- tmp %>% filter(abs(.data[[indep_var]]) < observed_z_threshold)
new_size <- nrow(tmp)
paste0('Removed ', orig_size - new_size, ' strange observations (', round((orig_size - new_size)*100/orig_size, 2), '%)')
indep_var <- paste0(indep_var, '_bracket')


weights <- tmp %>%
  group_by(across(all_of(strat_vars))) %>%
  summarise(cnt=n()) %>%
  ungroup() %>%
  mutate(weight=cnt/sum(cnt))

summary_df <- tmp %>% 
  group_by(across(all_of(c(strat_vars, indep_var)))) %>%
  summarise(across(all_of(user_dep_vars),
                   ~ mean(.x, na.rm=T), .names="mean_{.col}"),
            across(all_of(user_dep_vars),
                   ~ sd(.x, na.rm=T)/sqrt(n()), .names="se_{.col}")
            ) %>%
  ungroup() %>%
  inner_join(weights, by=strat_vars) %>%
  group_by(across(all_of(indep_var))) %>%
  summarise(across(all_of(paste0('mean_', user_dep_vars)), ~ weighted.mean(.x, weight)),
            across(all_of(paste0('se_', user_dep_vars)), ~ sqrt(sum((.x**2)*(weight**2)))),
            sum_weight=sum(weight))
assertthat::assert_that(all(summary_df$sum_weight > 0.999 & summary_df$sum_weight < 1.001))
summary_df

# compute pairwise ate p-values
pairs_df <- crossing(bracket1=summary_df[[indep_var]], bracket2=summary_df[[indep_var]]) %>%
  inner_join(summary_df, by=c('bracket1' = indep_var)) %>%
  inner_join(summary_df, by=c('bracket2' = indep_var), suffix=c('_1', '_2')) %>%
  mutate(across(matches("^mean_.*_1$"), ~ . - get(sub("_1$", "_2", cur_column())), .names='{col}_ate'),
         across(matches("^se_.*_1$"), ~ sqrt(.**2 + get(sub("_1$", "_2", cur_column()))**2), .names='{col}_se'),
         bracket1_bracket2=paste0(bracket1, '-', bracket2)) %>%
  rename_with(~sub("^mean_(.*)_1_ate$", "\\1_ate", .)) %>%
  rename_with(~sub("^se_(.*)_1_se$", "\\1_se", .)) %>%
  mutate(across(matches(".*_ate"), ~ 2*(1 - pnorm(abs(./get(sub("_ate", "_se", cur_column()))))),
                .names='{col}_p_val')) %>%
  rename_with(~sub("(.*)_ate_p_val", "\\1_p_val", .)) %>%
  select(matches("(bracket1|bracket2|_ate|_se|_p_val)$")) %>%
  filter(bracket2 != bracket1)
print(n=50,
      pairs_df %>%
        pivot_longer(cols=-c(bracket1, bracket2),
                     names_to = c("metric", ".value"),
                     names_pattern = "(.*)_(ate|se|p_val)") %>%
        filter(p_val < 0.05) %>% arrange(metric))

# plot average outcomes of each bracket
idx <- 1
user_dep_var <- user_dep_vars[idx]
mean_dep_var <- paste0('mean_', user_dep_var)
se_dep_var <- paste0('se_', user_dep_var)
p_vals_title <- pairs_df %>%
  filter(as.character(bracket1) > as.character(bracket2)) %>%
  mutate(across(matches(".*_p_val"), ~ round(., 4))) %>%
  select(all_of(c('bracket1_bracket2', paste0(user_dep_var, '_p_val')))) %>%
  unite(row_string, everything(), sep="  p value: ") %>%
  summarise(final_string = str_c(row_string, collapse = "\n"))
ggplot(summary_df, aes(x=.data[[indep_var]], y=.data[[mean_dep_var]])) +
  geom_point(position = position_dodge(width = 0.4)) +
  geom_errorbar(aes(ymin=.data[[mean_dep_var]] - 2*.data[[se_dep_var]],
                    ymax=.data[[mean_dep_var]] + 2*.data[[se_dep_var]]),
                width=0.2, position = position_dodge(width = 0.4)) +
  theme_bw() +
  ggtitle(p_vals_title) +
  theme(axis.text=element_text(colour="black", size=15),
        axis.title=element_text(colour="black", size=15),
        legend.title=element_text(size=12),
        legend.text=element_text(size=10),
        plot.title = element_text(size = 20, hjust=0.5))
pairs_df %>% select(c(matches('bracket'), matches(paste0('^',user_dep_var))))

###################
# Method 2: OLS
# overall regression
indep_var_reg <- function(x, reg_df) {
  res <- lm(as.formula(paste(x, '~ dalle_version +', indep_var)), data=reg_df)
  res <- coeftest(res, vcov=vcovHC(res, type="HC1"))
  res <- tidy(res) %>%
    filter(term == indep_var) %>%
    rename(indep_var=term) %>%
    mutate(dep_var=x) %>%
    select(dep_var, indep_var, estimate, p.value)
  return(res)
}
indep_var <- 'first_cosine_sim_observed_z'
#indep_var <- 'first_dream_sim_observed_z'
res <- list_rbind(map(user_dep_vars, \(x) indep_var_reg(x, reg_df=tmp))) %>%
  arrange(p.value)
cat('Coefs of regressing dep vars on', indep_var)
res %>% filter(p.value < 0.05)

# regression within pos and neg observed z
neg_data <- tmp %>% filter(.data[[indep_var]] <= 0)
pos_data <- tmp %>% filter(.data[[indep_var]] > 0)
neg_res <- list_rbind(map(user_dep_vars, \(x) indep_var_reg(x, reg_df=neg_data))) %>%
  arrange(p.value) %>%
  mutate(group = "negative")
pos_res <- list_rbind(map(user_dep_vars, \(x) indep_var_reg(x, reg_df=pos_data))) %>%
  arrange(p.value) %>%
  mutate(group = "positive")
res <- bind_rows(neg_res, pos_res)
cat('Coefs of regressing dep vars on positive/negative ranges of', indep_var)
res %>% filter(p.value < 0.05)
