pdf(NULL) # suppress Rplots.pdf in batch mode
# ============================================================================
# PRE-REGISTERED HYPOTHESIS 7: IMPROVEMENT OVER ITERATIONS
# Blocked/stratified two-sample tests comparing first vs best scores
# ============================================================================

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

# === LIBRARIES ===
library(tidyverse)
library(here)
library(fixest)
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_pct_rank_taskwide  = round(cume_dist(dream_sim_mean) * 100),
         cosine_sim_pct_rank_taskwide = round(cume_dist(cosine_sim_mean) * 100)) %>%
  ungroup()

# Sanity check: compare z-score definitions
are_columns_equal <- all.equal(main_df$z_dream_sim_mean,
                               main_df$dream_sim_zscore_taskwide,
                               tolerance = 1e-6)
if (isTRUE(are_columns_equal)) {
  cat("The columns are the same within the specified tolerance.\n")
} else {
  cat("The columns are not the same within the specified tolerance.\n")
}

# === HELPER: T-TEST ===
run_t_test <- function(sub_df, measure) {
  t_test_result <- t.test(score ~ score_group, data = sub_df)
  means <- as.numeric(t_test_result$estimate)
  mean_best  <- means[1]  # alphabetical: "best" before "first"
  mean_first <- means[2]
  mean_diff  <- mean_best - mean_first

  data.frame(
    dalle_version = unique(sub_df$dalle_version),
    measure       = measure,
    mean_best     = mean_best,
    mean_first    = mean_first,
    estimate      = mean_diff,
    se            = t_test_result$stderr,
    n_first       = sum(sub_df$score_group == "first"),
    n_best        = sum(sub_df$score_group == "best"),
    p.value       = t_test_result$p.value
  )
}

# === MAIN ANALYSIS FUNCTION ===
run_similarity_analysis <- function(main_df, sim_metric = "cosine_sim") {
  metric_cols <- list(
    mean     = paste0(sim_metric, "_mean"),
    zscore   = paste0(sim_metric, "_zscore_taskwide"),
    pct_rank = paste0(sim_metric, "_pct_rank_taskwide")
  )

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

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

  blocked_test_df <- pdat %>%
    arrange(userId, iteration) %>%
    group_by(userId, dalle_version, referenceImageName) %>%
    summarize(
      first_Score   = first(!!sym(metric_cols$mean)),
      best_Score    = max(!!sym(metric_cols$mean)),
      first_Zscore  = first(!!sym(metric_cols$zscore)),
      best_Zscore   = max(!!sym(metric_cols$zscore)),
      first_pctRank = first(!!sym(metric_cols$pct_rank)),
      best_pctRank  = max(!!sym(metric_cols$pct_rank)),
      .groups = "drop"
    )

  t_test_df <- blocked_test_df %>%
    pivot_longer(
      cols = starts_with("first_") | starts_with("best_"),
      names_to = c("score_group", ".value"),
      names_pattern = "(first|best)_(.*)"
    ) %>%
    mutate(score_group = if_else(score_group == "first", "first", "best"))

  measures <- c("Score", "Zscore", "pctRank")
  results_list <- list()

  for (measure in measures) {
    results <- data.frame(
      dalle_version = character(), measure = character(),
      mean_best = numeric(), mean_first = numeric(),
      estimate = numeric(), se = numeric(),
      n_first = integer(), n_best = integer(),
      p.value = numeric(), stringsAsFactors = FALSE
    )

    for (version in unique(t_test_df$dalle_version)) {
      sub_df <- t_test_df %>%
        filter(dalle_version == version) %>%
        select(userId, dalle_version, score_group, score = !!sym(measure))

      if (length(unique(sub_df$score_group)) == 2) {
        result <- run_t_test(sub_df, measure)
        results <- rbind(results, result)
      }
    }

    results$adjusted_p.value <- p.adjust(results$p.value, method = "BH")
    results_list[[measure]] <- results
  }

  cat("\nResults for", sim_metric, "analysis:\n")
  for (measure in measures) {
    cat("\nResults for", measure, "| Estimate = Best/Max - First", "\n")
    print(results_list[[measure]])
    cat("\n")
  }

  list(results_list = results_list, t_test_df = t_test_df, blocked_test_df = blocked_test_df)
}

# === COSINE SIMILARITY ANALYSIS ===
output <- run_similarity_analysis(main_df, "cosine_sim")
cosine_results  <- output$results_list
cosine_t_test_df <- output$t_test_df
blocked_test_df <- output$blocked_test_df

# === DREAMSIM ANALYSIS ===
dream_results <- run_similarity_analysis(main_df, "dream_sim")

