# Load relevant libraries

# The tidyverse is a bundle of packages that share a common design. This
# chapter draws on three of them: dplyr for reshaping data (mutate, filter,
# select), ggplot2 for every figure, and tidyr for handling missing values.
library(tidyverse)

# so the short file names below resolve correctly. This is the one line you
# must edit to run the chapter on your own machine.


# readRDS() reads a single saved R object back into memory. Each of these four
# files holds a LIST of imputed datasets -- not one table but several, which is
# why the plotting functions below all begin by stacking a list together.
#   AP_dist   Average Personalism, district level
#   TDE_dist  Total Duvergerian Effect, district level
#   AP_cty    Average Personalism, country level
#   TDE_cty   Total Duvergerian Effect, country level
# "dist" and "cty" matter: a country-level score averages over the districts
# within a system, so the two levels answer slightly different questions. Most
# figures in this chapter use the district-level data.
AP_dist <- readRDS("data/ch05/rds/district_AP_imputed.rds")
TDE_dist <- readRDS("data/ch05/rds/district_TDE_imputed.rds")
AP_cty <- readRDS("data/ch05/rds/country_AP_imputed.rds")
TDE_cty <- readRDS("data/ch05/rds/country_TDE_imputed.rds")

# FUNCTIONS

# Stack and Recode Electoral Data

# This function processes a list of datasets, applies factor recoding, selects relevant variables, and classifies electoral systems into families.

# List of parameters:
# data: a list of data frames containing electoral information.
# x: a character vector specifying the independent variable(s) to retain.
# y: a character vector specifying the dependent variable(s) to retain.
# level: a string specifying the level of aggregation

# The function returns a single stacked data frame with recoded categorical variables and a new variable classifying electoral systems into families.


