pdf(NULL) # suppress Rplots.pdf in batch mode
# ============================================================================
# PRE-REGISTERED HYPOTHESIS 5: PERFORMANCE ACROSS MODELS
# Pairwise t-tests, ANOVA, and contrasts for performance and NLP metrics
# ============================================================================

rm(list = ls())
set.seed(1324)

# === LIBRARIES ===
library(tidyverse)
library(here)
library(fixest)
library(broom)
library(gtsummary)
library(report)
library(modelbased)
library(tidylog, warn.conflicts = FALSE)

source(here("shared", "utils.R"))

# === DATA IMPORT ===
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

# === DATA PROCESSING ===
main_df <- user_iter_df %>%
  rename(image_id = originalID) %>%
  group_by(referenceImageName) %>%
  mutate(dream_sim_zscore_taskwide        = scale(dream_sim_mean, center = TRUE, scale = TRUE)[,1],
         cosine_sim_zscore_taskwide       = scale(cosine_sim_mean, center = TRUE, scale = TRUE)[,1],
         dream_sim_zscore_taskwide_noisy  = scale(dream_sim_observed, center = TRUE, scale = TRUE)[,1],
         cosine_sim_zscore_taskwide_noisy = scale(cosine_sim_observed, center = TRUE, scale = TRUE)[,1]) %>%
  ungroup()

selected_vars <- c(
  "referenceImageName", "userId", "iteration", "cosine_sim_mean", "cosine_sim_sd", "cosine_sim_observed",
  "cosine_sim_observed_z", "dream_sim_mean", "dream_sim_sd", "dream_sim_observed", "dream_sim_observed_z",
  "image_id", "sampled_cnt", "dalle_version", "unique_originalIDs", "unique_versions", "userIdreferenceImage",
  "referenceImageCategory", "originalOrder", "isDuplicate", "prompt", "revisedPrompt",
  "cosine_dist_from_center", "cosine_dist_from_center_quantile", "euclid_dist_from_center", "euclid_dist_from_center_quantile",
  "junk_flag", "unfiltered_prompt",
  "shifted_prompt", "shifted_junk_flag", "prev_prompt",
  "shifted_prompt_embedding_vector", "prev_prompt_embedding_vector", "tsr_prev_prompt", "contains_prev_prompt",
  "prompt_length_chars", "prompt_length_words", "dist_to_user_centroid_junkFiltered", "dist_to_task_centroid_junkFiltered",
  "cosine_sim_prev_prompt", "taskWide_prevCosineSimPctile", "originalReplication_sample_mean_dreamsim",
  "originalReplication_sample_var_dreamsim", "originalReplication_sample_count", "og_dreamsim_score", "tsr_revised_original",
  "dream_sim_zscore_taskwide", "cosine_sim_zscore_taskwide",
  "dream_sim_zscore_taskwide_noisy", "cosine_sim_zscore_taskwide_noisy"
)

main_df <- main_df %>%
  select(all_of(selected_vars)) %>%
  mutate(referenceImageName = factor(referenceImageName),
         userId = factor(userId))

pdat <- panel(main_df, panel.id = c('userId', 'iteration'))

# Initial user-level df (cosine sim)
userLvl_df <- pdat %>%
  arrange(userId, iteration) %>%
  group_by(userId, dalle_version, referenceImageName) %>%
  summarize(
    first_Score = first(cosine_sim_mean),
    best_Score  = max(cosine_sim_mean),
    mean_Score  = mean(cosine_sim_mean),
    .groups = "drop"
  ) %>%
  group_by(referenceImageName) %>%
  mutate(first_zscore_taskwide = scale(first_Score, center = TRUE, scale = TRUE)[,1],
         best_zscore_taskwide  = scale(best_Score,  center = TRUE, scale = TRUE)[,1],
         mean_zscore_taskwide  = scale(mean_Score,  center = TRUE, scale = TRUE)[,1],
         first_pct_rank = round(cume_dist(first_Score) * 100),
         best_pct_rank  = round(cume_dist(best_Score)  * 100),
         mean_pct_rank  = round(cume_dist(mean_Score)  * 100)) %>%
  ungroup()

# === HELPER FUNCTIONS ===
get_pairwise.t.test <- function(data, variable, by,
                                p.adjust.method = "BH", paired = FALSE, pool.sd = TRUE, ...) {
  pairwise_results <- pairwise.t.test(data[[variable]], data[[by]],
                                      p.adjust.method = p.adjust.method,
                                      paired = paired, pool.sd = pool.sd)
  broom::tidy(pairwise_results) %>%
    mutate(label = glue::glue("({group1} vs. {group2})")) %>%
    select(label, p.value) %>%
    spread(label, p.value)
}

