current_dir <- normalizePath(getwd(), winslash = "/", mustWork = TRUE)
project_root <- if (basename(current_dir) == "scripts") {
  dirname(current_dir)
} else if (dir.exists(file.path(current_dir, "scripts")) &&
           dir.exists(file.path(current_dir, "data"))) {
  current_dir
} else {
  dirname(current_dir)
}
setwd(project_root)
dir.create(file.path(project_root, "data"), showWarnings = FALSE)
dir.create(file.path(project_root, "outputs"), showWarnings = FALSE)
local_lib <- file.path(project_root, "r-lib")
if (dir.exists(local_lib)) {
  .libPaths(c(local_lib, .libPaths()))
}
local_lib <- file.path(getwd(), "r-lib")
.libPaths(c(local_lib, .libPaths()))

###### Voting for democracy - Average Treatment Moderation Effects (ATMEs)


library(dplyr)
library(tidyr)
library(tibble)
library(ggplot2)
library(sandwich)
library(lmtest)
library(here)

setwd(project_root)
dir.create(file.path(project_root, "data"), showWarnings = FALSE)
dir.create(file.path(project_root, "outputs"), showWarnings = FALSE)

# Read analytical datasets
us <- read.csv(file.path(project_root, "data", "us_analytic_v2.csv"))
sk <- read.csv(file.path(project_root, "data", "sk_analytic_v2.csv"))
pl <- read.csv(file.path(project_root, "data", "pl_analytic_v2.csv"))
pool <- read.csv(file.path(project_root, "data", "pool_analytic_v2.csv"))

# Read long conjoint datasets
us_long <- read.csv(file.path(project_root, "data", "us_long_v2.csv"))
sk_long <- read.csv(file.path(project_root, "data", "sk_long_v2.csv"))
pl_long <- read.csv(file.path(project_root, "data", "pl_long_v2.csv"))
pool_long <- read.csv(file.path(project_root, "data", "pool_long_v2.csv"))

# Ensure all eval vars exist in all four analytic datasets
eval_cols <- c("eval_dem_us","eval_rep_us",
               "eval_dpk_kr","eval_ppp_kr",
               "eval_pis_pl","eval_kon_pl",
               "eval_ko_pl","eval_lew_pl")

for (nm in c("pool","us","sk","pl")) {
  df <- get(nm)
  for (cl in eval_cols) {
    if (!cl %in% names(df)) {
      df[[cl]] <- NA_real_
    }
  }
  assign(nm, df)
}


#### Common setup ----


pc_vars <- c("demaut_pref","sup_ldem","inst_trust","pol_tolerance",
             "autonomy_values","emancip_val")

covars  <- c("gender_male","age_1","educ_5f","income_cat","region_pool",
             "pol_interest","gend_pol_left","redist_pol_left","top_party")

spec_levels <- c("main cntrls","all cntrls")

pro_dem_labels <- c("keep polling stations","respect court rulings",
                    "don't prosecute journalists")

# Order and labels for PC variables
pc_order <- c("demaut_pref","sup_ldem","inst_trust","pol_tolerance",
              "emancip_val","autonomy_values")

label_map <- c(
  demaut_pref = "Democ-autoc.\npreference",
  sup_ldem = "Support liberal\ndemocracy",
  inst_trust = "Institutional\ntrust",
  pol_tolerance = "Political\ntolerance",
  emancip_val = "Emancipative\nvalues",
  autonomy_values = "Authoritarian\nvalues (rev)"
)

# For cross-ideology plots
offsets  <- c("main cntrls" = +0.12, "all cntrls" = -0.12)

# For combined ATME plots
atme_pal <- c("Overall" = "#000000", "Cross-party" = "goldenrod",
              "Cross-ideology (redist)" = "dodgerblue3",
              "Cross-ideology (gend)"   = "dodgerblue3")

atme_shapes <- c("Overall" = 21, "Cross-party" = 22,
                 "Cross-ideology (redist)" = 24,
                 "Cross-ideology (gend)"   = 25)

series_offsets <- c("Overall" = +0.24, "Cross-party" = +0.08,
                    "Cross-ideology (redist)" = -0.08,
                    "Cross-ideology (gend)"   = -0.24)

dataset_names <- c("pool","us","sk","pl")



#### Loop 1: construct datasets----