data_stack <- function(data, x, y, level) {
  # ------------------------------------------------------------------------
  # The pattern here is worth learning, because it recurs throughout the book:
  #
  #   lapply(data, FUN = ...)   applies the same function to every element of
  #                             a list, returning a list of results.
  #   do.call(rbind, <list>)    then stacks those results into one data frame,
  #                             as if you had typed rbind(res1, res2, res3, ...).
  #
  # Together they mean "clean each imputed dataset the same way, then pile them
  # up". Reading the code inside out -- innermost parentheses first -- is the
  # only way to make sense of it.
  #
  # \(dta) { ... } is shorthand for function(dta) { ... }, available from
  # R 4.1 onward. It defines a small unnamed function on the spot; `dta` is
  # whichever dataset lapply() is currently handing it.
  # ------------------------------------------------------------------------
  do.call(
    rbind, # Combines the results of lapply into a single data frame
    lapply(data, # Applies function to each element of the list
      FUN = \(dta) {
        # If 'y' is "pers" and the level is "Country", use "pers_avg"
        # The country-level files store personalism already averaged
        # across districts, under a different column name. This
        # swaps in the right name so the caller does not have to.
        if (y == "pers" & level == "Country") {
          y <- "pers_avg"
        }

        # What follows is one long pipeline. Each %>% passes the
        # result of the previous step into the next, so the data are
        # transformed in sequence: recode labels, fill in missing
        # values, keep the needed columns, classify into families.
        dta %>%
          # Recode electoral function names for better readability
          # recode_factor(old = "New", ...) renames category labels.
          # Note that two inputs, stv and droop, both map to "Droop":
          # recoding can deliberately merge categories, not just
          # rename them.
          mutate(elec_fun_name = recode_factor(elec_fun_name,
            hare = "Hare",
            droop = "Droop",
            stv = "Droop",
            hagenbachbischoff = "Hagenbach-Bischoff",
            saintelague = "Sainte-Lague",
            modsaintlague = "Mod. Sainte-Lague",
            imperiali = "Imperiali",
            dhondt = "D'Hondt",
            fortified_pr = "Fortified PR",
            a_v = "Abs. Majority",
            lim_nom = "Plurality",
            plurality = "Plurality"
          )) %>%
          # Recode 'new.nvotes' variable to more interpretable labels
          mutate(new.nvotes = recode_factor(new.nvotes,
            One = "One",
            LessSeats = "Total Seats - 1",
            TotalSeats = "Total Seats",
            TotalCandidates = "Total Candidates"
          )) %>%
          # Replace NA values in 'new.nvotes' with "Total Seats - 1"
          # NA is R's marker for a missing value. replace_na()
          # substitutes a chosen value wherever one appears. Here the
          # substitutions encode a modelling decision: a system that
          # does not record this rule is treated as the stated
          # default, rather than being dropped from the figure.
          mutate(new.nvotes = replace_na(
            new.nvotes,
            "Total Seats - 1"
          )) %>%
          # Recode 'ballot_type' to reflect different voting systems
          mutate(ballot_type = recode_factor(ballot_type,
            open = "Open",
            flexible = "Flexible",
            closed = "Closed"
          )) %>%
          # Replace NA values in 'ballot_type' with "Open"
          mutate(ballot_type = replace_na(
            ballot_type,
            "Open"
          )) %>%
          # Recode 'pool_level' for different levels of vote aggregation
          mutate(pool_level = recode_factor(pool_level,
            candidate = "Candidate",
            party_list = "Sub-party List",
            party = "Party List"
          )) %>%
          # Replace NA values in 'pool_level' with "Sub-party List"
          mutate(pool_level = replace_na(
            pool_level,
            "Sub-party List"
          )) %>%
          # Recode electoral thresholds
          # The backticks around `0.01` are needed because a name
          # cannot begin with a digit in R; backticks let you use one
          # anyway. Turning the numbers into ordered labels means the
          # axis reads "0, 1%, 3%..." rather than as a continuous
          # scale.
          mutate(threshold = recode_factor(threshold,
            `0` = "0",
            `0.01` = "1%",
            `0.03` = "3%",
            `0.05` = "5%",
            `0.1` = "10%"
          )) %>%
          # Select only the relevant variables specified by the user
          # x and y are whatever the caller asked for, so this line
          # is what makes one function serve every figure.
          # all_of() insists that the named columns exist, turning a
          # silent mismatch into an immediate, informative error.
          dplyr::select(all_of(c(x, y, "system_name"))) %>%
          # Classify electoral systems into broad categories
          # case_when() is a multi-way if/else. Each line reads
          # "condition ~ value", checked from top to bottom; the
          # first match wins. %in% asks whether a value appears
          # anywhere in the list that follows it.
          # .default catches everything not matched above, so any
          # system that is neither PR nor mixed becomes Majoritarian.
          mutate(
            family = case_when(
              system_name %in% c(
                "FortifiedPR",
                "CLPR",
                "FrLPR",
                "OLPR",
                "STV",
                "FLPR"
              ) ~ "PR", # PR Systems
              system_name %in% c("MMC", "MMI") ~ "Mixed", # Mixed systems
              .default = "Majoritarian"
            ), # Majoritarian Systems
            # Record which level of aggregation this table came
            # from, so the two can be told apart once stacked.
            level = level
          ) %>%
          # Drop 'system_name' as it's no longer needed after classification
          # A minus sign inside select() removes a column instead of
          # keeping it.
          dplyr::select(-system_name)
      }
    )
  )
}


# Generate Electoral System Comparison Plots
# This function creates a comparative visualization of electoral system features and their effects on a specified outcome.

# List of parameters:
# data_cty: a list of country-level datasets to be processed and stacked.
# data_dist: a list of district-level datasets to be processed and stacked.
# rule: a string specifying the independent variable
# outcome: a string specifying the dependent variable to be plotted
# xlab: a string for the x-axis label.
# ylab: a string for the y-axis label.
# print: Logical; if TRUE (default), the plot is printed.
# file: for saving the plot).
# type: a string indicating whether the independent variable is "ordered" (categorical) or "continuous".

# The function returns a ggplot object visualizing the relationship between the rule and outcome across electoral system families.


