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)
#### Variable by variable plots, with samples, excluding pooled, in rows

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

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"))

# 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"))

# Ensure all eval vars exist in all three 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("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",
             "emancip_val", "autonomy_values")

pc_labels <- c(
  demaut_pref = "Democ-autoc. preference",
  sup_ldem = "Support liberal democracy",
  inst_trust = "Institutional trust",
  pol_tolerance = "Political tolerance",
  emancip_val = "Emancipative values",
  autonomy_values = "Authoritarian values (rev)"
)

sample_labels <- c(
  pl = "Poland",
  sk = "South Korea",
  us = "United States"
)

scenario_order <- c("Overall", "Cross-party",
                    "Cross-ideology (redist)", "Cross-ideology (gender)")

scenario_colors <- c(
  "Overall" = "#000000",
  "Cross-party" = "goldenrod",
  "Cross-ideology (redist)" = "dodgerblue3",
  "Cross-ideology (gender)" = "darkorchid3"
)

scenario_shapes <- c(
  "Overall" = 16,
  "Cross-party" = 15,
  "Cross-ideology (redist)" = 17,
  "Cross-ideology (gender)" = 25
)

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

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

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

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

for (nm in dataset_names) {

  message("Building datasets for: ", nm)

  base_df <- get(nm)
  df_long <- get(paste0(nm, "_long"))

  covars_base <- setdiff(covars, "top_party")
  covars_s <- switch(
    nm,
    "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")
  )

  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 <- 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
  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)
      )
    )

  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(
      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)),
      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)
    ) %>%
    filter(!drop) %>%
    select(RespondentID)

  long2 <- long2 %>%
    dplyr::semi_join(keep_resp, by = "RespondentID")

  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::filter(has_undem_in & has_dem_out) %>%
    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
  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(
      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_
      ),
      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_
      ),
      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()

  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::filter(has_undem_prox_r & has_dem_dist_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)

  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::filter(has_undem_prox_g & has_dem_dist_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)
}

#### Loop 2: compute ATMEs ----

atme_plot_res <- NULL

for (nm in dataset_names) {

  message("ATME estimation for: ", 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))

  covars_base <- setdiff(covars, "top_party")
  covars_s <- switch(
    nm,
    "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")
  )

  scenario_data <- list(
    "Overall" = dat_long,
    "Cross-party" = dat_h2,
    "Cross-ideology (redist)" = dat_h3_r,
    "Cross-ideology (gender)" = dat_h3_g
  )

  for (series_lab in names(scenario_data)) {
    dat_use <- scenario_data[[series_lab]]

    ## Minimal controls
    for (v in pc_vars) {
      preds <- c(v, covars_s)
      dsub <- dat_use %>%
        dplyr::select(chosen, RespondentID, 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)

      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)

      est1 <- unname(ct1[v, 1]); se1 <- unname(ct1[v, 2])
      est0 <- unname(ct0[v, 1]); se0 <- unname(ct0[v, 2])

      atme <- est1 - est0
      se   <- sqrt(se1^2 + se0^2)

      atme_plot_res <- bind_rows(
        atme_plot_res,
        tibble(
          sample = nm,
          sample_label = sample_labels[[nm]],
          scenario = series_lab,
          spec = "Minimal controls",
          pc_var = v,
          pc_label = pc_labels[[v]],
          atme = atme,
          se = se,
          lo = atme - 1.96 * se,
          hi = atme + 1.96 * se
        )
      )
    }

    ## Full controls
    preds_all <- c(pc_vars, covars_s)
    dsub_all <- dat_use %>%
      dplyr::select(chosen, RespondentID, is_cand_dem, dplyr::all_of(preds_all)) %>%
      tidyr::drop_na()

    d1_all <- dsub_all %>% dplyr::filter(is_cand_dem == 1L)
    m1_all <- lm(reformulate(preds_all, "chosen"), data = d1_all)
    V1_all <- vcovCL(m1_all, cluster = d1_all$RespondentID)
    ct1_all <- coeftest(m1_all, vcov = V1_all)

    d0_all <- dsub_all %>% dplyr::filter(is_cand_dem == 0L)
    m0_all <- lm(reformulate(preds_all, "chosen"), data = d0_all)
    V0_all <- vcovCL(m0_all, cluster = d0_all$RespondentID)
    ct0_all <- coeftest(m0_all, vcov = V0_all)

    for (v in pc_vars) {
      est1 <- unname(ct1_all[v, 1]); se1 <- unname(ct1_all[v, 2])
      est0 <- unname(ct0_all[v, 1]); se0 <- unname(ct0_all[v, 2])

      atme <- est1 - est0
      se   <- sqrt(se1^2 + se0^2)

      atme_plot_res <- bind_rows(
        atme_plot_res,
        tibble(
          sample = nm,
          sample_label = sample_labels[[nm]],
          scenario = series_lab,
          spec = "Full controls",
          pc_var = v,
          pc_label = pc_labels[[v]],
          atme = atme,
          se = se,
          lo = atme - 1.96 * se,
          hi = atme + 1.96 * se
        )
      )
    }
  }
}