for (nm in dataset_names) {
  
  message("Building datasets for: ", nm)
  
  base_df <- get(nm)                 
  df_long <- get(paste0(nm, "_long"))
  
  # Choose covariates set for this dataset
  covars_base <- setdiff(covars, "top_party")
  covars_s <- switch(
    nm,
    "pool" = covars,  
    "us"   = c(covars_base, "eval_dem_us", "eval_rep_us"),
    "pl"   = c(covars_base, "eval_pis_pl", "eval_kon_pl", "eval_ko_pl", "eval_lew_pl"),
    "sk"   = c(covars_base, "eval_dpk_kr", "eval_ppp_kr")
  )
  
  # Build the join data: always include pc_vars, covars_s, AND any missing eval_* columns
  to_join <- base_df %>%
    dplyr::select(
      RespondentID,
      dplyr::all_of(pc_vars),
      dplyr::all_of(covars_s),
      dplyr::all_of(setdiff(eval_cols, c(pc_vars, covars_s)))
    )
  
  ## Overall: dat_long_
  
  dat_long <- df_long %>%
    dplyr::mutate(
      is_cand_dem = as.integer(cand_system_change %in% pro_dem_labels),
      chosen = as.integer(chosen)
    ) %>%
    dplyr::left_join(to_join, by = "RespondentID")
  
  assign(paste0("dat_long_", nm), dat_long)
  
  ## Cross-party: dat_h2_
  
  long2 <- dat_long %>%
    mutate(
      party_dem_us = country == "US" & cand_party == "Democrat",
      party_rep_us = country == "US" & cand_party == "Republican",
      party_dpk_kr = country == "KR" & cand_party == "DPK",
      party_ppp_kr = country == "KR" & cand_party == "PPP",
      party_pis_pl = country == "PL" & cand_party == "PiS",
      party_kon_pl = country == "PL" & cand_party == "Konfederacja",
      party_ko_pl  = country == "PL" & cand_party == "KO",
      party_lew_pl = country == "PL" & cand_party == "Lewica"
    ) %>%
    mutate(
      in_party = (
        (party_dem_us & eval_dem_us >= 6) | (party_rep_us & eval_rep_us >= 6) |
          (party_dpk_kr & eval_dpk_kr >= 6) | (party_ppp_kr & eval_ppp_kr >= 6) |
          (party_pis_pl & eval_pis_pl >= 6) | (party_kon_pl & eval_kon_pl >= 6) |
          (party_ko_pl  & eval_ko_pl  >= 6) | (party_lew_pl & eval_lew_pl >= 6)
      ),
      out_party = (
        (party_dem_us & eval_dem_us <= 4) | (party_rep_us & eval_rep_us <= 4) |
          (party_dpk_kr & eval_dpk_kr <= 4) | (party_ppp_kr & eval_ppp_kr <= 4) |
          (party_pis_pl & eval_pis_pl <= 4) | (party_kon_pl & eval_kon_pl <= 4) |
          (party_ko_pl  & eval_ko_pl  <= 4) | (party_lew_pl & eval_lew_pl <= 4)
      )
    )
  
  # Drop respondents with ambiguous positioning within-country
  keep_resp <- base_df %>%
    mutate(
      dem_in_us  = if ("eval_dem_us" %in% names(.))  eval_dem_us >= 6 else NA,
      rep_in_us  = if ("eval_rep_us" %in% names(.))  eval_rep_us >= 6 else NA,
      dem_out_us = if ("eval_dem_us" %in% names(.))  eval_dem_us <= 4 else NA,
      rep_out_us = if ("eval_rep_us" %in% names(.))  eval_rep_us <= 4 else NA,
      
      dpk_in_kr  = if ("eval_dpk_kr" %in% names(.))  eval_dpk_kr >= 6 else NA,
      ppp_in_kr  = if ("eval_ppp_kr" %in% names(.))  eval_ppp_kr >= 6 else NA,
      dpk_out_kr = if ("eval_dpk_kr" %in% names(.))  eval_dpk_kr <= 4 else NA,
      ppp_out_kr = if ("eval_ppp_kr" %in% names(.))  eval_ppp_kr <= 4 else NA,
      
      pis_in_pl  = if ("eval_pis_pl" %in% names(.))  eval_pis_pl >= 6 else NA,
      kon_in_pl  = if ("eval_kon_pl" %in% names(.))  eval_kon_pl >= 6 else NA,
      ko_in_pl   = if ("eval_ko_pl"  %in% names(.))  eval_ko_pl  >= 6 else NA,
      lew_in_pl  = if ("eval_lew_pl" %in% names(.))  eval_lew_pl >= 6 else NA,
      pis_out_pl = if ("eval_pis_pl" %in% names(.))  eval_pis_pl <= 4 else NA,
      kon_out_pl = if ("eval_kon_pl" %in% names(.))  eval_kon_pl <= 4 else NA,
      ko_out_pl  = if ("eval_ko_pl"  %in% names(.))  eval_ko_pl  <= 4 else NA,
      lew_out_pl = if ("eval_lew_pl" %in% names(.))  eval_lew_pl <= 4 else NA
    ) %>%
    mutate(
      # US & KR: both in or both out
      drop_us = country == "US" & (
        (dem_in_us  & rep_in_us) |
          (dem_out_us & rep_out_us)
      ),
      drop_kr = country == "KR" & (
        (dpk_in_kr  & ppp_in_kr) |
          (dpk_out_kr & ppp_out_kr)
      ),
      # PL: only drop if ALL FOUR in or ALL FOUR out, and treat NAs as FALSE
      all_in_pl  = rowSums(cbind(pis_in_pl, kon_in_pl, ko_in_pl, lew_in_pl),
                           na.rm = FALSE) == 4,
      all_out_pl = rowSums(cbind(pis_out_pl, kon_out_pl, ko_out_pl, lew_out_pl),
                           na.rm = FALSE) == 4,
      all_in_pl  = tidyr::replace_na(all_in_pl,  FALSE),
      all_out_pl = tidyr::replace_na(all_out_pl, FALSE),
      drop_pl    = country == "PL" & (all_in_pl | all_out_pl),
      
      drop = (drop_us | drop_kr | drop_pl)
    ) %>%
    dplyr::filter(!drop) %>%
    dplyr::select(RespondentID)
  
  long2 <- long2 %>%
    dplyr::semi_join(keep_resp, by = "RespondentID")
  
  # Keep only pairs with undem in-party AND pro-dem out-party
  pair_flags <- long2 %>%
    dplyr::group_by(RespondentID, task) %>%
    dplyr::summarise(
      has_undem_in = any(is_cand_dem == 0 & in_party,  na.rm = TRUE),
      has_dem_out  = any(is_cand_dem == 1 & out_party, na.rm = TRUE),
      .groups = "drop"
    ) %>%
    dplyr::mutate(keep_pair = has_undem_in & has_dem_out) %>%
    dplyr::filter(keep_pair) %>%
    dplyr::select(RespondentID, task)
  
  dat_h2 <- long2 %>%
    dplyr::inner_join(pair_flags, by = c("RespondentID","task"))
  
  assign(paste0("dat_h2_", nm), dat_h2)
  
  ## Cross-ideology: dat_h3_r_, dat_h3_g_
  
  long3 <- dat_long
  ideonames  <- c("redist_1","redist_2","gend_pol_1","gend_pol_2")
  missing_ideo <- setdiff(ideonames, names(long3))
  if (length(missing_ideo) > 0) {
    long3 <- long3 %>%
      dplyr::left_join(
        base_df %>% dplyr::select(RespondentID, dplyr::all_of(ideonames)),
        by = "RespondentID"
      )
  }
  
  long3 <- long3 %>%
    dplyr::mutate(
      # Candidate policy codings (redistribution)
      cand_redist_1 = dplyr::case_when(
        cand_tax_spend == "spend more"    ~ 4.5,
        cand_tax_spend == "spend less"    ~ 1.5,
        TRUE ~ NA_real_
      ),
      cand_redist_2 = dplyr::case_when(
        cand_tax_spend == "tax rich more" ~ 4.5,
        cand_tax_spend == "tax rich less" ~ 1.5,
        TRUE ~ NA_real_
      ),
      # Candidate policy codings (gender)
      cand_gend_1 = dplyr::case_when(
        cand_gend_policy == "require hire pref" ~ 4.5,
        cand_gend_policy == "no hire pref"      ~ 1.5,
        TRUE ~ NA_real_
      ),
      cand_gend_2 = dplyr::case_when(
        cand_gend_policy == "require quotas" ~ 4.5,
        cand_gend_policy == "no quotas"      ~ 1.5,
        TRUE ~ NA_real_
      ),
      # Absolute distances
      dist_r1 = abs(redist_1   - cand_redist_1),
      dist_r2 = abs(redist_2   - cand_redist_2),
      dist_g1 = abs(gend_pol_1 - cand_gend_1),
      dist_g2 = abs(gend_pol_2 - cand_gend_2)
    ) %>%
    dplyr::rowwise() %>%
    dplyr::mutate(
      prox_r    = any(c(dist_r1, dist_r2) == 0.5, na.rm = TRUE),
      any_r_far = any(c(dist_r1, dist_r2) %in% c(1.5, 2.5, 3.5), na.rm = TRUE),
      dist_r    = (!prox_r) & isTRUE(any_r_far),
      prox_g    = any(c(dist_g1, dist_g2) == 0.5, na.rm = TRUE),
      any_g_far = any(c(dist_g1, dist_g2) %in% c(1.5, 2.5, 3.5), na.rm = TRUE),
      dist_g    = (!prox_g) & isTRUE(any_g_far)
    ) %>%
    dplyr::ungroup()
  
  # Redistribution: Anti-dem proximate & Pro-dem distal
  pair_flags_r <- long3 %>%
    dplyr::group_by(RespondentID, task) %>%
    dplyr::summarise(
      has_undem_prox_r = any(is_cand_dem == 0 & prox_r, na.rm = TRUE),
      has_dem_dist_r   = any(is_cand_dem == 1 & dist_r, na.rm = TRUE),
      .groups = "drop"
    ) %>%
    dplyr::mutate(keep_pair_r = has_undem_prox_r & has_dem_dist_r) %>%
    dplyr::filter(keep_pair_r) %>%
    dplyr::select(RespondentID, task)
  
  dat_h3_r <- long3 %>%
    dplyr::inner_join(pair_flags_r, by = c("RespondentID","task"))
  
  assign(paste0("dat_h3_r_", nm), dat_h3_r)
  
  # Gender: Anti-dem proximate & Pro-dem distal
  pair_flags_g <- long3 %>%
    dplyr::group_by(RespondentID, task) %>%
    dplyr::summarise(
      has_undem_prox_g = any(is_cand_dem == 0 & prox_g, na.rm = TRUE),
      has_dem_dist_g   = any(is_cand_dem == 1 & dist_g, na.rm = TRUE),
      .groups = "drop"
    ) %>%
    dplyr::mutate(keep_pair_g = has_undem_prox_g & has_dem_dist_g) %>%
    dplyr::filter(keep_pair_g) %>%
    dplyr::select(RespondentID, task)
  
  dat_h3_g <- long3 %>%
    dplyr::inner_join(pair_flags_g, by = c("RespondentID","task"))
  
  assign(paste0("dat_h3_g_", nm), dat_h3_g)
}