plot_one <- function(data_cty,
                     data_dist,
                     rule,
                     outcome,
                     xlab, ylab,
                     print = TRUE,
                     file,
                     type = c("ordered", "continuous")) {
  # Stack country-level and district-level datasets
  all_data_cty <- data_stack(data_cty, rule, outcome, "Country")
  all_data_dist <- data_stack(data_dist, rule, outcome, "District")

  # Ensure consistent column names across both datasets
  names(all_data_cty) <- names(all_data_dist)

  # Combine both datasets into a single data frame
  all_data <- rbind(all_data_cty, all_data_dist)


  # Apply data filtering conditions based on the selected rule
  if (rule == "M") {
    all_data <- all_data %>%
      dplyr::filter(!((family == "Majoritarian") & (M > 43))) # Exclude majoritarian systems with M > 43
  }
  if (rule == "new.nvotes") {
    all_data <- all_data %>%
      dplyr::filter(!(family == "Mixed")) # Exclude mixed systems for this rule
  }
  if (rule == "ballot_type") {
    all_data <- all_data %>%
      dplyr::filter(!(family == "Mixed")) # Exclude mixed systems for this rule
  }
  if (rule == "threshold") {
    all_data <- all_data %>%
      dplyr::filter(!(family == "Majoritarian")) # Exclude majoritarian systems for this rule
  }


  # Initialize ggplot object with rule on x-axis and outcome on y-axis
  # -----------------------------------------------------------------------
  # HOW ggplot2 WORKS
  # A ggplot is assembled by adding layers with +, much as base R graphics
  # are built up by successive commands. The pieces here:
  #   ggplot(aes(...))  declares the "aesthetic mapping": which column drives
  #                     the x position, which the y, which the colour.
  #   .data[[rule]]     is the pronoun trick that makes these functions
  #                     general. `rule` holds a column NAME as text, e.g.
  #                     "M"; .data[[rule]] looks up that column at run time.
  #                     Writing plain `rule` would plot the literal word.
  #   scale_color_grey  greyscale, since the book prints in black and white.
  #   facet_wrap(...)   splits the plot into one panel per electoral family.
  #                     scales = "free_x" lets each panel set its own x-axis,
  #                     which matters because majoritarian and PR systems do
  #                     not share the same range of magnitudes or formulas.
  #   labs / theme_bw   axis labels and a plain black-and-white theme.
  # Note that the object is only BUILT here, not drawn; it is stored in `plt`
  # and drawn further down by print().
  # -----------------------------------------------------------------------
  plt <- all_data %>%
    ggplot(aes(x = .data[[rule]], y = .data[[outcome]], color = level)) +
    scale_color_grey(start = 0.7, end = 0) + # Use grayscale colors
    facet_wrap(vars(family), scales = "free_x") + # Create facet grid by electoral family
    labs(x = xlab, y = ylab, color = "Level") + # Set axis labels and legend title
    theme_bw()

  # Add visualization type based on 'type' argument
  # The right display depends on the kind of rule being plotted: a trend line
  # for a numeric rule such as district magnitude, a box plot for a categorical
  # one such as ballot type. Adding a layer to an existing plot object is done
  # with plt <- plt + <layer>.
  if (type == "continuous") {
    plt <- plt +
      geom_smooth(se = TRUE) # Add a smoothed line with confidence intervals for continuous variables
  } else {
    plt <- plt +
      # scale_x_discrete(guide = guide_axis(n.dodge = 3)) +
      theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1)) + # Rotate x-axis labels for readability
      geom_boxplot() # Use box plots for categorical independent variables
  }
  # A ggplot object must be printed to appear. Inside a function this has to be
  # explicit, because only the value of the LAST line is auto-printed.
  if (print) {
    print(plt) # Print the plot if 'print' argument is TRUE
  }
}


# Generate Interaction Plots for Electoral Systems

# This function creates a visualization to examine the interaction between two electoral system characteristics and their effects on a specified outcome.

