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 - conjoint coding and analysis

library(here)
library(dplyr)
library(forcats)
library(tibble)
library(ggplot2)
library(cregg)

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


#### Conjoint: reshape wide -> long for each dataset---

id <- "RespondentID"
tasks <- 12L
atts <- 7L

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

for (nm in dataset_names) {

  df <- get(nm)

  long_list <- vector("list", length = tasks * 2L)
  idx <- 1L

  for (tk in seq_len(tasks)) {
    choice_var <- sprintf("CJ_Task_%d_choice", tk)

    # Candidate 1
    c1_vars <- c(id, choice_var,
                 sprintf("CJ_Task_%d_Cand_1_Att%d", tk, seq_len(atts)))
    c1 <- df[, c1_vars, drop = FALSE]
    names(c1)[-(1:2)] <- paste0("Att", seq_len(atts))
    c1$task <- tk
    c1$cand <- 1L
    c1$chosen <- as.integer(c1[[choice_var]] == 1L)
    c1 <- c1[, c(id, "task", "cand", "chosen", paste0("Att", 1:atts)), drop = FALSE]
    long_list[[idx]] <- c1
    idx <- idx + 1L

    # Candidate 2
    c2_vars <- c(id, choice_var,
                 sprintf("CJ_Task_%d_Cand_2_Att%d", tk, seq_len(atts)))
    c2 <- df[, c2_vars, drop = FALSE]
    names(c2)[-(1:2)] <- paste0("Att", seq_len(atts))
    c2$task <- tk
    c2$cand <- 2L
    c2$chosen <- as.integer(c2[[choice_var]] == 2L)
    c2 <- c2[, c(id, "task", "cand", "chosen", paste0("Att", 1:atts)), drop = FALSE]
    long_list[[idx]] <- c2
    idx <- idx + 1L
  }

  df_long <- dplyr::bind_rows(long_list)
  df_long <- dplyr::arrange(df_long, RespondentID, task, cand)

  # Add country
  df_long <- dplyr::left_join(
    df_long,
    dplyr::select(df, RespondentID, country),
    by = "RespondentID"
  )

  # Attribute names
  df_long <- dplyr::rename(
    df_long,
    cand_age = Att1,
    cand_gender = Att2,
    cand_background = Att3,
    cand_party = Att4,
    cand_tax_spend = Att5,
    cand_gend_policy = Att6,
    cand_system_change = Att7
  )

  # Attribute levels
  df_long <- dplyr::mutate(
    df_long,
    cand_age = factor(cand_age, levels = 1:6,
                      labels = c("35", "44", "51", "59", "64", "68")),
    cand_gender = factor(cand_gender, levels = 1:2,
                         labels = c("male", "female")),
    cand_background = factor(cand_background, levels = 1:7,
                             labels = c("lawyer", "politician", "journalist",
                                        "founder", "military", "athlete", "scientist")),
    cand_tax_spend = factor(cand_tax_spend, levels = 1:4,
                            labels = c("spend more", "spend less",
                                       "tax rich more", "tax rich less")),
    cand_gend_policy = factor(cand_gend_policy, levels = 1:4,
                              labels = c("require hire pref", "no hire pref",
                                         "require quotas", "no quotas")),
    cand_system_change = factor(cand_system_change, levels = 1:6,
                                labels = c("cut polling stations", "keep polling stations",
                                           "ignore court rulings", "respect court rulings",
                                           "prosecute journalists", "don't prosecute journalists"))
  )

  # Party mapping (country x party code -> label)
  party_map <- tibble::tribble(
    ~country, ~cand_party, ~party_label,
    "US", 1L, "Democrat",
    "US", 2L, "Republican",
    "KR", 1L, "DPK",
    "KR", 2L, "PPP",
    "PL", 1L, "PiS",
    "PL", 2L, "KO",
    "PL", 3L, "Konfederacja",
    "PL", 4L, "Lewica"
  )

  if (is.character(df_long$cand_party)) {
    df_long$cand_party <- as.integer(df_long$cand_party)
  }

  df_long <- dplyr::left_join(df_long, party_map, by = c("country", "cand_party"))
  df_long <- dplyr::mutate(
    df_long,
    cand_party = factor(
      party_label,
      levels = c("Democrat", "Republican", "DPK", "PPP",
                 "PiS", "KO", "Konfederacja", "Lewica")
    )
  )
  df_long <- dplyr::select(df_long, -party_label)

  assign(paste0(nm, "_long"), df_long)
}

# Quick checks
table(pool_long$country, useNA = "ifany")
table(pool_long$country, pool_long$cand_party, useNA = "ifany")


#### Basic conjoint analysis: marginal means & plots----

attr_labs <- c(
  cand_age = "Age",
  cand_gender = "Gender",
  cand_background = "Background",
  cand_party = "Party",
  cand_tax_spend = "Redist policy",
  cand_gend_policy = "Gender policy",
  cand_system_change = "System change"
)

feature_order <- c("Age", "Gender", "Background", "Party", "Redist policy",
                   "Gender policy", "System change")

wrapped_labs <- c(
  "Age" = "Age",
  "Gender" = "Gen",
  "Background" = "Background",
  "Party" = "Party",
  "Redist policy" = "Redist\npolicy",
  "Gender policy" = "Gender\npolicy",
  "System change" = "System\nchange"
)

mm_results <- list()
mm_plots <- list()

for (nm in dataset_names) {

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

  # Marginal means
  mm_res <- mm(
    chosen ~ cand_age + cand_gender + cand_background + cand_party +
      cand_tax_spend + cand_gend_policy + cand_system_change,
    data = df_long,
    id = ~ RespondentID
  )

  mm_plot <- dplyr::mutate(
    mm_res,
    feature = dplyr::recode(feature, !!!attr_labs),
    feature = factor(feature, levels = feature_order),
    feature_lab = factor(feature,
                         levels = feature_order,
                         labels = wrapped_labs[feature_order]),
    level = fct_rev(fct_inorder(level))
  )

  p <- ggplot(mm_plot, aes(x = estimate, y = level)) +
    geom_errorbarh(aes(xmin = lower, xmax = upper), height = 0) +
    geom_point() +
    facet_grid(feature_lab ~ ., scales = "free_y",
               space = "free_y", switch = "y") +
    labs(x = "Probability of being chosen", y = NULL) +
    theme_minimal(base_size = 12) +
    theme(
      legend.position = "none",
      strip.placement = "outside",
      strip.text.y.right = element_text(face = "bold", angle = 0),
      panel.spacing.y = grid::unit(0.8, "lines")
    )

  # Save plot
  ggsave(file.path(project_root, "outputs", sprintf("conjoint_mm_%s.png", nm)), p,
         width = 7, height = 5.5, dpi = 250)

  mm_results[[nm]] <- mm_res
  mm_plots[[nm]] <- p
}

# Print plots in the session
mm_plots$pool
mm_plots$us
mm_plots$sk
mm_plots$pl


#### Save long-format datasets----

write.csv(pool_long, file.path(project_root, "data", "pool_long_v2.csv"), row.names = FALSE)
write.csv(us_long, file.path(project_root, "data", "us_long_v2.csv"), row.names = FALSE)
write.csv(sk_long, file.path(project_root, "data", "sk_long_v2.csv"), row.names = FALSE)
write.csv(pl_long, file.path(project_root, "data", "pl_long_v2.csv"), row.names = FALSE)