#### Loop 3: output table and plots ----

if (!is.null(atme_plot_res) && nrow(atme_plot_res) > 0) {

  write.csv(atme_plot_res, file.path(project_root, "outputs", "ATME_results_faceted_v2.csv"), row.names = FALSE)

  atme_plot_res$scenario <- factor(
    atme_plot_res$scenario,
    levels = scenario_order
  )
  atme_plot_res$sample <- factor(atme_plot_res$sample, levels = c("pl", "sk", "us"))
  atme_plot_res$sample_label <- factor(
    atme_plot_res$sample_label,
    levels = c("Poland", "South Korea", "United States")
  )
  atme_plot_res$pc_label <- factor(
    atme_plot_res$pc_label,
    levels = unname(pc_labels[pc_vars])
  )
  atme_plot_res$spec <- factor(
    atme_plot_res$spec,
    levels = c("Full controls", "Minimal controls")
  )

  dodge_pos <- position_dodge(width = 0.55, reverse = TRUE)

  make_plot <- function(dat_plot, out_file) {
    p <- ggplot(
      dat_plot,
      aes(x = atme, y = sample_label, color = scenario, fill = scenario, shape = scenario)
    ) +
      geom_vline(xintercept = 0, linetype = "dotted", color = "gray40") +
      geom_errorbar(
        aes(xmin = lo, xmax = hi),
        orientation = "y",
        width = 0.12,
        linewidth = 0.6,
        position = dodge_pos
      ) +
      geom_point(size = 2.7, stroke = 0.7, position = dodge_pos) +
      facet_wrap(~ pc_label, ncol = 2) +
      scale_color_manual(values = scenario_colors) +
      scale_fill_manual(values = scenario_colors) +
      scale_shape_manual(values = scenario_shapes) +
      labs(
        x = "ATME",
        y = NULL,
        color = NULL,
        fill = NULL
      ) +
      guides(
        color = guide_legend(
          override.aes = list(shape = unname(scenario_shapes),
                              fill = unname(scenario_colors),
                              size = 3.2)
        ),
        fill = "none",
        shape = "none"
      ) +
      theme_bw(base_size = 11) +
      theme(
        panel.grid.minor = element_blank(),
        panel.grid.major.x = element_line(color = "grey85", linewidth = 0.3),
        panel.grid.major.y = element_line(color = "grey85", linewidth = 0.3),
        strip.background = element_blank(),
        strip.text = element_text(size = 10, face = "plain"),
        legend.position = "bottom",
        legend.title = element_blank(),
        legend.margin = margin(t = 2, r = 0, b = 2, l = 0),
        legend.box.spacing = grid::unit(0.35, "lines"),
        legend.key.height = grid::unit(0.32, "lines"),
        legend.text = element_text(size = 10),
        panel.spacing.x = grid::unit(0.8, "lines"),
        panel.spacing.y = grid::unit(0.8, "lines"),
        plot.margin = margin(4, 4, 10, 4)
      )

    print(p)
    ggsave(file.path(project_root, "outputs", out_file), p, width = 8.6, height = 7.6, dpi = 250)
  }

  make_plot(
    atme_plot_res %>% dplyr::filter(spec == "Full controls"),
    "ATMEs_faceted_allcovar_v2.png"
  )

  make_plot(
    atme_plot_res %>% dplyr::filter(spec == "Minimal controls"),
    "ATMEs_faceted_maincovar_v2.png"
  )
}