# List of parameters:
# data_cty: a list of country-level datasets to be processed and stacked.
# data_dist: a list of district-level datasets to be processed and stacked.
# rule: a character vector of length 2 specifying the two independent variables for interaction.
# outcome: a string specifying the dependent variable to be plotted.
# xlab: a string for the x-axis label.
# ylab: a string for the y-axis label.
# llab: a string for the legend label.
# print: Logical; if TRUE (default), the plot is printed.
# file: (Unused in the current version, but could be implemented for saving the plot).
# type: a string indicating whether the independent variable is "ordered" (categorical) or "continuous".

# The function returns a ggplot object visualizing the interaction between two rules and the electoral outcome across different electoral system families.


plot_interact <- function(data_cty,
                          data_dist,
                          rule,
                          outcome,
                          xlab, ylab, llab,
                          print = TRUE,
                          file,
                          type = c("ordered", "continuous")) {
  # Stack country-level and district-level datasets
  all_data_cty <- data_stack(data_cty, rule, outcome, "Country")
  all_data_dist <- data_stack(data_dist, rule, outcome, "District")

  # Ensure column names match between both datasets
  names(all_data_cty) <- names(all_data_dist)

  # Combine both datasets into a single data frame
  all_data <- rbind(all_data_cty, all_data_dist)

  # Filter the data to include only relevant observations
  plt <- all_data %>%
    dplyr::filter(
      level == "District", # Consider only district-level data
      ballot_type %in% c("Closed", "Open"), # Include only these ballot types
      !((family == "Majoritarian") & (M > 43))
    ) %>% # Exclude majoritarian systems with M > 43
    ggplot(aes(x = .data[[rule[1]]], y = .data[[outcome]], color = .data[[rule[2]]])) +
    scale_color_grey(start = 0.7, end = 0) + # Use grayscale color
    facet_wrap(vars(family), scales = "free_x") + # Facet plot by electoral family
    labs(x = xlab, y = ylab, color = llab) + # Add axis labels and legend title
    theme_bw()

  # Add visualization type based on 'type' argument
  if (type == "continuous") {
    plt <- plt +
      geom_smooth(se = TRUE) # Add a smoothed trend line with confidence intervals when the variable is continuous
  } else {
    plt <- plt +
      # scale_x_discrete(guide = guide_axis(n.dodge = 3)) +
      theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1)) + # Rotate x-axis labels for readability
      geom_boxplot() # Use box plots for categorical  variables
  }

  # Print the plot if 'print' argument is TRUE
  if (print) {
    print(plt)
  }
}


# Generate District-Level Electoral System Plots
# This function creates a visualization to examine the relationship between an electoral rule and an outcome at the district level.

# data_dist: a list of district-level datasets to be processed and stacked.
# rule: a string specifying the independent variable (e.g., "M", "new.nvotes", "ballot_type", "threshold").
# outcome: a string specifying the dependent variable to be plotted.
# xlab: a string for the x-axis label.
# ylab: a string for the y-axis label.
# print: Logical; if TRUE (default), the plot is printed.
# file: could be implemented for saving the plot.
# type: a string indicating whether the independent variable is "ordered" (categorical) or "continuous".

# The function returns a ggplot object visualizing the relationship between the rule and outcome across electoral system families at the district level.


