pdf(NULL) # suppress Rplots.pdf in batch mode
library(ggplot2)
library(lmtest)
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'
result <- generate_data_from_user_iter_file(user_iter_file)
user_iter_df <- result$user_iter_df
user_df <- result$user_df

# user-level nlp vars
topical_trans_col <- "n_total_topics_tsr_tasklvl_mean"
nlp_cols_cont <- c("embedding_var", "mean_tsr_prev_prompt", "mean_cosine_sim_prev_prompt",
                   "mean_prompt_length_chars", "mean_prompt_length_words")
nlp_cols_disc <- c(topical_trans_col, "mean_contains_prev_prompt")
nlp_cols <- c(nlp_cols_disc, nlp_cols_cont)

# user-iter-level nlp vars
iter_nlp_cols_cont <- c("tsr_prev_prompt", "cosine_sim_prev_prompt")
iter_nlp_cols_disc <- c("flag_topical_transition_tsr_tasklvl", "contains_prev_prompt")
iter_nlp_cols <- c(iter_nlp_cols_cont, iter_nlp_cols_disc)

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')
idx <- 1
user_dep_var <- dep_vars[idx]
iter_dep_var <- iter_dep_vars[idx]
cat('Dependent variables are:', user_dep_var, 'and', iter_dep_var)


# collapse rare and large transitions to a single category
cap_topical_trans <- 4
user_df_comp <- user_df %>%
  mutate(across(all_of(topical_trans_col), ~ifelse(. > cap_topical_trans, cap_topical_trans, .))) %>%
  filter(if_any(all_of(nlp_cols), ~ !is.na(.))) %>%
  ungroup()

user_iter_df_comp <- user_iter_df %>%
  filter(iteration > 0 & iteration < 10) %>%
  filter(if_any(all_of(iter_nlp_cols), ~ !is.na(.))) %>%
  ungroup()


##############################
# Method 1: Linear Model for User-level
indep_var_res <- function(x) {
  res <- lm(user_df_comp[[user_dep_var]] ~ user_df_comp[[x]] + user_df_comp$dalle_version)
  res <- coeftest(res, vcov=vcovHC(res, type="HC1"))
  rownames(res) <- sub('user_df_comp[[x]]', x, rownames(res), fixed=TRUE)
  res <- as.data.frame(res[,])
  rnames <- rownames(res)
  res <- res[grepl(x, rnames, fixed=TRUE),c('Estimate', 'Pr(>|t|)')]
  res <- res %>% rename(p_value=`Pr(>|t|)` , estimate=Estimate)
  return(res)
}
res <- list_rbind(map(nlp_cols, indep_var_res))
res$p_value_adjusted <- p.adjust(res$p_value, method = "BH")
res %>% arrange(p_value_adjusted)



##############################
# Method 2: Matching for User-Iteration-level
treatment_var_idx <- 1
matching_vars <- c(paste0(iter_dep_var, "_prev"), "referenceImageName", "dalle_version", "iteration")
group_var <- paste0(iter_dep_var, '_prev_bracket')
# give large weight to categorical vars that get encoded as numbers so they get matched exactly
matching_weights <- c(1,1000, 1000, 1000)
treatment_vars <- c(paste0(iter_nlp_cols_cont, '_bracket'), iter_nlp_cols_disc)
cat('All possible treatment vars:\n', paste(treatment_vars, collapse='\n'))
treatment_var <- treatment_vars[treatment_var_idx]
cat(paste('Matching vars:', paste(matching_vars, collapse=', '), '\n',
          'Treatment var:', treatment_var, '\n',
          'Group var:', group_var, '\n',
          'Dependent var:', iter_dep_var))

num_nlp_brackets <- 7
num_prev_iter_brackets <- 6
tmp_df <- user_iter_df_comp %>%
  mutate(across(all_of(iter_nlp_cols_cont),
                ~factor(cut(.x, breaks=quantile(.x, probs=seq(0, 1, length.out=num_nlp_brackets), na.rm=TRUE), include.lowest = TRUE), ordered = TRUE),
                .names='{col}_bracket'),
         across(all_of(paste0(iter_dep_vars, '_prev')),
                ~factor(cut(.x, breaks=quantile(.x, probs=seq(0, 1, length.out=num_prev_iter_brackets), na.rm=TRUE), include.lowest = TRUE), ordered = TRUE),
                .names='{col}_bracket')) %>%
  drop_na(all_of(c(matching_vars, treatment_vars)))

groups <- sort(unique(as.factor(tmp_df[[group_var]])))
res_list <- list()
# which treatment is base case?
base_case_id <- 1
for(group in groups) {
  tmp_group_df <- tmp_df[(tmp_df[[group_var]]==group),]
  tmp_group_df$iteration <- as.factor(tmp_group_df$iteration)
  tmp_group_df <- droplevels(tmp_group_df)
  if(nrow(tmp_group_df) < 400) {
    next
  }
  distances <- distances(data.matrix(tmp_group_df[, matching_vars]), weights=matching_weights)#, normalize="studentize") 
  treatments <- unique(as.factor(tmp_group_df[[treatment_var]]))
  # ensure there are enough units from each treatment in a match group 
  matching_ate <- quickmatch(distances, tmp_group_df[[treatment_var]], treatment_constraints=setNames(rep(5, length(treatments)), treatments))
  matchfit <- lm_match(tmp_group_df[[iter_dep_var]], tmp_group_df[[treatment_var]], matching_ate)
    
  treatment_levels <- rownames(matchfit$effects)
  plot_df <- data.frame(group=group,
                        treatment_bracket=factor(treatment_levels, levels=treatment_levels, ordered = TRUE),
                        ate=matchfit$effects[,base_case_id],
                        se=sqrt(matchfit$effect_variances[,base_case_id]))
  res_list[[length(res_list) + 1]] <- plot_df
}
plot_df <- do.call("rbind", res_list)
plot_df$group <- factor(plot_df$group, levels=groups, ordered = TRUE)
p <- ggplot(plot_df, aes(x=treatment_bracket, y=ate, color=group)) +
  geom_point(position = position_dodge(0.2)) +
  geom_pointrange(aes(ymin=ate - 2*se, ymax=ate + 2*se),
                position=position_dodge(0.2)) +
  xlab(treatment_var) +
  scale_color_manual(values = scales::hue_pal()(5)) +
  guides(color=guide_legend(title='Previous prompt z-score')) +
  theme_minimal() + # Increase base text size by 50%
  labs(y=paste('ATE of', iter_dep_var)) +
  theme(
    text = element_text(size = 11),
    legend.text = element_text(size = 9),
    legend.position='bottom'
  )
p

# generate a table of data used for the plot: ATE wrt base case, SE and P-value
plot_df$z <- plot_df$ate / plot_df$se
plot_df$p_value <- 2*(1-pnorm(abs(plot_df$z)))
# exclude the base case
table_data <- plot_df %>% filter(treatment_bracket != levels(plot_df$treatment_bracket)[base_case_id])
table_data <- table_data %>%
  arrange(group, group) %>%
  select(group, treatment_bracket, ate, se, p_value) %>%
  rename(
    `Previous prompt performance bracket` = group,
    `Cosine similarity bracket` = treatment_bracket,
    `ATE` = ate,
    `SE` = se,
    `p` = p_value
  )
table_data
