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 - Logit predicted-probability plots

library(here)
library(dplyr)
library(tidyr)
library(forcats)
library(ggplot2)
library(sandwich)
library(grid)

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

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

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

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

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

choice_overall_list <- vector("list", length(samples))
names(choice_overall_list) <- samples
choice_party_list <- vector("list", length(samples))
names(choice_party_list) <- samples
choice_ideol_redist_list <- vector("list", length(samples))
names(choice_ideol_redist_list) <- samples
choice_ideol_gender_list <- vector("list", length(samples))
names(choice_ideol_gender_list) <- samples

fd_results <- NULL


#### Loop 1: build task-level choice datasets with choose_dem outcome ----

for (s in samples) {
  cat("Preparing task-level choice datasets for", s, "...\n")

  base_df  <- wide_list[[s]]
  dat_long <- long_list[[s]]

  if (!("is_cand_dem" %in% names(dat_long))) {
    dat_long$is_cand_dem <- as.integer(dat_long$cand_system_change %in% pro_dem_labels)
  }
  dat_long$chosen <- as.integer(dat_long$chosen)

  non_cj_cols <- names(base_df)[!grepl("^CJ_Task_", names(base_df))]

  ## Overall
  choice_overall <- dat_long %>%
    group_by(RespondentID, task) %>%
    summarise(
      choose_dem = as.integer(any(chosen == 1L & is_cand_dem == 1L)),
      .groups = "drop"
    ) %>%
    left_join(base_df %>% select(all_of(non_cj_cols)), by = "RespondentID") %>%
    relocate(country, RespondentID, task, choose_dem)

  choice_overall_list[[s]] <- choice_overall

  ## Cross-party
  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 %>%
    left_join(base_df %>% select(RespondentID, all_of(eval_cols)), by = "RespondentID") %>%
    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_keep <- long2 %>%
    semi_join(keep_resp, by = "RespondentID")

  pair_flags_party <- long2_keep %>%
    group_by(RespondentID, task) %>%
    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"
    ) %>%
    filter(has_undem_in & has_dem_out) %>%
    select(RespondentID, task)

  choice_party <- long2_keep %>%
    inner_join(pair_flags_party, by = c("RespondentID", "task")) %>%
    group_by(RespondentID, task) %>%
    summarise(
      choose_dem = as.integer(any(chosen == 1L & is_cand_dem == 1L & out_party)),
      .groups = "drop"
    ) %>%
    left_join(base_df %>% select(all_of(non_cj_cols)), by = "RespondentID") %>%
    relocate(country, RespondentID, task, choose_dem)

  choice_party_list[[s]] <- choice_party

  ## Cross-ideology
  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 %>%
      left_join(base_df %>% select(RespondentID, all_of(missing_ideo)), by = "RespondentID")
  }

  long3 <- long3 %>%
    mutate(
      cand_redist_1 = case_when(
        cand_tax_spend == "spend more" ~ 4.5,
        cand_tax_spend == "spend less" ~ 1.5,
        TRUE ~ NA_real_
      ),
      cand_redist_2 = case_when(
        cand_tax_spend == "tax rich more" ~ 4.5,
        cand_tax_spend == "tax rich less" ~ 1.5,
        TRUE ~ NA_real_
      ),
      cand_gend_1 = case_when(
        cand_gend_policy == "require hire pref" ~ 4.5,
        cand_gend_policy == "no hire pref"      ~ 1.5,
        TRUE ~ NA_real_
      ),
      cand_gend_2 = 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)
    ) %>%
    rowwise() %>%
    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)
    ) %>%
    ungroup()

  pair_flags_r <- long3 %>%
    group_by(RespondentID, task) %>%
    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"
    ) %>%
    filter(has_undem_prox_r & has_dem_dist_r) %>%
    select(RespondentID, task)

  choice_ideol_redist <- long3 %>%
    inner_join(pair_flags_r, by = c("RespondentID", "task")) %>%
    group_by(RespondentID, task) %>%
    summarise(
      choose_dem = as.integer(any(chosen == 1L & is_cand_dem == 1L & dist_r)),
      .groups = "drop"
    ) %>%
    left_join(base_df %>% select(all_of(non_cj_cols)), by = "RespondentID") %>%
    relocate(country, RespondentID, task, choose_dem)

  choice_ideol_redist_list[[s]] <- choice_ideol_redist

  pair_flags_g <- long3 %>%
    group_by(RespondentID, task) %>%
    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"
    ) %>%
    filter(has_undem_prox_g & has_dem_dist_g) %>%
    select(RespondentID, task)

  choice_ideol_gender <- long3 %>%
    inner_join(pair_flags_g, by = c("RespondentID", "task")) %>%
    group_by(RespondentID, task) %>%
    summarise(
      choose_dem = as.integer(any(chosen == 1L & is_cand_dem == 1L & dist_g)),
      .groups = "drop"
    ) %>%
    left_join(base_df %>% select(all_of(non_cj_cols)), by = "RespondentID") %>%
    relocate(country, RespondentID, task, choose_dem)

  choice_ideol_gender_list[[s]] <- choice_ideol_gender
}


#### Loop 2: fit logits and build predictions ----

x_seq <- seq(-2, 2, by = 0.05)