# This is the function that produces Figures 5.1 through 5.6. It is the simplest
# of the four: one rule, one outcome, district level only.
plot_district <- function(data_dist,
                          rule,
                          outcome,
                          xlab, ylab,
                          print = TRUE,
                          file,
                          type = c("ordered", "continuous")) {
  # Stack district-level data
  all_data <- data_stack(data_dist, rule, outcome, "District")

  # Apply data filtering conditions based on the selected rule
  # Each rule needs a different exclusion, because some combinations are not
  # meaningful: a majoritarian system with 43+ seats per district is a
  # simulation artefact, and vote rules or ballot types are not well defined
  # for mixed systems. The ! operator negates a condition, so filter(!(...))
  # reads "keep everything EXCEPT those cases".
  if (rule == "M") {
    all_data <- all_data %>%
      dplyr::filter(!((family == "Majoritarian") & (M > 43))) # Exclude majoritarian systems with M > 43
  }
  if (rule == "new.nvotes") {
    all_data <- all_data %>%
      dplyr::filter(!(family == "Mixed")) # Exclude mixed systems for this rule
  }
  if (rule == "ballot_type") {
    all_data <- all_data %>%
      dplyr::filter(!(family == "Mixed")) # Exclude mixed systems for this rule
  }
  if (rule == "threshold") {
    all_data <- all_data %>%
      dplyr::filter(!(family == "Majoritarian")) # Exclude majoritarian systems for this rule
  }

  # Initialize ggplot object with rule on x-axis and outcome on y-axis
  plt <- all_data %>%
    ggplot(aes(x = .data[[rule]], y = .data[[outcome]])) +
    scale_color_grey(start = 0.7, end = 0) + # Use grayscale color scale
    facet_wrap(vars(family), scales = "free_x") + # Facet plot by electoral family
    labs(x = xlab, y = ylab) + # Add axis labels
    theme_bw()

  # Add visualization type based on 'type' argument
  if (type == "continuous") {
    plt <- plt +
      geom_smooth(se = TRUE, color = "black") # Add a smoothed trend line with confidence intervals for continuous  variables
  } else {
    plt <- plt +
      # scale_x_discrete(guide = guide_axis(n.dodge = 3)) +
      theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1)) + # Rotate x-axis labels for readability
      geom_boxplot() # Use box plots for categorical variables
  }

  # Print the plot if 'print' argument is TRUE
  if (print) {
    print(plt)
  }
}


# Generate Interaction Plots for Electoral Systems
# This function creates a visualization to examine the interaction between two electoral system characteristics and their effects on a specified outcome.

# data_cty: a list of country-level datasets to be processed and stacked.
# data_dist: a list of district-level datasets to be processed and stacked.
# rule: a character vector of length 2 specifying the two independent variables for interaction.
# outcome: a string specifying the dependent variable to be plotted.
# xlab: a string for the x-axis label.
# ylab: a string for the y-axis label.
# llab: a string for the legend label.
# print: Logical; if TRUE (default), the plot is printed.
# file: could be implemented for saving the plot.
# type: a string indicating whether the independent variable is "ordered" (categorical) or "continuous".

# The function returns a ggplot object visualizing the interaction between two rules and the electoral outcome across different electoral system families.

# NOTE: this is a second definition of plot_interact(), identical to the one
# above. Re-defining a function simply overwrites the earlier version, so the
# behaviour of Figure 5.7 is unaffected -- the two are the same recipe. The
# duplicate is left in place because this file reproduces the code as run.
plot_interact <- function(data_cty,
                          data_dist,
                          rule,
                          outcome,
                          xlab, ylab, llab,
                          print = TRUE,
                          file,
                          type = c("ordered", "continuous")) {
  # Stack country-level and district-level datasets
  all_data_cty <- data_stack(data_cty, rule, outcome, "Country")
  all_data_dist <- data_stack(data_dist, rule, outcome, "District")

  # Ensure column names match between both datasets
  names(all_data_cty) <- names(all_data_dist)

  # Combine both datasets into a single data frame
  all_data <- rbind(all_data_cty, all_data_dist)

  # Filter the data to include only relevant observations
  plt <- all_data %>%
    dplyr::filter(
      level == "District", # Consider only district-level data
      ballot_type %in% c("Closed", "Open"), # Include only these ballot types
      !((family == "Majoritarian") & (M > 43))
    ) %>% # Exclude majoritarian systems with M > 43
    ggplot(aes(x = .data[[rule[1]]], y = .data[[outcome]], color = .data[[rule[2]]])) +
    scale_color_grey(start = 0.7, end = 0) + # Use grayscale color scale
    facet_wrap(vars(family), scales = "free_x") + # Facet plot by electoral family
    labs(x = xlab, y = ylab, color = llab) + # Add axis labels and legend title
    theme_bw()

  # Add visualization type based on 'type' argument
  if (type == "continuous") {
    plt <- plt +
      geom_smooth(se = TRUE) # Add a smoothed trend line with confidence intervals for continuous variables
  } else {
    plt <- plt +
      # scale_x_discrete(guide = guide_axis(n.dodge = 3)) +
      theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1)) + # Rotate x-axis labels for readability
      geom_boxplot() # Use box plots for categorical variables
  }

  # Print the plot if 'print' argument is TRUE
  if (print) {
    print(plt)
  }
}

