pdf(NULL) # suppress Rplots.pdf in batch mode
# ============================================================================
# BALANCE TABLE: MAIN EXPERIMENT
# Covariate balance checks across treatment groups (DALLE version)
# ============================================================================

# ============================================================================
# SECTION 1: SETUP AND CONFIGURATION
# ============================================================================

# Clear environment for reproducible execution
rm(list = ls())

# Log execution timestamp
cat("=== BALANCE TABLE: MAIN EXPERIMENT - START ===",
    format(Sys.time()), "\n")

# Set random seed for reproducibility
set.seed(1324)

# ============================================================================
# SECTION 2: PACKAGE LOADING
# ============================================================================
suppressPackageStartupMessages({
  library(dplyr)          # Data manipulation
  library(readr)          # CSV reading/writing
  library(here)           # Path management
  library(gtsummary)      # Summary tables with statistical tests
})

# ============================================================================
# SECTION 4: PATH VALIDATION
# ============================================================================

# Verify script location and project structure
here::i_am("similarity_exp_replication/code/analysis/prod_balance_tbl_main_exp.R")

# Verify critical directories
cat("Working directory:", getwd(), "\n")

stopifnot(
  "inputs/ directory not found" = dir.exists(here("similarity_exp_replication", "inputs")),
  "outputs/tables/ directory needed" = {
    if (!dir.exists(here("similarity_exp_replication", "outputs", "tables"))) dir.create(here("similarity_exp_replication", "outputs", "tables"), recursive = TRUE)
    TRUE
  }
)

# ============================================================================
# SECTION 5: DATA IMPORT & VARIABLE DEFINITIONS
# ============================================================================

cat("\nLoading main experiment data...\n")

# Load and process main experiment data
# user_iter_file <- here("similarity_exp_replication", "inputs",
#                        "main_df_promptLvl_2025-02-12_PROD_HASHED.csv")

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

user_iter_df = read_csv(user_iter_file) %>%
  mutate(
    GPT_usage4 = factor(ifelse(is.na(GPT_usage4), "NA/Missing", GPT_usage4),
                        levels = c("None", "Little", "Some", "A lot", "NA/Missing")),
    gender = factor(ifelse(is.na(gender), "NA/Missing", gender),
                    levels = c("Female", "Male", "Other or prefer not to say", "NA/Missing")),
    age5 = factor(ifelse(is.na(age5), "NA/Missing", age5),
                  levels = c("[18,26)", "[26,32)", "[32,39)", "[39,49)", "[49,83]", "NA/Missing"))
  )

# Define skill-related covariates
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"
)

# Define demographic and usage covariates
group_vars <- c(
  "GPT_usage4", "degree3", "gender", "age5", "GenAI_outlook3",
  "programming_skill3", "program_frequency3", "imagery_skill3",
  "write_image_use3", "instruction_skill3", "write_instr_usage3"
)

# Combine all covariates
group_vars <- c(group_vars, skill_vars)
id_vars <- c("userId", "referenceImageName", "dalle_version")

# Extract user-level data from iteration-level data
# Take first occurrence of each (userId, referenceImageName) combination
# This preserves user-level covariates which don't vary across iterations
user_df <- user_iter_df %>%
  distinct(userId, referenceImageName, dalle_version, .keep_all = TRUE) %>%
  select(all_of(c(id_vars, group_vars)))

cat("Loaded", n_distinct(user_df$userId), "unique users\n")
cat("Loaded", nrow(user_df), "user-image combinations\n")
cat("Loaded", nrow(user_iter_df), "total iterations\n\n")


# ============================================================================
# SECTION 6: COVARIATE BALANCE ASSESSMENT
# ============================================================================

# Define balance assessment covariates
balance_covars <- c("GPT_usage4", "gender", "age5")
# GenAI_outlook3 temporarily commented out
# balance_covars <- c("GPT_usage4", "gender", "age5", "GenAI_outlook3")

cat("🔄 Assessing covariate balance across treatment arms...\n\n")

# Create arm mapping for display (DALLE versions -> labels)
arm_label_map <- c(
  "2" = "D2",
  "3-revised-prompt" = "D3r",
  "3-verbatim-prompt" = "D3v"
)

# Subset data to balance covariates only (we need to make sure the user count matches 1893 as reported in paper
# dont need to match prompt count bc this is not a prompt analysis)
user_df_balance <- user_df[, c("dalle_version", balance_covars)]