# ============================================================================
# STRATIFIED T-TESTS: INVERSE-VARIANCE WEIGHTS
# ============================================================================
run_stratified_t_tests <- function(data,
                                   strat_var = "dalle_version",
                                   measures = c("Score", "Zscore", "pctRank"),
                                   first_prefix = "first",
                                   best_prefix = "best") {
  results_list <- list()

  for (measure in measures) {
    first_col <- paste0(first_prefix, "_", measure)
    best_col  <- paste0(best_prefix, "_", measure)

    sub_df <- data %>%
      select(all_of(strat_var), all_of(first_col), all_of(best_col)) %>%
      pivot_longer(
        cols = c(all_of(first_col), all_of(best_col)),
        names_to = "score_group", values_to = "score"
      ) %>%
      mutate(score_group = if_else(str_detect(score_group, paste0("^", first_prefix)),
                                   "first", "best"))

    strata_summary <- sub_df %>%
      group_by(across(all_of(strat_var))) %>%
      summarise(
        n_first    = sum(score_group == "first"),
        n_best     = sum(score_group == "best"),
        mean_first = mean(score[score_group == "first"], na.rm = TRUE),
        mean_best  = mean(score[score_group == "best"], na.rm = TRUE),
        var_first  = var(score[score_group == "first"], na.rm = TRUE),
        var_best   = var(score[score_group == "best"], na.rm = TRUE),
        .groups = "drop"
      ) %>%
      mutate(
        delta     = mean_best - mean_first,
        var_delta = var_first / n_first + var_best / n_best,
        weight    = 1 / var_delta,
        df_strata = (var_delta^2) / (
          ((var_first / n_first)^2 / (n_first - 1)) +
          ((var_best  / n_best)^2  / (n_best  - 1))
        )
      )

    print(paste("Strata summary for measure:", measure))
    print(strata_summary)

    pooled_delta <- sum(strata_summary$delta * strata_summary$weight, na.rm = TRUE) /
                    sum(strata_summary$weight, na.rm = TRUE)
    pooled_se    <- sqrt(1 / sum(strata_summary$weight, na.rm = TRUE))
    t_stat       <- pooled_delta / pooled_se
    df_pooled    <- (sum(strata_summary$weight, na.rm = TRUE)^2) /
                    sum((strata_summary$weight^2) / strata_summary$df_strata, na.rm = TRUE)
    p_value      <- 2 * pt(-abs(t_stat), df = df_pooled)

    results_list[[measure]] <- data.frame(
      measure = measure, pooled_diff = pooled_delta, pooled_se = pooled_se,
      t_stat = t_stat, df = df_pooled, p_value = p_value,
      stringsAsFactors = FALSE
    )
  }

  results_list
}

results <- run_stratified_t_tests(blocked_test_df, strat_var = "dalle_version")

print("Overall, across strata t-tests for ....")
print("Score (group means: best - first):")
print(results$Score)
print("Zscore (group means: best - first):")
print(results$Zscore)
print("PctRank (group means: best - first):")
print(results$pctRank)

# ============================================================================
# SANITY CHECK: STRATIFIED T-TESTS WITH PROPORTIONAL WEIGHTS
# ============================================================================
run_stratified_t_tests <- function(data,
                                   strat_var = "dalle_version",
                                   measures = c("Score", "Zscore", "pctRank"),
                                   first_prefix = "first",
                                   best_prefix = "best") {
  results_list <- list()

  for (measure in measures) {
    first_col <- paste0(first_prefix, "_", measure)
    best_col  <- paste0(best_prefix, "_", measure)

    sub_df <- data %>%
      select(all_of(strat_var), all_of(first_col), all_of(best_col)) %>%
      pivot_longer(
        cols = c(all_of(first_col), all_of(best_col)),
        names_to = "score_group", values_to = "score"
      ) %>%
      mutate(score_group = if_else(str_detect(score_group, paste0("^", first_prefix)),
                                   "first", "best"))

    strata_summary <- sub_df %>%
      group_by(across(all_of(strat_var))) %>%
      summarise(
        n_first    = sum(score_group == "first"),
        n_best     = sum(score_group == "best"),
        mean_first = mean(score[score_group == "first"], na.rm = TRUE),
        mean_best  = mean(score[score_group == "best"], na.rm = TRUE),
        var_first  = var(score[score_group == "first"], na.rm = TRUE),
        var_best   = var(score[score_group == "best"], na.rm = TRUE),
        .groups = "drop"
      ) %>%
      mutate(
        n_total   = n_first + n_best,
        delta     = mean_best - mean_first,
        var_delta = var_first / n_first + var_best / n_best,
        df_strata = (var_delta^2) / (
          ((var_first / n_first)^2 / (n_first - 1)) +
          ((var_best  / n_best)^2  / (n_best  - 1))
        )
      ) %>%
      mutate(weight = n_total / sum(n_total))

    print(paste("Strata summary for measure:", measure))
    print(strata_summary)

    pooled_diff <- sum(strata_summary$delta * strata_summary$weight, na.rm = TRUE)
    pooled_var  <- sum((strata_summary$weight^2) * strata_summary$var_delta, na.rm = TRUE)
    pooled_se   <- sqrt(pooled_var)
    t_stat      <- pooled_diff / pooled_se
    df_pooled   <- 1 / sum((strata_summary$weight^2) / strata_summary$df_strata, na.rm = TRUE)
    p_value     <- 2 * pt(-abs(t_stat), df = df_pooled)

    results_list[[measure]] <- data.frame(
      measure = measure, pooled_diff = pooled_diff, pooled_se = pooled_se,
      t_stat = t_stat, df = df_pooled, p_value = p_value,
      stringsAsFactors = FALSE
    )
  }

  results_list
}

results <- run_stratified_t_tests(blocked_test_df, strat_var = "dalle_version")

print("Overall, across strata t-tests for ....")
print("Score (group means: best - first):")
print(results$Score)
print("Zscore (group means: best - first):")
print(results$Zscore)
print("PctRank (group means: best - first):")
print(results$pctRank)