for (s in samples) {
  cat("Building logit prediction plots for", s, "...\n")

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

  version_specs <- list(
    redist = list(
      data_list = list(
        overall = choice_overall_list[[s]],
        party   = choice_party_list[[s]],
        cross   = choice_ideol_redist_list[[s]]
      ),
      series_labels = c(
        overall = "Overall",
        party   = "Cross-party",
        cross   = "Cross-ideology (redist)"
      ),
      outfile = file.path(project_root, "outputs", paste0("logit_preds_redist_", s, ".png"))
    ),
    gender = list(
      data_list = list(
        overall = choice_overall_list[[s]],
        party   = choice_party_list[[s]],
        cross   = choice_ideol_gender_list[[s]]
      ),
      series_labels = c(
        overall = "Overall",
        party   = "Cross-party",
        cross   = "Cross-ideology (gend)"
      ),
      outfile = file.path(project_root, "outputs", paste0("logit_preds_gender_", s, ".png"))
    )
  )

  for (version_name in names(version_specs)) {
    version_spec <- version_specs[[version_name]]
    plot_df <- NULL

    for (series_name in names(version_spec$data_list)) {
      dat_choice <- version_spec$data_list[[series_name]]
      rhs_all <- paste(c(pc_vars, covars_s), collapse = " + ")
      f_all <- as.formula(paste("choose_dem ~", rhs_all))
      req_all <- unique(c(all.vars(f_all), "RespondentID"))
      rows_all <- complete.cases(dat_choice[, req_all, drop = FALSE])
      dat_model <- dat_choice[rows_all, , drop = FALSE]

      cat_cols <- intersect(c("educ_5f", "income_cat", "region_pool", "top_party"), names(dat_model))
      for (cc in cat_cols) {
        dat_model[[cc]] <- factor(dat_model[[cc]])
      }

      m_all <- glm(f_all, data = dat_model, family = binomial(link = "logit"))
      V_all <- vcovCL(m_all, cluster = dat_model$RespondentID)
      trm <- delete.response(terms(m_all))
      beta_all <- coef(m_all)
      beta_all <- beta_all[!is.na(beta_all)]
      vn_all <- colnames(V_all)

      for (v in pc_vars) {
        for (x in x_seq) {
          dat_pred <- dat_model

          for (pv in pc_vars) {
            if (pv == v) {
              dat_pred[[pv]] <- x
            } else {
              dat_pred[[pv]] <- 0
            }
          }

          Xfull <- model.matrix(trm, data = dat_pred)
          comm <- Reduce(intersect, list(colnames(Xfull), names(beta_all), vn_all))
          if (length(comm) < 2L) next

          X <- Xfull[, comm, drop = FALSE]
          beta <- beta_all[comm]
          V_use <- V_all[comm, comm, drop = FALSE]

          eta <- as.vector(X %*% beta)
          p_hat <- plogis(eta)
          avg_p <- mean(p_hat)

          grad_w <- p_hat * (1 - p_hat)
          G <- X * grad_w
          grad <- colMeans(G)
          se_p <- sqrt(as.numeric(t(grad) %*% V_use %*% grad))
          lo_p <- max(0, avg_p - 1.96 * se_p)
          hi_p <- min(1, avg_p + 1.96 * se_p)

          plot_df <- bind_rows(
            plot_df,
            tibble(
              pc_var = v,
              x = x,
              pred = avg_p,
              lo = lo_p,
              hi = hi_p,
              series = version_spec$series_labels[[series_name]]
            )
          )
        }
      }
    }

    plot_df <- plot_df %>%
      mutate(
        pc_label = factor(recode(pc_var, !!!label_map), levels = unname(label_map[pc_vars])),
        series = factor(series,
                        levels = unname(version_spec$series_labels))
      )

    fd_now <- plot_df %>%
      filter(x %in% c(-1, 1)) %>%
      mutate(x_tag = ifelse(x == -1, "pred_m1", "pred_p1")) %>%
      select(pc_var, pc_label, series, x_tag, pred) %>%
      tidyr::pivot_wider(names_from = x_tag, values_from = pred) %>%
      mutate(
        sample = s,
        version = version_name,
        outcome = as.character(series),
        first_diff = pred_p1 - pred_m1
      ) %>%
      select(sample, version, outcome, pc_var, pc_label, pred_m1, pred_p1, first_diff)

    fd_results <- bind_rows(fd_results, fd_now)

    show_legend <- (s == "pool")
    plot_h <- if (show_legend) 2.80 else 2.50

    p_out <- ggplot(plot_df,
                    aes(x = x, y = pred, color = series, fill = series, group = series)) +
      geom_ribbon(aes(ymin = lo, ymax = hi), alpha = 0.14, linewidth = 0) +
      geom_line(linewidth = 0.6) +
      scale_x_continuous(limits = c(-2, 2), breaks = c(-2, -1, 0, 1, 2), minor_breaks = NULL) +
      scale_y_continuous(limits = c(0, 1), breaks = seq(0, 1, by = 0.1), expand = c(0, 0)) +
      scale_color_manual(values = atme_pal,
                         guide = guide_legend(title = NULL, nrow = 1)) +
      scale_fill_manual(values = atme_pal, guide = "none") +
      facet_wrap(~ pc_label, nrow = 1) +
      labs(
        title = sample_titles[s],
        x = "Political culture measure (z-score)",
        y = "Predicted probability 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_out)
    ggsave(version_spec$outfile, p_out, width = 8, height = plot_h, dpi = 250)
  }
}


if (!is.null(fd_results) && nrow(fd_results) > 0) {
  fd_results <- fd_results %>%
    mutate(
      sample = factor(sample, levels = samples),
      version = factor(version, levels = c("redist", "gender")),
      pc_label = factor(pc_label, levels = unname(label_map[pc_vars]))
    ) %>%
    arrange(version, sample, outcome, pc_label)

  write.csv(fd_results, file.path(project_root, "outputs", "logit_first_differences_v2.csv"), row.names = FALSE)
}