style_pvalue_stars <- function(x) {
  dplyr::case_when(
    x < 0.001 ~ paste0(style_pvalue(x, digits = 3, prepend_p = TRUE), "***"),
    x < 0.01  ~ paste0(style_pvalue(x, digits = 3, prepend_p = TRUE), "**"),
    x < 0.05  ~ paste0(style_pvalue(x, digits = 3, prepend_p = TRUE), "*"),
    TRUE       ~ style_pvalue(x, digits = 3, prepend_p = TRUE)
  )
}

# === PERFORMANCE SUMMARY TABLE ===
get_pairwise.t.test(userLvl_df, "first_zscore_taskwide", "dalle_version", pool.sd = TRUE)
pairwise_results_temp <- get_pairwise.t.test(userLvl_df, "first_zscore_taskwide", "dalle_version")
colnames_list <- colnames(pairwise_results_temp)

table <- tbl_summary(
    userLvl_df,
    include = c(first_zscore_taskwide, best_zscore_taskwide, mean_zscore_taskwide),
    by = dalle_version,
    missing = "no",
    statistic = everything() ~ "{mean} ({sd})"
  ) %>%
  add_n() %>%
  add_stat(fns = everything() ~ get_pairwise.t.test) %>%
  modify_header(label = "**Variable**") %>%
  modify_fmt_fun(all_of(colnames_list) ~ style_pvalue_stars) %>%
  bold_labels()

table

# === CONTRAST ANALYSIS HELPER ===
run_contrast_loop <- function(userLvl_df, variables, metric_label) {
  contr_list        <- list()
  report_list       <- list()
  report_table_list <- list()
  aov_report_list   <- list()

  for (var in variables) {
    formula <- as.formula(paste(var, "~ dalle_version"))
    model   <- lm(formula, data = userLvl_df)

    # TODO: p_adjust uses "fdr" (lowercase alias for Benjamini-Hochberg, identical to "BH")
    # because modelbased 0.13.1 has a tolower() bug that breaks "BH" in estimate_contrasts()
    contr              <- estimate_contrasts(model, p_adjust = "fdr")
    contr_report       <- suppressWarnings(report(contr))
    contr_report_table <- report_table(contr)

    contr_list[[var]]        <- contr
    report_list[[var]]       <- contr_report
    report_table_list[[var]] <- contr_report_table

    aov_model <- aov(formula, data = userLvl_df)
    aov_report <- report(aov_model)
    aov_report_list[[var]] <- aov_report

    cat("\n\n", strrep("-", 50), "\n")
    cat("ANOVA Results for variable:", var, "and metric:", metric_label, "\n")
    cat(strrep("-", 50), "\n\n")
    print(aov_report)

    cat("\n\n", strrep("-", 50), "\n")
    cat("Results for variable:", var, "and metric:", metric_label, "\n")
    cat(strrep("-", 50), "\n\n")
    print(contr_report)
    print(contr_report_table)
  }

  list(contr_list = contr_list, report_list = report_list,
       report_table_list = report_table_list, aov_report_list = aov_report_list)
}

# ============================================================================
# COSINE SIM REPORTING
# ============================================================================
metric <- 'cosine'

userLvl_df <- pdat %>%
  arrange(userId, iteration) %>%
  group_by(userId, dalle_version, referenceImageName) %>%
  summarize(
    first_Score = first(cosine_sim_mean),
    best_Score  = max(cosine_sim_mean),
    mean_Score  = mean(cosine_sim_mean),
    .groups = "drop"
  ) %>%
  group_by(referenceImageName) %>%
  mutate(first_zscore_taskwide = scale(first_Score, center = TRUE, scale = TRUE)[,1],
         best_zscore_taskwide  = scale(best_Score,  center = TRUE, scale = TRUE)[,1],
         mean_zscore_taskwide  = scale(mean_Score,  center = TRUE, scale = TRUE)[,1],
         first_pct_rank = round(cume_dist(first_Score) * 100),
         best_pct_rank  = round(cume_dist(best_Score)  * 100),
         mean_pct_rank  = round(cume_dist(mean_Score)  * 100)) %>%
  ungroup()

perf_variables <- c("first_zscore_taskwide", "best_zscore_taskwide", "mean_zscore_taskwide",
                     "first_pct_rank", "best_pct_rank", "mean_pct_rank")

cosine_contrasts <- run_contrast_loop(userLvl_df, perf_variables, metric)
contr_list        <- cosine_contrasts$contr_list
report_list       <- cosine_contrasts$report_list
report_table_list <- cosine_contrasts$report_table_list
aov_report_list   <- cosine_contrasts$aov_report_list