# Reading the call: use the district-level TDE data; put district magnitude (M)
# on the x-axis and the TDE score on the y-axis; label the axes; and because M
# is a number rather than a category, draw a smoothed trend line
# (type = "continuous") rather than box plots.
# `file` is accepted by the function but never used -- the figure is displayed
# inline rather than written to disk.
plot_district(TDE_dist, "M", "TDE_enp", "Magnitude", "Total Duvergerian Effect",
  file = "Figures/M_TDE_dist.pdf", type = "continuous"
)

# Same function, different rule. Because the seat allocation formula is a
# category rather than a number, type = "ordered" switches the display from a
# trend line to box plots -- one box per formula.
plot_district(TDE_dist, "elec_fun_name", "TDE_enp", "Seat Allocation Formula", "Total Duvergerian Effect",
  file = "Figures/Formula_TDE_dist.pdf", type = "ordered"
)

# Same function and the same display settings as Figures 5.2 and 5.4, so the
# three figures are directly comparable: district-level data, a categorical rule
# on the x-axis (type = "ordered", hence box plots), TDE on the y-axis.
#
# Two things happen automatically here, both set up much earlier in the file:
#   1. plot_district() carries a rule-specific filter. For "threshold" it drops
#      the Majoritarian family, since systems that award a single seat by
#      plurality have no threshold to speak of. That is why the figure resolves
#      into one panel rather than three.
#   2. data_stack() recoded the raw threshold values into ordered labels
#      (0, 1%, 3%, 5%, 10%). Because recode_factor() fixes the level order, the
#      boxes appear in ascending order along the axis rather than alphabetically.
plot_district(TDE_dist, "threshold", "TDE_enp", "Electoral Threshold", "Total Duvergerian Effect",
  file = "Figures/Threshold_TDE_dist.pdf", type = "ordered"
)

# The outcome switches from TDE to AP here, so the data argument changes from
# TDE_dist to AP_dist and the outcome column from "TDE_enp" to "pers".
plot_district(AP_dist, "ballot_type", "pers", "Ballot Type", "Personalism",
  file = "Figures/Ballot_pers_dist.pdf", type = "ordered"
)

# "new.nvotes" is the recoded number-of-votes rule, whose labels were set in
# data_stack() above: One, Total Seats - 1, Total Seats, Total Candidates.
plot_district(AP_dist, "new.nvotes", "pers", "Nr. of Votes", "Personalism",
  file = "Figures/NVotes_pers_dist.pdf", type = "ordered"
)

# "pool_level" was recoded in data_stack() to Candidate, Sub-party List and
# Party List -- ordered from the narrowest pool to the broadest.
plot_district(AP_dist, "pool_level", "pers", "Vote Pooling Level", "Personalism",
  file = "Figures/Pool_pers_dist.pdf", type = "ordered"
)

# This is the only figure using plot_interact(), and the only one that needs
# BOTH datasets, country and district. The key argument is c("M","ballot_type"):
# the first name goes on the x-axis, the second becomes the line colour. That is
# what turns one trend line into several -- one per ballot type -- so the two
# rules can be seen acting together.
plot_interact(AP_cty, AP_dist,
  c("M", "ballot_type"), "pers", "Magnitude", "Personalism", "Ballot type",
  file = "Figures/M_ballot_type_pers.pdf", type = "continuous"
)