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 <- '~/research/prompteng_analysis/data/finalized_datasets_for_analysis/main_df_promptLvl_2024-07-29_notWinsorized_PROD.csv'

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

result <- generate_data_from_user_iter_file(user_iter_file)
user_iter_df <- result$user_iter_df
user_df <- result$user_df

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)

# remove any row with NA
user_df_comp <- user_df[complete.cases(user_df[,group_vars]),]

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')
dep_var <- dep_vars[6]
cat('Dependent variable is:', dep_var)


##########################
# Method 1: ANOVA
p_values_twoway <- user_df_comp %>% summarise(across(all_of(group_vars), ~ summary(aov(user_df_comp[[dep_var]] ~ .x + user_df_comp$dalle_version))[[1]]$`Pr(>F)`[1]))
effects_twoway <- user_df_comp %>% summarise(across(all_of(group_vars), ~ summary(aov(user_df_comp[[dep_var]] ~ .x + user_df_comp$dalle_version))[[1]]$`F value`[1]))
p_values_threeway <- user_df_comp %>% summarise(across(all_of(group_vars), ~ summary(aov(user_df_comp[[dep_var]] ~ .x + 
                                                                                           user_df_comp$dalle_version + user_df_comp$referenceImageName))[[1]]$`Pr(>F)`[1]))
effects_threeway <- user_df_comp %>% summarise(across(all_of(group_vars), ~ summary(aov(user_df_comp[[dep_var]] ~ .x + 
                                                                                          user_df_comp$dalle_version + user_df_comp$referenceImageName))[[1]]$`F value`[1]))

p_values_twoway <- p_values_twoway %>% pivot_longer(everything(), names_to="Variable", values_to="p_value_twoway")
effects_twoway <- effects_twoway %>% pivot_longer(everything(), names_to="Variable", values_to="effects_twoway")
p_values_threeway <- p_values_threeway %>% pivot_longer(everything(), names_to="Variable", values_to="p_value_threeway")
effects_threeway <- effects_threeway %>% pivot_longer(everything(), names_to="Variable", values_to="effects_threeway")

p_values <- inner_join(p_values_twoway, p_values_threeway, by="Variable")
p_values$p_value_twoway_adjusted <- p.adjust(p_values$p_value_twoway, method = "BH")
p_values$p_value_threeway_adjusted <- p.adjust(p_values$p_value_threeway, method = "BH")

effects <- inner_join(effects_twoway, effects_threeway, by="Variable")
effects_and_p_values <- inner_join(p_values, effects, by="Variable")

effects_and_p_values %>% filter(p_value_threeway_adjusted <= 0.05) %>%
  arrange(p_value_threeway_adjusted) %>%
  select(Variable, p_value_threeway, p_value_threeway_adjusted, effects_threeway)

##########################
# Method 2: Kruskal-Wallis
p_values2 <- user_df_comp %>% summarise(across(all_of(group_vars), ~ (kruskal.test(user_df_comp[[dep_var]] ~ .x))$p.value))
p_values2 <- p_values2 %>% pivot_longer(everything(), names_to="Variable", values_to="p_value_kw")
p_values2$p_value_kw_adjusted <- p.adjust(p_values2$p_value_kw, method = "BH")

p_values2 %>% filter(p_value_kw_adjusted <= 0.05) %>%
  arrange(p_value_kw_adjusted)

##########################
# Method 3: Linear Regression
factor_p_vals <- 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]]', paste0(x, '_'), rownames(res), fixed=TRUE)
  rnames <- rownames(res)
  res <- res[grepl(x, rnames, fixed=TRUE),'Pr(>|t|)']
  res <- setNames(res, rnames[grepl(x, rnames, fixed=TRUE)])
  return(res)
}

factor_estimates<- 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]]', paste0(x, '_'), rownames(res), fixed=TRUE)
  rnames <- rownames(res)
  res <- res[grepl(x, rnames, fixed=TRUE),'Estimate']
  res <- setNames(res, rnames[grepl(x, rnames, fixed=TRUE)])
  return(res)
}

factor_ses <- function(x) {
  res <- lm(user_df_comp[[dep_var]] ~ user_df_comp[[x]] + user_df_comp$dalle_version)
  print(res)
  res <- coeftest(res, vcov=vcovHC(res, type="HC1"))
  print(res)
  rownames(res) <- sub('user_df_comp[[x]]', paste0(x, '_'), rownames(res), fixed=TRUE)
  rnames <- rownames(res)
  res <- res[grepl(x, rnames, fixed=TRUE),'Std. Error']
  res <- setNames(res, rnames[grepl(x, rnames, fixed=TRUE)])
  return(res)
}

p_values3 <- c()
estimates3 <- c()
ses3 <- c()

invisible(capture.output({
  for(group_var in group_vars) {
    res <- factor_p_vals(group_var)
    res2 <- factor_estimates(group_var)
    res3 <- factor_ses(group_var)
    p_values3 <- c(p_values3, res)
    estimates3 <- c(estimates3, res2)
    ses3 <- c(ses3, res3)
  }
}))

p_values3 <- tibble(Variable=names(p_values3), p_value_lm=p_values3)
estimates3 <- tibble(Variable=names(estimates3), estimates_lm=estimates3)
ses3 <- tibble(Variable=names(ses3), ses_lm=ses3)
p_values3$p_value_lm_adjusted <- p.adjust(p_values3$p_value_lm, method = "BH")
p_values_and_estimates3 <- inner_join(p_values3, estimates3, by="Variable")

p_values_and_estimates3 %>% filter(p_value_lm_adjusted <= 0.05) %>%
  arrange(estimates_lm)



