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 - Bivariate nonparametric plots

# libraries
library(dplyr)
library(tidyr)
library(forcats)
library(ggplot2)
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"))

wide_list <- list(pool = pool, us = us, sk = sk, pl = pl)
long_list <- list(pool = pool_long, us = us_long, sk = sk_long, pl = pl_long)
samples   <- names(wide_list)

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

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

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

## Containers for per-respondent rates
overall_resp_list    <- vector("list", length(samples)); names(overall_resp_list)    <- samples
H2_resp_list         <- vector("list", length(samples)); names(H2_resp_list)         <- samples
H3_resp_redist_list  <- vector("list", length(samples)); names(H3_resp_redist_list)  <- samples
H3_resp_gender_list  <- vector("list", length(samples)); names(H3_resp_gender_list)  <- samples


#### Calculate rates for each sample ----


for (s in samples) {
  cat("Building respondent-level rates for", s, "...\n")
  
  base_df  <- wide_list[[s]]
  dat_long <- long_list[[s]]
  
  ## Ensure basic variables exist / types
  if (!("is_cand_dem" %in% names(dat_long))) {
    dat_long <- dat_long %>%
      dplyr::mutate(is_cand_dem = as.integer(cand_system_change %in% pro_dem_labels))
  }
  dat_long <- dat_long %>%
    dplyr::mutate(chosen = as.integer(chosen))
  
  non_cj_cols <- names(base_df)[!grepl("^CJ_Task_", names(base_df))]
  
  ## 1) Overall rate per respondent
  
  overall_resp <- dat_long %>%
    dplyr::mutate(pro_dem = cand_system_change %in% pro_dem_labels) %>%
    dplyr::group_by(RespondentID, task) %>%
    dplyr::summarise(
      choose_dem = as.integer(any(chosen == 1L & pro_dem)),
      .groups = "drop"
    ) %>%
    dplyr::group_by(RespondentID) %>%
    dplyr::summarise(
      rate_overall = mean(choose_dem),
      .groups = "drop"
    )
  
  overall_resp_list[[s]] <- overall_resp
  
  ## 2) Cross-party (H2) rate per respondent 

  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")
  eval_cols <- intersect(eval_cols, names(base_df))
  
  long2 <- dat_long %>%
    dplyr::left_join(base_df %>% dplyr::select(RespondentID, dplyr::all_of(eval_cols)),
                     by = "RespondentID") %>%
    dplyr::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"
    ) %>%
    dplyr::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 %>%
    dplyr::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
    ) %>%
    dplyr::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)
      ),
      # PL: drop only if ALL four parties in or ALL four out
      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_keep <- long2 %>%
    dplyr::semi_join(keep_resp, by = "RespondentID")
  
  pair_flags_h2 <- long2_keep %>%
    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)
  
  H2_resp <- long2_keep %>%
    dplyr::semi_join(pair_flags_h2, by = c("RespondentID","task")) %>%
    dplyr::group_by(RespondentID, task) %>%
    dplyr::summarise(
      choose_dem_out_vs_undem_in =
        as.integer(any(chosen == 1L & is_cand_dem == 1L & out_party)),
      .groups = "drop"
    ) %>%
    dplyr::group_by(RespondentID) %>%
    dplyr::summarise(
      rate_h2_specific = mean(choose_dem_out_vs_undem_in),
      .groups = "drop"
    )
  
  H2_resp_list[[s]] <- H2_resp
  
  ## 3) Cross-ideology (H3) – build long3 with distances
  
  ideonames <- c("redist_1","redist_2","gend_pol_1","gend_pol_2")
  missing_ideo <- setdiff(ideonames, names(dat_long))
  
  long3 <- dat_long
  if (length(missing_ideo) > 0) {
    long3 <- long3 %>%
      dplyr::left_join(base_df %>% dplyr::select(RespondentID, dplyr::all_of(missing_ideo)),
                       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::mutate(
      dist_r = dplyr::coalesce(dist_r1, dist_r2),
      dist_g = dplyr::coalesce(dist_g1, dist_g2)
    ) %>%
    dplyr::mutate(
      prox_redist   = dplyr::coalesce(dist_r == 0.5, FALSE),
      distal_redist = dplyr::coalesce(dist_r %in% c(1.5, 2.5, 3.5), FALSE),
      prox_gender   = dplyr::coalesce(dist_g == 0.5, FALSE),
      distal_gender = dplyr::coalesce(dist_g %in% c(1.5, 2.5, 3.5), FALSE)
    )
  
  ## 3a) Cross-ideology, redistribution
  
  pair_flags_h3_redist <- long3 %>%
    dplyr::group_by(RespondentID, task) %>%
    dplyr::summarise(
      has_undem_prox_r = any(is_cand_dem == 0 & prox_redist,   na.rm = TRUE),
      has_dem_dist_r   = any(is_cand_dem == 1 & distal_redist, na.rm = TRUE),
      .groups = "drop"
    ) %>%
    dplyr::filter(has_undem_prox_r & has_dem_dist_r) %>%
    dplyr::select(RespondentID, task)
  
  H3_resp_redist <- long3 %>%
    dplyr::semi_join(pair_flags_h3_redist, by = c("RespondentID","task")) %>%
    dplyr::group_by(RespondentID, task) %>%
    dplyr::summarise(
      choose_dem_dist_vs_undem_prox_redist =
        as.integer(any(chosen == 1L & is_cand_dem == 1L & distal_redist)),
      .groups = "drop"
    ) %>%
    dplyr::group_by(RespondentID) %>%
    dplyr::summarise(
      rate_h3_redist = mean(choose_dem_dist_vs_undem_prox_redist),
      .groups = "drop"
    )
  
  H3_resp_redist_list[[s]] <- H3_resp_redist
  
  ## 3b) Cross-ideology, gender
  
  pair_flags_h3_gender <- long3 %>%
    dplyr::group_by(RespondentID, task) %>%
    dplyr::summarise(
      has_undem_prox_g = any(is_cand_dem == 0 & prox_gender,   na.rm = TRUE),
      has_dem_dist_g   = any(is_cand_dem == 1 & distal_gender, na.rm = TRUE),
      .groups = "drop"
    ) %>%
    dplyr::filter(has_undem_prox_g & has_dem_dist_g) %>%
    dplyr::select(RespondentID, task)
  
  H3_resp_gender <- long3 %>%
    dplyr::semi_join(pair_flags_h3_gender, by = c("RespondentID","task")) %>%
    dplyr::group_by(RespondentID, task) %>%
    dplyr::summarise(
      choose_dem_dist_vs_undem_prox_g =
        as.integer(any(chosen == 1L & is_cand_dem == 1L & distal_gender)),
      .groups = "drop"
    ) %>%
    dplyr::group_by(RespondentID) %>%
    dplyr::summarise(
      rate_h3_gender = mean(choose_dem_dist_vs_undem_prox_g),
      .groups = "drop"
    )
  
  H3_resp_gender_list[[s]] <- H3_resp_gender
}

## Aesthetics (match 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)

sample_titles <- c(pool = "Pooled samples",
                   us   = "United States",
                   sk   = "South Korea",
                   pl   = "Poland")



#### Loop 2: build plotting frames + plots for each sample

for (s in samples) {
  cat("Building bivariate plots for", s, "...\n")
  
  base_df      <- wide_list[[s]]
  overall_resp <- overall_resp_list[[s]]
  H2_resp      <- H2_resp_list[[s]]
  H3_r         <- H3_resp_redist_list[[s]]
  H3_g         <- H3_resp_gender_list[[s]]
  
  # Assemble plotting base (one row per respondent)
  plot_base <- base_df %>%
    dplyr::select(RespondentID, dplyr::all_of(pc_vars)) %>%
    dplyr::left_join(overall_resp, by = "RespondentID") %>%
    dplyr::left_join(H2_resp,      by = "RespondentID") %>%
    dplyr::left_join(H3_r,         by = "RespondentID") %>%
    dplyr::left_join(H3_g,         by = "RespondentID")
  
  ## Redist version (overall + H2 + H3_redist) ----
  
  plot_long_redist <- NULL
  for (v in pc_vars) {
    x <- plot_base[[v]]
    if (sum(!is.na(x)) >= 2L) {
      dec <- dplyr::ntile(x, 10L)
    } else {
      dec <- rep(NA_integer_, length(x))
    }
    tmp <- plot_base %>%
      dplyr::mutate(pc_var = v, pc_decile = dec) %>%
      tidyr::pivot_longer(
        cols = c(rate_overall, rate_h2_specific, rate_h3_redist),
        names_to = "series", values_to = "rate"
      ) %>%
      dplyr::select(RespondentID, pc_var, pc_decile, series, rate)
    plot_long_redist <- dplyr::bind_rows(plot_long_redist, tmp)
  }
  
  plot_sum_redist <- plot_long_redist %>%
    dplyr::filter(!is.na(pc_decile), !is.na(rate)) %>%
    dplyr::mutate(
      pc_label = factor(dplyr::recode(pc_var, !!!label_map),
                        levels = unname(label_map[pc_vars])),
      series   = factor(series,
                        levels = c("rate_overall","rate_h2_specific","rate_h3_redist"),
                        labels = c("Overall","Cross-party","Cross-ideology (redist)"))
    ) %>%
    dplyr::group_by(pc_label, pc_decile, series) %>%
    dplyr::summarise(
      n    = dplyr::n(),
      mean = mean(rate),
      sd   = sd(rate),
      se   = sd / sqrt(n),
      lo   = pmax(0, mean - 1.96*se),
      hi   = pmin(1, mean + 1.96*se),
      .groups = "drop"
    ) %>%
    dplyr::mutate(
      x0 = as.numeric(pc_decile),
      x  = dplyr::case_when(
        series == "Cross-party" ~ x0,
        series == "Cross-ideology (redist)" ~ x0,
        TRUE ~ x0
      )
    ) %>%
    dplyr::mutate(series = forcats::fct_drop(series))
  
  show_legend <- (s == "pool")
  plot_h      <- if (show_legend) 2.80 else 2.50   
  
  series_present <- levels(droplevels(plot_sum_redist$series))
  legend_override <- list(
    shape    = unname(atme_shapes[series_present]),
    fill     = unname(atme_pal[series_present]),
    linetype = 1
  )
  
  p_three_redist <- ggplot(plot_sum_redist,
                           aes(x = x, y = mean,
                               color = series, fill = series,
                               shape = series, group = series)) +
    geom_line(linewidth = 0.6) +
    geom_point(size = 1.6, stroke = 0.5) +
    geom_errorbar(aes(ymin = lo, ymax = hi), width = 0, linewidth = 0.6) +
    scale_x_continuous(breaks = 1:10, labels = 1:10, minor_breaks = NULL) +
    scale_y_continuous(limits = c(0, 1), expand = c(0, 0)) +
    scale_color_manual(values = atme_pal, 
                       guide  = guide_legend(title = NULL, nrow  = 1,
                                             override.aes = legend_override)) +
    scale_fill_manual(values = atme_pal, guide = "none") +
    scale_shape_manual(values = atme_shapes, guide = "none") +  
    facet_wrap(~ pc_label, nrow = 1) +
    labs(
      title = sample_titles[s],
      x = "Political culture decile (low \u2192 high)",
      y = "Rate of choosing\npro-democratic candidate"
    ) +
    theme_bw(base_size = 11) +
    theme(
      panel.grid.minor   = element_blank(),
      legend.position    = if (show_legend) "bottom" else "none",
      legend.box.spacing = grid::unit(0.1, "lines"),
      legend.key.height  = grid::unit(0.4, "lines"),
      legend.text        = element_text(size = 11)
    )
  
  print(p_three_redist)
  
  ggsave(file.path(project_root, "outputs", paste0("bivar_prodem_redist_", s, "_2.png")),
         p_three_redist, width = 8, height = plot_h, dpi = 250)
  
  ## Gender version (overall + H2 + H3_gender)----
  
  plot_long_gender <- NULL
  for (v in pc_vars) {
    x <- plot_base[[v]]
    if (sum(!is.na(x)) >= 2L) {
      dec <- dplyr::ntile(x, 10L)
    } else {
      dec <- rep(NA_integer_, length(x))
    }
    tmp <- plot_base %>%
      dplyr::mutate(pc_var = v, pc_decile = dec) %>%
      tidyr::pivot_longer(
        cols = c(rate_overall, rate_h2_specific, rate_h3_gender),
        names_to = "series", values_to = "rate"
      ) %>%
      dplyr::select(RespondentID, pc_var, pc_decile, series, rate)
    plot_long_gender <- dplyr::bind_rows(plot_long_gender, tmp)
  }
  
  plot_sum_gender <- plot_long_gender %>%
    dplyr::filter(!is.na(pc_decile), !is.na(rate)) %>%
    dplyr::mutate(
      pc_label = factor(dplyr::recode(pc_var, !!!label_map),
                        levels = unname(label_map[pc_vars])),
      series   = factor(series,
                        levels = c("rate_overall","rate_h2_specific","rate_h3_gender"),
                        labels = c("Overall","Cross-party","Cross-ideology (gend)"))
    ) %>%
    dplyr::group_by(pc_label, pc_decile, series) %>%
    dplyr::summarise(
      n    = dplyr::n(),
      mean = mean(rate),
      sd   = sd(rate),
      se   = sd / sqrt(n),
      lo   = pmax(0, mean - 1.96*se),
      hi   = pmin(1, mean + 1.96*se),
      .groups = "drop"
    ) %>%
    dplyr::mutate(
      x0 = as.numeric(pc_decile),
      x  = dplyr::case_when(
        series == "Cross-party" ~ x0,
        series == "Cross-ideology (gend)" ~ x0,
        TRUE ~ x0
      )
    ) %>%
    dplyr::mutate(series = forcats::fct_drop(series))
  
  show_legend <- (s == "pool")
  plot_h      <- if (show_legend) 2.80 else 2.50   
  
  series_present <- levels(droplevels(plot_sum_gender$series))
  legend_override <- list(
    shape    = unname(atme_shapes[series_present]),
    fill     = unname(atme_pal[series_present]),
    linetype = 1
  )
  
  p_three_gender <- ggplot(plot_sum_gender,
                           aes(x = x, y = mean,
                               color = series, fill = series,
                               shape = series, group = series)) +
    geom_line(linewidth = 0.6) +
    geom_point(size = 1.6, stroke = 0.5) +
    geom_errorbar(aes(ymin = lo, ymax = hi), width = 0, linewidth = 0.6) +
    scale_x_continuous(breaks = 1:10, labels = 1:10, minor_breaks = NULL) +
    scale_y_continuous(limits = c(0, 1), expand = c(0, 0)) +
    scale_color_manual(values = atme_pal,
                       guide = guide_legend(title=NULL, nrow=1,
                                            override.aes = legend_override)) +
    scale_fill_manual(values = atme_pal, guide="none") +
    scale_shape_manual(values = atme_shapes, guide="none") +
    facet_wrap(~ pc_label, nrow = 1) +
    labs(
      title = sample_titles[s],
      x = "Political culture decile (low \u2192 high)",
      y = "Rate of choosing\npro-democratic candidate"
    ) +
    theme_bw(base_size = 11) +
    theme(
      panel.grid.minor   = element_blank(),
      legend.position    = if (show_legend) "bottom" else "none",
      legend.box.spacing = grid::unit(0.1, "lines"),
      legend.key.height  = grid::unit(0.4, "lines"),
      legend.text        = element_text(size = 11)
    )
  
  print(p_three_gender)
  
  ggsave(file.path(project_root, "outputs", paste0("bivar_prodem_gender_", s, "_2.png")),
         p_three_gender, width = 8, height = plot_h, dpi = 250)
}