long_list <- list(dat_long_pool, dat_long_us, dat_long_sk, dat_long_pl)
lapply(long_list, dim)

party_list <- list(dat_h2_pool, dat_h2_us, dat_h2_sk, dat_h2_pl)
lapply(party_list, dim)

ideo_r_list <- list(dat_h3_r_pool, dat_h3_r_us, dat_h3_r_sk, dat_h3_r_pl)
lapply(ideo_r_list, dim)

ideo_g_list <- list(dat_h3_g_pool, dat_h3_g_us, dat_h3_g_sk, dat_h3_g_pl)
lapply(ideo_g_list, dim)



#### Loop 2: compute ATMEs and plots-----


atme_res_list <- list()  # overall
atme_res_h2_list <- list()  # cross-party
atme_r_list <- list()  # cross-ideology (redist)
atme_g_list <- list()  # cross-ideology (gender)

p_atme_combo_all_list  <- list()
p_atme_combo_main_list <- list()

for (nm in dataset_names) {
  
  message("ATME estimation for: ", nm)
  
  base_df  <- get(nm)
  dat_long <- get(paste0("dat_long_", nm))
  dat_h2   <- get(paste0("dat_h2_", nm))
  dat_h3_r <- get(paste0("dat_h3_r_", nm))
  dat_h3_g <- get(paste0("dat_h3_g_", nm))
  
  ## Sample-specific covariates

  covars_base <- setdiff(covars, "top_party")
  covars_s <- switch(
    nm,
    "pool" = covars,  
    "us" = c(covars_base, "eval_dem_us", "eval_rep_us"),
    "pl" = c(covars_base, "eval_pis_pl", "eval_kon_pl", "eval_ko_pl", "eval_lew_pl"),
    "sk" = c(covars_base, "eval_dpk_kr", "eval_ppp_kr")
  )
  
  ## ATMEs - overall
  
  slopes_cov_list <- lapply(pc_vars, function(v) {
    preds <- c(v, covars_s)
    dsub  <- dat_long %>%
      dplyr::select(chosen, RespondentID, is_cand_dem,
                    dplyr::all_of(preds)) %>%
      tidyr::drop_na()
    
    # T=1
    m1 <- lm(reformulate(preds, "chosen"),
             data = dplyr::filter(dsub, is_cand_dem == 1))
    V1 <- vcovCL(m1, cluster = ~ RespondentID)
    ct1 <- coeftest(m1, vcov = V1)
    row1 <- tibble(
      term     = v,
      estimate = unname(ct1[v, 1]),
      se       = unname(ct1[v, 2]),
      arm      = 1L,
      spec     = "main cntrls"
    )
    
    # T=0
    m0 <- lm(reformulate(preds, "chosen"),
             data = dplyr::filter(dsub, is_cand_dem == 0))
    V0 <- vcovCL(m0, cluster = ~ RespondentID)
    ct0 <- coeftest(m0, vcov = V0)
    row0 <- tibble(
      term     = v,
      estimate = unname(ct0[v, 1]),
      se       = unname(ct0[v, 2]),
      arm      = 0L,
      spec     = "main cntrls"
    )
    
    dplyr::bind_rows(row1, row0)
  })
  slopes_cov <- dplyr::bind_rows(slopes_cov_list)
  
  preds_all <- c(pc_vars, covars_s)
  dsub_all  <- dat_long %>%
    dplyr::select(chosen, RespondentID, is_cand_dem,
                  dplyr::all_of(preds_all)) %>%
    tidyr::drop_na()
  
  m1_all  <- lm(reformulate(preds_all, "chosen"),
                data = dplyr::filter(dsub_all, is_cand_dem == 1))
  V1_all  <- vcovCL(m1_all, cluster = ~ RespondentID)
  ct1_all <- coeftest(m1_all, vcov = V1_all)
  rows1 <- tibble(
    term     = pc_vars,
    estimate = unname(ct1_all[pc_vars, 1]),
    se       = unname(ct1_all[pc_vars, 2]),
    arm      = 1L,
    spec     = "all cntrls"
  )
  
  m0_all  <- lm(reformulate(preds_all, "chosen"),
                data = dplyr::filter(dsub_all, is_cand_dem == 0))
  V0_all  <- vcovCL(m0_all, cluster = ~ RespondentID)
  ct0_all <- coeftest(m0_all, vcov = V0_all)
  rows0 <- tibble(
    term     = pc_vars,
    estimate = unname(ct0_all[pc_vars, 1]),
    se       = unname(ct0_all[pc_vars, 2]),
    arm      = 0L,
    spec     = "all cntrls"
  )
  
  slopes_all <- dplyr::bind_rows(rows1, rows0)
  
  slopes_df <- dplyr::bind_rows(slopes_cov, slopes_all) %>%
    dplyr::mutate(
      var  = factor(dplyr::recode(term, !!!label_map),
                    levels = rev(unname(label_map))),
      arm  = factor(
        arm,
        levels = c(1, 0),
        labels = c("pro-dem statement (T=1)",
                   "undem statement (T=0)")
      ),
      spec = factor(spec, levels = spec_levels),
      lo   = estimate - 1.96 * se,
      hi   = estimate + 1.96 * se
    )
  
  t1 <- slopes_df %>%
    dplyr::filter(arm == "pro-dem statement (T=1)") %>%
    dplyr::select(term, spec, estimate_T1 = estimate, se_T1 = se)
  
  t0 <- slopes_df %>%
    dplyr::filter(arm == "undem statement (T=0)") %>%
    dplyr::select(term, spec, estimate_T0 = estimate, se_T0 = se)
  
  atme_res <- dplyr::left_join(t1, t0, by = c("term", "spec")) %>%
    dplyr::mutate(
      atme = estimate_T1 - estimate_T0,
      se   = sqrt(se_T1^2 + se_T0^2),
      lo   = atme - 1.96 * se,
      hi   = atme + 1.96 * se,
      p    = 2 * pnorm(abs(atme / se), lower.tail = FALSE),
      var  = factor(dplyr::recode(term, !!!label_map),
                    levels = rev(unname(label_map))),
      spec = factor(spec, levels = spec_levels)
    )
  
  atme_res_list[[nm]] <- atme_res
  
  
  ## ATMEs - cross-party (H2) 
  
  slopes_cov_h2_list <- list(); idx <- 0L
  for (v in pc_vars) {
    preds <- c(v, covars_s)
    dsub  <- dat_h2 %>%
      dplyr::select(chosen, RespondentID, is_cand_dem,
                    dplyr::all_of(preds)) %>%
      tidyr::drop_na()
    for (arm in c(1L, 0L)) {
      d_arm <- dsub %>% dplyr::filter(is_cand_dem == arm)
      m     <- lm(reformulate(preds, "chosen"), data = d_arm)
      V     <- vcovCL(m, cluster = d_arm$RespondentID)
      ct    <- coeftest(m, vcov = V)
      est   <- unname(ct[v, 1]); se <- unname(ct[v, 2])
      idx   <- idx + 1L
      slopes_cov_h2_list[[idx]] <- tibble(
        term     = v,
        estimate = est,
        se       = se,
        arm      = arm,
        spec     = "main cntrls"
      )
    }
  }
  slopes_cov_h2 <- dplyr::bind_rows(slopes_cov_h2_list)
  
  preds_all_h2 <- c(pc_vars, covars_s)
  dsub_all_h2  <- dat_h2 %>%
    dplyr::select(RespondentID, chosen, is_cand_dem,
                  dplyr::all_of(preds_all_h2)) %>%
    tidyr::drop_na()
  
  slopes_all_h2_list <- list(); idx <- 0L
  for (arm in c(1L, 0L)) {
    d_arm <- dsub_all_h2 %>% dplyr::filter(is_cand_dem == arm)
    m     <- lm(reformulate(preds_all_h2, "chosen"), data = d_arm)
    V     <- vcovCL(m, cluster = d_arm$RespondentID)
    ct    <- coeftest(m, vcov = V)
    for (v in pc_vars) {
      est <- unname(ct[v, 1]); se <- unname(ct[v, 2])
      idx <- idx + 1L
      slopes_all_h2_list[[idx]] <- tibble(
        term     = v,
        estimate = est,
        se       = se,
        arm      = arm,
        spec     = "all cntrls"
      )
    }
  }
  slopes_all_h2 <- dplyr::bind_rows(slopes_all_h2_list)
  
  slopes_df_h2 <- dplyr::bind_rows(slopes_cov_h2, slopes_all_h2) %>%
    dplyr::mutate(
      var  = factor(dplyr::recode(term, !!!label_map),
                    levels = rev(unname(label_map))),
      arm  = factor(
        arm,
        levels = c(1, 0),
        labels = c("Pro-dem out-party (T=1)",
                   "Anti-dem in-party (T=0)")
      ),
      spec = factor(spec, levels = spec_levels),
      lo   = estimate - 1.96 * se,
      hi   = estimate + 1.96 * se
    )
  
  t1_h2 <- slopes_df_h2 %>%
    dplyr::filter(arm == "Pro-dem out-party (T=1)") %>%
    dplyr::select(term, spec, estimate_T1 = estimate, se_T1 = se)
  
  t0_h2 <- slopes_df_h2 %>%
    dplyr::filter(arm == "Anti-dem in-party (T=0)") %>%
    dplyr::select(term, spec, estimate_T0 = estimate, se_T0 = se)
  
  atme_res_h2 <- dplyr::left_join(t1_h2, t0_h2, by = c("term", "spec")) %>%
    dplyr::mutate(
      atme = estimate_T1 - estimate_T0,
      se   = sqrt(se_T1^2 + se_T0^2),
      lo   = atme - 1.96 * se,
      hi   = atme + 1.96 * se,
      p    = 2 * pnorm(abs(atme / se), lower.tail = FALSE),
      var  = factor(dplyr::recode(term, !!!label_map),
                    levels = rev(unname(label_map))),
      spec = factor(spec, levels = spec_levels)
    )
  
  atme_res_h2_list[[nm]] <- atme_res_h2
  
  
  ## ATMEs - cross-ideology (H3: redist)
  
  slopes_cov_r_list <- list()
  k <- 0L
  for (v in pc_vars) {
    preds <- c(v, covars_s)
    dsub  <- dat_h3_r %>%
      dplyr::select(RespondentID, chosen, is_cand_dem,
                    dplyr::all_of(preds)) %>%
      tidyr::drop_na()
    d1 <- dsub %>% dplyr::filter(is_cand_dem == 1L)
    m1 <- lm(reformulate(preds, "chosen"), data = d1)
    V1 <- vcovCL(m1, cluster = d1$RespondentID)
    ct1 <- coeftest(m1, vcov = V1); e1 <- unname(ct1[v,1]); s1 <- unname(ct1[v,2])
    k <- k + 1L; slopes_cov_r_list[[k]] <- tibble(term=v, estimate=e1,
                                                  se=s1, arm=1, spec="main cntrls")
    d0 <- dsub %>% dplyr::filter(is_cand_dem == 0L)
    m0 <- lm(reformulate(preds, "chosen"), data = d0)
    V0 <- vcovCL(m0, cluster = d0$RespondentID)
    ct0 <- coeftest(m0, vcov = V0); e0 <- unname(ct0[v,1]); s0 <- unname(ct0[v,2])
    k <- k + 1L; slopes_cov_r_list[[k]] <- tibble(term=v, estimate=e0,
                                                  se=s0, arm=0, spec="main cntrls")
  }
  slopes_cov_r <- dplyr::bind_rows(slopes_cov_r_list)
  
  preds_all_h3 <- c(pc_vars, covars_s)
  dsub_all_r   <- dat_h3_r %>%
    dplyr::select(RespondentID, chosen, is_cand_dem,
                  dplyr::all_of(preds_all_h3)) %>%
    tidyr::drop_na()
  
  slopes_all_r_list <- list(); k <- 0L
  for (arm in c(1L, 0L)) {
    d_arm <- dsub_all_r %>% dplyr::filter(is_cand_dem == arm)
    m <- lm(reformulate(preds_all_h3, "chosen"), data = d_arm)
    V <- vcovCL(m, cluster = d_arm$RespondentID)
    ct <- coeftest(m, vcov = V)
    for (v in pc_vars) {
      e <- unname(ct[v,1]); s <- unname(ct[v,2])
      k <- k + 1L
      slopes_all_r_list[[k]] <- tibble(term=v, estimate=e, se=s,
                                       arm=arm, spec="all cntrls")
    }
  }
  slopes_all_r <- dplyr::bind_rows(slopes_all_r_list)
  
  slopes_df_r <- dplyr::bind_rows(slopes_cov_r, slopes_all_r) %>%
    dplyr::mutate(
      var  = factor(dplyr::recode(term, !!!label_map),
                    levels = rev(unname(label_map))),
      arm  = factor(
        arm,
        levels = c(1, 0),
        labels = c("Pro-dem distal (T=1)",
                   "Anti-dem proximate (T=0)")
      ),
      spec = factor(spec, levels = spec_levels),
      lo   = estimate - 1.96*se,
      hi   = estimate + 1.96*se
    )
  
  t1_r <- slopes_df_r %>%
    dplyr::filter(arm == "Pro-dem distal (T=1)") %>%
    dplyr::select(term, spec, estimate_T1=estimate, se_T1=se)
  t0_r <- slopes_df_r %>%
    dplyr::filter(arm == "Anti-dem proximate (T=0)") %>%
    dplyr::select(term, spec, estimate_T0=estimate, se_T0=se)
  
  atme_r <- dplyr::left_join(t1_r, t0_r, by = c("term","spec")) %>%
    dplyr::mutate(
      atme = estimate_T1 - estimate_T0,
      se   = sqrt(se_T1^2 + se_T0^2),
      lo   = atme - 1.96*se,
      hi   = atme + 1.96*se,
      p    = 2*pnorm(abs(atme/se), lower.tail = FALSE),
      var  = factor(dplyr::recode(term, !!!label_map),
                    levels = rev(unname(label_map))),
      spec = factor(spec, levels = spec_levels)
    ) %>%
    dplyr::mutate(
      y_base = as.numeric(var),
      y      = y_base + offsets[as.character(spec)]
    )
  
  atme_r_list[[nm]] <- atme_r
  
  
  ## ATMEs - cross-ideology (H3: gender)
  
  slopes_cov_g_list <- list()
  k <- 0L
  for (v in pc_vars) {
    preds <- c(v, covars_s)
    dsub  <- dat_h3_g %>%
      dplyr::select(RespondentID, chosen, is_cand_dem,
                    dplyr::all_of(preds)) %>%
      tidyr::drop_na()
    d1 <- dsub %>% dplyr::filter(is_cand_dem == 1L)
    m1 <- lm(reformulate(preds, "chosen"), data = d1)
    V1 <- vcovCL(m1, cluster = d1$RespondentID)
    ct1 <- coeftest(m1, vcov = V1); e1 <- unname(ct1[v,1]); s1 <- unname(ct1[v,2])
    k <- k + 1L; slopes_cov_g_list[[k]] <- tibble(term=v, estimate=e1,
                                                  se=s1, arm=1, spec="main cntrls")
    d0 <- dsub %>% dplyr::filter(is_cand_dem == 0L)
    m0 <- lm(reformulate(preds, "chosen"), data = d0)
    V0 <- vcovCL(m0, cluster = d0$RespondentID)
    ct0 <- coeftest(m0, vcov = V0); e0 <- unname(ct0[v,1]); s0 <- unname(ct0[v,2])
    k <- k + 1L; slopes_cov_g_list[[k]] <- tibble(term=v, estimate=e0,
                                                  se=s0, arm=0, spec="main cntrls")
  }
  slopes_cov_g <- dplyr::bind_rows(slopes_cov_g_list)
  
  dsub_all_g <- dat_h3_g %>%
    dplyr::select(RespondentID, chosen, is_cand_dem,
                  dplyr::all_of(preds_all_h3)) %>%
    tidyr::drop_na()
  
  slopes_all_g_list <- list(); k <- 0L
  for (arm in c(1L, 0L)) {
    d_arm <- dsub_all_g %>% dplyr::filter(is_cand_dem == arm)
    m <- lm(reformulate(preds_all_h3, "chosen"), data = d_arm)
    V <- vcovCL(m, cluster = d_arm$RespondentID)
    ct <- coeftest(m, vcov = V)
    for (v in pc_vars) {
      e <- unname(ct[v,1]); s <- unname(ct[v,2])
      k <- k + 1L
      slopes_all_g_list[[k]] <- tibble(term=v, estimate=e, se=s,
                                       arm=arm, spec="all cntrls")
    }
  }
  slopes_all_g <- dplyr::bind_rows(slopes_all_g_list)
  
  slopes_df_g <- dplyr::bind_rows(slopes_cov_g, slopes_all_g) %>%
    dplyr::mutate(
      var  = factor(dplyr::recode(term, !!!label_map),
                    levels = rev(unname(label_map))),
      arm  = factor(
        arm,
        levels = c(1, 0),
        labels = c("Pro-dem distal (T=1)",
                   "Anti-dem proximate (T=0)")
      ),
      spec = factor(spec, levels = spec_levels),
      lo   = estimate - 1.96*se,
      hi   = estimate + 1.96*se
    )
  
  t1_g <- slopes_df_g %>%
    dplyr::filter(arm == "Pro-dem distal (T=1)") %>%
    dplyr::select(term, spec, estimate_T1=estimate, se_T1=se)
  t0_g <- slopes_df_g %>%
    dplyr::filter(arm == "Anti-dem proximate (T=0)") %>%
    dplyr::select(term, spec, estimate_T0=estimate, se_T0=se)
  
  atme_g <- dplyr::left_join(t1_g, t0_g, by = c("term","spec")) %>%
    dplyr::mutate(
      atme = estimate_T1 - estimate_T0,
      se   = sqrt(se_T1^2 + se_T0^2),
      lo   = atme - 1.96*se,
      hi   = atme + 1.96*se,
      p    = 2*pnorm(abs(atme/se), lower.tail = FALSE),
      var  = factor(dplyr::recode(term, !!!label_map),
                    levels = rev(unname(label_map))),
      spec = factor(spec, levels = spec_levels)
    ) %>%
    dplyr::mutate(
      y_base = as.numeric(var),
      y      = y_base + offsets[as.character(spec)]
    )
  
  atme_g_list[[nm]] <- atme_g
  
  ## Combined ATME plots
  
  title_label <- switch(
    nm,
    "pool" = "Pooled",
    "us"   = "United States",
    "sk"   = "South Korea",
    "pl"   = "Poland",
    nm
  )
  
  # All covariates (including culture vars)
  atme_combo_all <- dplyr::bind_rows(
    atme_res    %>% dplyr::filter(spec == "all cntrls") %>%
      dplyr::mutate(series = "Overall"),
    atme_res_h2 %>% dplyr::filter(spec == "all cntrls") %>%
      dplyr::mutate(series = "Cross-party"),
    atme_r      %>% dplyr::filter(spec == "all cntrls") %>%
      dplyr::mutate(series = "Cross-ideology (redist)"),
    atme_g      %>% dplyr::filter(spec == "all cntrls") %>%
      dplyr::mutate(series = "Cross-ideology (gend)")
  ) %>%
    dplyr::mutate(
      var    = factor(dplyr::recode(term, !!!label_map),
                      levels = rev(unname(label_map))),
      y_base = as.numeric(var),
      y      = y_base + series_offsets[series]
    )
  
  atme_combo_all$var <- factor(
    atme_combo_all$var,
    levels = rev(label_map[pc_order])
  )
  
  p_atme_combo_all <- ggplot(atme_combo_all,
                             aes(x = atme, y = y,
                                 color = series, shape = series, fill = series)) +
    geom_vline(xintercept = 0, linetype = 3, linewidth = 0.4) +
    geom_errorbarh(aes(xmin = lo, xmax = hi), height = 0, linewidth = 1) +
    geom_point(size = 2.5, stroke = 0.6) +
    scale_y_continuous(
      breaks = sort(unique(atme_combo_all$y_base)),
      labels = levels(atme_combo_all$var)
    ) +
    scale_color_manual(values = atme_pal, breaks = names(atme_pal), name = NULL) +
    scale_fill_manual(values = atme_pal,  breaks = names(atme_pal), name = NULL) +
    scale_shape_manual(values = atme_shapes, breaks = names(atme_shapes), name = NULL) +
    guides(color = guide_legend(nrow = 1), shape = guide_legend(nrow = 1)) +
    labs(
      title = title_label,
      x     = "ATME",
      y     = NULL,
      color = NULL,
      shape = NULL,
      fill  = NULL
    ) +
    theme_bw(base_size = 14) +
    theme(
      panel.grid.minor  = element_blank(),
      legend.position   = "bottom",
      legend.box.spacing= grid::unit(0.1, "lines"),
      legend.key.height = grid::unit(0.4, "lines"),
      legend.text       = element_text(size = 11)
    )
  
  p_atme_combo_all_list[[nm]] <- p_atme_combo_all
  
  ggsave(file.path(project_root, "outputs", sprintf("ATMEs_combined_allcovar_%s_2.png", nm)),
         p_atme_combo_all, width = 7, height = 5, dpi = 250)
  
  
  # Main covariates only (“main cntrls”)
  atme_combo_main <- dplyr::bind_rows(
    atme_res    %>% dplyr::filter(spec == "main cntrls") %>%
      dplyr::mutate(series = "Overall"),
    atme_res_h2 %>% dplyr::filter(spec == "main cntrls") %>%
      dplyr::mutate(series = "Cross-party"),
    atme_r      %>% dplyr::filter(spec == "main cntrls") %>%
      dplyr::mutate(series = "Cross-ideology (redist)"),
    atme_g      %>% dplyr::filter(spec == "main cntrls") %>%
      dplyr::mutate(series = "Cross-ideology (gend)")
  ) %>%
    dplyr::mutate(
      var    = factor(dplyr::recode(term, !!!label_map),
                      levels = rev(unname(label_map))),
      y_base = as.numeric(var),
      y      = y_base + series_offsets[series]
    )
  
  atme_combo_main$var <- factor(
    atme_combo_main$var,
    levels = rev(label_map[pc_order])
  )
  
  p_atme_combo_main <- ggplot(atme_combo_main,
                              aes(x = atme, y = y,
                                  color = series, shape = series, fill = series)) +
    geom_vline(xintercept = 0, linetype = 3, linewidth = 0.4) +
    geom_errorbarh(aes(xmin = lo, xmax = hi), height = 0, linewidth = 1) +
    geom_point(size = 2.5, stroke = 0.6) +
    scale_y_continuous(
      breaks = sort(unique(atme_combo_main$y_base)),
      labels = levels(atme_combo_main$var)
    ) +
    scale_color_manual(values = atme_pal, breaks = names(atme_pal), name = NULL) +
    scale_fill_manual(values = atme_pal,  breaks = names(atme_pal), name = NULL) +
    scale_shape_manual(values = atme_shapes, breaks = names(atme_shapes), name = NULL) +
    guides(color = guide_legend(nrow = 1), shape = guide_legend(nrow = 1)) +
    labs(
      title = title_label,
      x     = "ATME",
      y     = NULL,
      color = NULL,
      shape = NULL,
      fill  = NULL
    ) +
    theme_bw(base_size = 14) +
    theme(
      panel.grid.minor  = element_blank(),
      legend.position   = "bottom",
      legend.box.spacing= grid::unit(0.1, "lines"),
      legend.key.height = grid::unit(0.4, "lines"),
      legend.text       = element_text(size = 11)
    )
  
  p_atme_combo_main_list[[nm]] <- p_atme_combo_main
  
  ggsave(file.path(project_root, "outputs", sprintf("ATMEs_combined_maincovar_%s_2.png", nm)),
         p_atme_combo_main, width = 7, height = 5, dpi = 250)
}