cat("User counts by arm:\n\n")
table(user_df_balance$dalle_version)

cat("Creating pairwise balance tables with gtsummary...\n\n")

# D2 vs D3r comparison
balance_d2_d3r <- user_df_balance %>%
  filter(dalle_version %in% c("2", "3-revised-prompt")) %>%
  tbl_summary(
    by = dalle_version,
    missing = "no",
    label = list(
      GPT_usage4 ~ "GPT Usage",
      gender ~ "Gender",
      age5 ~ "Age Group"
      # GenAI_outlook3 ~ "GenAI Outlook"
    )
  ) %>%
  add_p(test = list(all_categorical() ~ "chisq.test"))

# D2 vs D3v comparison
balance_d2_d3v <- user_df_balance %>%
  filter(dalle_version %in% c("2", "3-verbatim-prompt")) %>%
  tbl_summary(
    by = dalle_version,
    missing = "no",
    label = list(
      GPT_usage4 ~ "GPT Usage",
      gender ~ "Gender",
      age5 ~ "Age Group"
      # GenAI_outlook3 ~ "GenAI Outlook"
    )
  ) %>%
  add_p(test = list(all_categorical() ~ "chisq.test"))

# D3r vs D3v comparison
balance_d3r_d3v <- user_df_balance %>%
  filter(dalle_version %in% c("3-revised-prompt", "3-verbatim-prompt")) %>%
  tbl_summary(
    by = dalle_version,
    missing = "no",
    label = list(
      GPT_usage4 ~ "GPT Usage",
      gender ~ "Gender",
      age5 ~ "Age Group"
      # GenAI_outlook3 ~ "GenAI Outlook"
    )
  ) %>%
  add_p(test = list(all_categorical() ~ "chisq.test"))

cat("✅ Gtsummary tables created\n\n")

# Helper function to extract balance table data from gtsummary object
# Extracts: proportions, chi-squared statistic, degrees of freedom, and p-value
# col1_name and col2_name are the group labels for proper naming
extract_balance_data <- function(tbl_object, comparison_name, col1_name, col2_name) {
  # Extract the table body which contains all test results
  tbl_data <- tbl_object$table_body %>%
    select(
      label,
      stat_1,
      stat_2,
      statistic,    # chi-squared statistic
      parameter,    # degrees of freedom
      p.value
    ) %>%
    # Rename stat columns with proper group labels
    rename(
      !!col1_name := stat_1,
      !!col2_name := stat_2
    ) %>%
    # Round statistic and p.value to 3 decimal places
    mutate(
      statistic = round(statistic, 3),
      p.value = round(p.value, 3)
    )

  # Add comparison name
  tbl_data$comparison <- comparison_name

  return(tbl_data)
}

# Extract data from each comparison with proper column naming
df_d2_d3r <- extract_balance_data(balance_d2_d3r, "D2_vs_D3r", "N(prop)_D2", "N(prop)_D3r")
df_d2_d3v <- extract_balance_data(balance_d2_d3v, "D2_vs_D3v", "N(prop)_D2", "N(prop)_D3v")
df_d3r_d3v <- extract_balance_data(balance_d3r_d3v, "D3r_vs_D3v", "N(prop)_D3r", "N(prop)_D3v")

cat("✅ Results extracted to DataFrames\n\n")

output_dir <- here("similarity_exp_replication", "outputs", "tables")

# Write CSV files
write_csv(df_d2_d3r, file.path(output_dir, "balance_table_main_exp_d2_vs_d3r.csv"))
cat(paste("✅ Saved D2 vs D3r comparison to:", file.path(output_dir, "balance_table_main_exp_d2_vs_d3r.csv"), "\n"))

write_csv(df_d2_d3v, file.path(output_dir, "balance_table_main_exp_d2_vs_d3v.csv"))
cat(paste("✅ Saved D2 vs D3v comparison to:", file.path(output_dir, "balance_table_main_exp_d2_vs_d3v.csv"), "\n"))

write_csv(df_d3r_d3v, file.path(output_dir, "balance_table_main_exp_d3r_vs_d3v.csv"))
cat(paste("✅ Saved D3r vs D3v comparison to:", file.path(output_dir, "balance_table_main_exp_d3r_vs_d3v.csv"), "\n"))

cat("\n=== BALANCE TABLE: MAIN EXPERIMENT - COMPLETE ===")
cat(format(Sys.time()), "\n")