# ============================================================================
# DREAMSIM REPORTING
# ============================================================================
metric <- 'dreamsim'

userLvl_df <- pdat %>%
  arrange(userId, iteration) %>%
  group_by(userId, dalle_version, referenceImageName) %>%
  summarize(
    first_Score = first(dream_sim_mean),
    best_Score  = max(dream_sim_mean),
    mean_Score  = mean(dream_sim_mean),
    .groups = "drop"
  ) %>%
  group_by(referenceImageName) %>%
  mutate(first_zscore_taskwide = scale(first_Score, center = TRUE, scale = TRUE)[,1],
         best_zscore_taskwide  = scale(best_Score,  center = TRUE, scale = TRUE)[,1],
         mean_zscore_taskwide  = scale(mean_Score,  center = TRUE, scale = TRUE)[,1],
         first_pct_rank = round(cume_dist(first_Score) * 100),
         best_pct_rank  = round(cume_dist(best_Score)  * 100),
         mean_pct_rank  = round(cume_dist(mean_Score)  * 100)) %>%
  ungroup()

dreamsim_contrasts <- run_contrast_loop(userLvl_df, perf_variables, metric)
contr_list        <- dreamsim_contrasts$contr_list
report_list       <- dreamsim_contrasts$report_list
report_table_list <- dreamsim_contrasts$report_table_list
aov_report_list   <- dreamsim_contrasts$aov_report_list

# ============================================================================
# NLP BEHAVIOR METRICS
# ============================================================================
userLvl_df <- pdat %>%
  arrange(userId, iteration) %>%
  group_by(userId, dalle_version, referenceImageName) %>%
  summarize(
    userLvl_mean_prompt_length         = mean(prompt_length_words, na.rm = TRUE),
    userLvl_mean_dist_to_user_centroid = mean(dist_to_user_centroid_junkFiltered, na.rm = TRUE),
    userLvl_mean_cosine_sim_prev_prompt = mean(cosine_sim_prev_prompt, na.rm = TRUE),
    userLvl_mean_tsr_prev_prompt       = mean(tsr_prev_prompt, na.rm = TRUE),
    userLvl_mean_contains_prev_prompt  = mean(contains_prev_prompt, na.rm = TRUE),
    .groups = "drop"
  ) %>%
  group_by(referenceImageName) %>%
  mutate(prompt_length_zscore_taskwide           = scale(userLvl_mean_prompt_length, center = TRUE, scale = TRUE)[,1],
         dist_to_user_centroid_zscore_taskwide   = scale(userLvl_mean_dist_to_user_centroid, center = TRUE, scale = TRUE)[,1],
         cosine_sim_prev_zscore_taskwide         = scale(userLvl_mean_cosine_sim_prev_prompt, center = TRUE, scale = TRUE)[,1],
         tsr_prev_zscore_taskwide                = scale(userLvl_mean_tsr_prev_prompt, center = TRUE, scale = TRUE)[,1],
         contains_prev_prompt_zscore_taskwide    = scale(userLvl_mean_contains_prev_prompt, center = TRUE, scale = TRUE)[,1]) %>%
  ungroup()

# Behavior summary table
get_pairwise.t.test(userLvl_df, "prompt_length_zscore_taskwide", "dalle_version")
pairwise_results_temp <- get_pairwise.t.test(userLvl_df, "prompt_length_zscore_taskwide", "dalle_version")
colnames_list <- colnames(pairwise_results_temp)

variables <- c("prompt_length_zscore_taskwide", "dist_to_user_centroid_zscore_taskwide",
               "cosine_sim_prev_zscore_taskwide", "tsr_prev_zscore_taskwide",
               "contains_prev_prompt_zscore_taskwide")

table <- tbl_summary(
    userLvl_df,
    include = all_of(variables),
    by = dalle_version,
    missing = "no",
    statistic = everything() ~ "{mean} ({sd})"
  ) %>%
  add_n() %>%
  add_stat(fns = everything() ~ get_pairwise.t.test) %>%
  modify_header(label = "**Variable**") %>%
  modify_fmt_fun(all_of(colnames_list) ~ style_pvalue_stars) %>%
  bold_labels()

table

# Behavior contrasts
behavior_contrasts <- run_contrast_loop(userLvl_df, variables, "behavior")
contr_list        <- behavior_contrasts$contr_list
report_list       <- behavior_contrasts$report_list
report_table_list <- behavior_contrasts$report_table_list
aov_report_list   <- behavior_contrasts$aov_report_list
