pdf(NULL) # suppress Rplots.pdf in batch mode
library(ggplot2)
library(lmtest)
library(sandwich)
library(tidyverse)
library(here)
# source("~/research/prompteng_analysis/replication/utils.R")

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")

# 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)

skill_vars <- c("Reading Comprehension", "Active Listening", "Writing", "Speaking", "Critical Thinking",                       
                "Social Perceptiveness", "Coordination", "Instructing", "Programming",
                "Judgment and Decision Making", "Systems Evaluation", "Science", "Active Learning",
                "Learning Strategies", "Monitoring", "Complex Problem Solving", "Operations Analysis",
                "Technology Design", "Troubleshooting", "Quality Control Analysis", "Systems Analysis")
group_vars <- c('GPT_usage4', 'degree3', 'gender', 'age5', 'GenAI_outlook3',
                'programming_skill3', 'program_frequency3', 'imagery_skill3', 'write_image_use3',
                'instruction_skill3', 'write_instr_usage3')
group_vars <- c(group_vars, skill_vars)



# 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_df_comp <- user_df_comp[complete.cases(user_df_comp[,c(group_vars, nlp_cols)]),]


##############################
# Method 1: Linear Regression
dep_var <- nlp_cols[1]
indep_var_res <- function(x) {
  res <- lm(user_df_comp[[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(group_vars, indep_var_res))
res$p_value_adjusted <- p.adjust(res$p_value, method = "BH")
res <- res %>% arrange(p_value_adjusted)

sig_vars <- res %>% filter(p_value_adjusted < 0.05)
cat('Significant variables for the dependent variable:', dep_var)
sig_vars

# plot the indep variable against nlp variables
group_var <- "age5"
#group_var <- "program_frequency3"
#group_var <- "write_image_use3"
#group_var <- "write_instr_usage3"
#group_var <- "GenAI_outlook3"
#group_var <- "GPT_usage4"
#group_var <- "gender"
#group_var <- "degree"
summary_df <- user_df_comp %>% group_by(across(all_of(group_var))) %>%
  summarise(across(all_of(dep_var), list(mean=mean, se=sd, n=~n()), .names='{.fn}_{.col}'))
summary_df[[paste0('se_',dep_var)]] <- summary_df[[paste0('se_',dep_var)]] / sqrt(summary_df[[paste0('n_',dep_var)]])
summary_df

mean_dep_col <- paste0('mean_', dep_var)
se_dep_col <- paste0('se_', dep_var)
ggplot(summary_df, aes(x=.data[[group_var]], y=.data[[mean_dep_col]])) +
  geom_point(position = position_dodge(width = 0.4)) +
  geom_errorbar(aes(ymin=.data[[mean_dep_col]] - 2*.data[[se_dep_col]],
                    ymax=.data[[mean_dep_col]] + 2*.data[[se_dep_col]]),
                width=0.2, position = position_dodge(width = 0.4)) +
  theme_bw() +
  theme(axis.text=element_text(colour="black", size=10),
        axis.title=element_text(colour="black", size=10),
        legend.title=element_text(size=7),
        legend.text=element_text(size=8))