# Load relevant libraries
library(igraph) # network construction and manipulation
library(Matrix) # sparse matrices -- retweet networks are mostly zeros
library(sna) # social network analysis routines
library(intergraph) # converts between igraph and sna/network formats
library(furrr) # parallel versions of purrr's map functions
library(progressr) # progress bars for the long-running network fits
library(RColorBrewer) # colour palettes for the network figures
library(gbm) # gradient boosting machines: supply the AP and TDE predictions
library(amen) # Additive and Multiplicative Effects models for networks --
# the estimator behind the disunity scores
library(texreg) # formats regression results into tables
library(ggraph) # grammar-of-graphics plotting for networks
library(lme4) # multilevel / mixed-effects models
library(broom.mixed) # tidies mixed-model output for tables
library(mice) # multiple imputation; pools across the five datasets
library(tidyverse) # bundle of data-handling packages; loaded last on purpose

# No setwd() is needed anywhere in this chapter: paths are relative to the
# repository root, which is where this document sits.

# Load retweet data: each row is a retweet from one MP to another
full_df <- read.csv(
  file = "data/ch13/csv/rtnet_sn.csv",
  colClasses = c("src_screen_name" = "character", "rt_screen_name" = "character")
)


# Load metadata on legislators (social media identifiers and other info)
meta <- as_tibble(read.csv("data/ch13/csv/metaFull.csv", colClasses = c("uid" = "character")))

# Load party ID data and keep only member_id and party id (used for matching)
party_ln_data <- as_tibble(read.csv("data/ch13/csv/members_2021_04_28.csv", colClasses = c("uid" = "character"))) %>%
  dplyr::select(member_id, mp_party_id) %>%
  distinct()

# Merge metadata with party ID info and filter out legislators without known party
meta <- meta %>%
  left_join(party_ln_data, multiple = "first") %>%
  dplyr::filter(!is.na(mp_party_id)) %>%
  # Add manifesto and seat share data (MPDataset)
  left_join(
    read_csv("data/ch13/csv/MPDataset_MPDS2024a.csv") %>%
      mutate(elecyear = as.numeric(substr(date, 1, 4))) %>%
      mutate(perseat = absseat / totseats) %>%
      dplyr::select(mp_party_id = party, rile, totseats, perseat),
    multiple = "last"
  ) %>%
  # Add party-level government and family information (PPEG data)
  left_join(
    read_csv("data/ch13/csv/ppeg_comb_2024v1.csv") %>%
      dplyr::select(mp_party_id = cmp, gov_party, parfam = cmp_parfam) %>%
      mutate(gov_party = if_else(gov_party == "yes", 1, 0)),
    multiple = "last"
  )


# Load contextual country-level variables (e.g., democracy, presidentialism)
load("data/ch13/RData/merged_data.RData")
# Keep only the latest entry per country (to avoid duplicates)
bfc_mag <- merged_files %>%
  dplyr::select(country, start_year, democracy, regionalism, presidentialism) %>%
  mutate(max_year = max(start_year), .by = start_year) %>%
  dplyr::filter(start_year == max_year) %>%
  dplyr::select(country, start_year, democracy, regionalism, presidentialism) %>%
  distinct()

# Remove large object to save memory
rm(list = c("merged_files"))

# Convert full_df to tibble and merge country info for sender and retweeted MP
full_df <- as_tibble(full_df) %>%
  left_join(dplyr::select(meta, c(src_screen_name = uid, ctry_sender = country)), relationship = "many-to-many") %>%
  left_join(dplyr::select(meta, c(rt_screen_name = uid, ctry_rt = country)), relationship = "many-to-many") %>%
  # Keep only retweets where both sender and receiver are matched to countries
  dplyr::filter(!is.na(ctry_sender) & !is.na(ctry_rt)) %>%
  # Group by MP dyads and count total retweets (edge weight), keeping country info
  group_by(src_screen_name, rt_screen_name) %>%
  summarize(
    edge = n(),
    ctry_sender = first(ctry_sender), ctry_rt = first(ctry_rt)
  ) %>%
  # Keep only retweets that occur within the same country
  dplyr::filter(ctry_sender == ctry_rt) %>%
  # Remove self-retweets
  dplyr::filter(src_screen_name != rt_screen_name)


# Filter and retain only metadata for MPs who are involved in at least one retweet (sender or target)
meta <- meta %>%
  dplyr::select(vertex = uid, party, country, member_id, magnitude, rile, parfam, perseat, totseats, gov_party) %>%
  dplyr::filter((vertex %in% full_df$src_screen_name) | (vertex %in% full_df$rt_screen_name)) %>%
  distinct(vertex, .keep_all = TRUE)

## Function to analyze retweet networks for a given country

# Defines a recipe applied separately to each country: build that country's
# retweet network, fit the AME model to it, and return one row per legislator.
analyze_nets <- function(cty, metadata, rt_df) {
  # Subset the retweet data to the given country
  cty_df <- rt_df %>% dplyr::filter(ctry_sender == cty)

  # Subset metadata for MPs in the given country
  cty_meta <- metadata %>% dplyr::filter(country == cty)

  # Create an igraph object from country-specific retweet data
  rt_net <- graph_from_data_frame(
    cty_df %>% dplyr::select(src_screen_name, rt_screen_name, weight = edge, Country = ctry_sender) %>%
      mutate(weight = log(weight)), # Log-transform the retweet frequency
    vertices = cty_meta %>% dplyr::select(vertex, Party = party)
  ) # Attach party info to each node


  # Create a sociomatrix (adjacency matrix) with edge weights
  # A sociomatrix is a square table with one row and one column per legislator;
  # the cell at [i, j] records how much i retweeted j. This is the format the
  # ame() function expects, as opposed to the edge list used above.
  socmat <- as_adjacency_matrix(rt_net, sparse = FALSE, attr = "weight")

  # Fit a 1-dimensional Additive and Multiplicative Effects (AME) model
  net_model <- ame(socmat,
    family = "nrm", # Normal response model (appropriate for weighted networks)
    symmetric = FALSE, # Retweet networks are directional
    R = 1, # One latent dimension (for position estimation)
    burn = 2000, # Burn-in iterations
    nscan = 500, # Sampling iterations
    odens = 1, # Output every iteration
    print = FALSE, # Suppress console printing
    plot = FALSE
  ) # Suppress plotting

  # Build dataframe with estimated latent positions and covariates
  df <- data.frame(
    pos = net_model$U, # Estimated latent positions (unity dimension)
    id = cty_meta$vertex, # MP ID
    party = cty_meta$party, # Party ID
    magnitude = cty_meta$magnitude, # District magnitude
    parfam = cty_meta$parfam, # Party family
    party_ideo = cty_meta$rile, # Party ideology (RILE index)
    perseat = cty_meta$perseat, # % of seats held
    totseats = cty_meta$totseats, # Total seats
    gov_party = cty_meta$gov_party, # Government party indicator
    country = cty_meta$country
  ) # Country name

  # Compute party-level average latent positions
  df <- df %>%
    group_by(party) %>%
    mutate(party_mean = mean(pos)) %>%
    ungroup() %>%
    # Calculate normalized squared distance from party average (disunity score)
    mutate(dist = ((pos - party_mean)^2) / sd(pos)) %>%
    # Keep relevant variables with clear names
    dplyr::select(id, country, party, dist, pos, party_mean,
      M = magnitude, parfam, party_ideo, perseat,
      totseats, gov_party
    )
  # Save the resulting dataframe to file
  saveRDS(df, file = paste0("data/", cty, ".rds"))
}

## Analyze networks for all countries in the dataset
countries <- unique(full_df$ctry_sender)

# Use parallel processing for faster execution
plan(multisession, workers = 8)

# Progress-aware future map to iterate over all countries
with_progress({
  p <- progressor(steps = length(countries)) # Initialize progress bar
  future_map(countries,
    ~ {
      analyze_nets(.x, metadata = meta, rt_df = full_df)
      p()
    } # Run function and update progress
    ,
    .options = furrr_options(seed = 831213)
  )
})

# Load and merge disunity data from all countries

all_data_orig <- list.files(path = "data/ch13/rds/", pattern = "*.rds", full.names = TRUE) %>%
  map_dfr(readRDS) # Read and bind all into a single dataframe

# Merge with system-level variables (democracy, presidentialism, etc.)
all_data <- all_data_orig %>%
  inner_join(bfc_mag, multiple = "last") %>%
  dplyr::select(
    id, country, party, dist, M, democracy, regionalism, presidentialism,
    party_ideo, parfam, perseat, gov_party, totseats, pos, party_mean
  )

# Merge with electoral system component rules
realCases <- read.csv("data/ch13/csv/realCases.csv") %>%
  dplyr::select(country, year, formula, # Electoral formula
    ballot_type, new.nvotes, # Ballot structure and number of votes
    pool_level,
    threshold = threshold_tier1 # Pooling level and electoral threshold
  ) %>%
  # Harmonize country name for the U.S.
  mutate(country = if_else(country == "United States of America",
    "United States", country
  )) %>%
  # Keep only countries in disunity data
  dplyr::filter(country %in% all_data$country) %>%
  # Keep only the most recent election year for each country
  mutate(filter = if_else(year == max(year), 1, 0), .by = country) %>%
  dplyr::filter(filter == 1) %>%
  dplyr::select(-c(filter, year)) %>%
  # Normalize thresholds: treat any positive threshold as 5%, else 0%
  mutate(threshold = if_else(threshold > 0, 0.05, 0.0)) %>%
  distinct() %>%
  # Remove duplicate Greece case with inconsistent ballot_type = "open"
  dplyr::filter(!(country == "Greece" & ballot_type == "open"))


# Fix known inconsistencies manually
realCases <- realCases %>%
  mutate(formula = case_match(country,
    c("Greece", "Italy") ~ "hare", # coded as using Hare formula
    .default = formula
  )) %>%
  mutate(threshold = case_match(country,
    "Austria" ~ 0.05, # Austria threshold is 5%
    "Greece" ~ 0.03, # Greece threshold is 3%
    .default = threshold
  ))

# Merge electoral rules with full legislator-level dataset
all_data <- right_join(realCases, all_data, by = "country")

# Drop rows with missing M (district magnitude) and finalize variable names
all_data <- all_data %>%
  tidyr::drop_na(M) %>%
  rename(Distance = dist) %>% # Rename disunity score to "Distance"
  distinct(id, .keep_all = TRUE) # Remove duplicates by legislator ID

# Manually correct component rule values for single-member districts in Germany and NZ
all_data$new.nvotes[all_data$M == 1 & all_data$country == "New Zealand" | all_data$M == 1 & all_data$country == "Germany"] <- "One"
all_data$ballot_type[all_data$M == 1 & all_data$country == "New Zealand" | all_data$M == 1 & all_data$country == "Germany"] <- "closed"
all_data$formula[all_data$M == 1 & all_data$country == "New Zealand" | all_data$M == 1 & all_data$country == "Germany"] <- "plurality"

## Predict AP and TDE scores

# Load GBM objects trained at the district level for TDE and AP
load(file = "data/shared/TDE_district_objects_t305_d15.RData")
load(file = "data/shared/AP_district_objects_t515_d7.RData")

# Assign to new variable names for clarity and to avoid overwriting other objects
totalAP.objects.district <- totalAP.objects
totalEffENP.objects.district <- totalEffENP.objects

# Remove original objects to free up memory and avoid confusion
rm(totalAP.objects, totalEffENP.objects)

# Convert input variables to factors with correct levels based on GBM training
all_data$new.nvotes <- factor(all_data$new.nvotes, levels = totalEffENP.objects.district$optimalGBM[[1]]$var.levels[[2]])
all_data$pool_level <- factor(all_data$pool_level, levels = totalEffENP.objects.district$optimalGBM[[1]]$var.levels[[3]])
all_data$ballot_type <- factor(all_data$ballot_type, levels = totalEffENP.objects.district$optimalGBM[[1]]$var.levels[[4]])
all_data$formula <- factor(all_data$formula, levels = totalEffENP.objects.district$optimalGBM[[1]]$var.levels[[6]])

# Initialize empty lists to store predicted values
## Using the district-level GBM object ##
totalEff.hat <- pers.hat <- list()

# Do we need pers.hat <- list() ?


# Predict TDE and AP for each of the 5 imputed models (ensembling for robustness)
for (i in 1:5) {
  optimalGBMIntra <- totalAP.objects.district$optimalGBM[[i]]
  optimalGBM <- totalEffENP.objects.district$optimalGBM[[i]]

  # Predict AP (intraparty personalism)
  all_data[[paste0("AP", i)]] <- predict(optimalGBMIntra, all_data,
    n.trees = optimalGBMIntra$n.trees
  )
  pers.hat[[i]] <- all_data[[paste0("AP", i)]]

  # Predict TDE (interparty fragmentation)
  all_data[[paste0("TDE", i)]] <- predict(optimalGBM, all_data,
    n.trees = optimalGBM$n.trees
  )
  totalEff.hat[[i]] <- all_data[[paste0("TDE", i)]]
}

# Compute average predicted values across the 5 imputations
all_data$TDE_avg <- Reduce("+", totalEff.hat) / 5
all_data$AP_avg <- Reduce("+", pers.hat) / 5

## Networks for illustrative purposes

# Define the countries to include in the comparison
cty <- c("Netherlands", "France", "Slovenia")

# Filter retweet data to only those observations where the sender is from the selected countries
cty_df <- full_df %>% dplyr::filter(ctry_sender %in% cty)

# Filter metadata for MPs in those countries
cty_meta <- meta %>% dplyr::filter(country %in% cty)

# Build igraph object:
# - Only keep retweets that occurred within the same country
# - Define edges using screen names
# - Attach country and party metadata to vertices (nodes)

rt_net <- graph_from_data_frame(
  cty_df %>%
    dplyr::filter(ctry_sender == ctry_rt) %>%
    dplyr::select(src_screen_name, rt_screen_name, Country = ctry_sender),
  vertices = cty_meta %>%
    dplyr::select(vertex, Party = party, Country = country)
)

# Set a random seed for consistent layout
set.seed(831213)

# Use ggraph to plot the network using the Fruchterman-Reingold layout algorithm
ggraph(rt_net, layout = "igraph", algorithm = "fr") +
  geom_edge_link0(edge_alpha = 0.1) + # Draw light gray edges (retweets)
  geom_node_point() + # Draw black points for each MP
  # geom_node_point(aes(color = Party)) +     # Optional: add party color
  guides(color = "none") + # Hide legend (not used here)
  facet_wrap(
    ~ factor(Country,
      levels = c("Netherlands", "France", "Slovenia")
    ),
    scales = "free"
  ) + # One panel per country
  theme_bw() + # Use clean theme
  theme( # Remove axes and gridlines for clarity
    axis.text.x = element_blank(), axis.ticks.x = element_blank(),
    axis.text.y = element_blank(), axis.ticks.y = element_blank(),
    panel.grid.major = element_blank(), panel.grid.minor = element_blank()
  ) +
  ylab("") + xlab("") # Remove axis labels

## Strength of connections within party vs. across party in prototypical countries

# Define the countries to plot
plot_ctr <- c("Netherlands", "France", "Slovenia")

# Filter metadata to those countries and merge with electoral variables
cty_meta <- meta %>%
  dplyr::filter(country %in% plot_ctr) %>%
  left_join(all_data %>% dplyr::select(vertex = id, country, AP = AP_avg, TDE = TDE_avg, regionalism, democracy, presidentialism)) %>%
  dplyr::filter(complete.cases(.))

# Filter retweet data to dyads where both sender and receiver are in selected metadata
cty_df <- full_df %>%
  dplyr::filter(src_screen_name %in% cty_meta$vertex & rt_screen_name %in% cty_meta$vertex) %>%
  dplyr::filter(ctry_sender %in% plot_ctr)

# Create a graph from dyads with weight = number of retweets
rt_net <- graph_from_data_frame(
  cty_df %>%
    dplyr::select(src_screen_name, rt_screen_name, weight = edge),
  vertices = cty_meta %>%
    dplyr::select(vertex, Party = party, country, magnitude, parfam, perseat, AP, TDE)
)


# Extract node and edge information from the graph
node_list <- igraph::as_data_frame(rt_net, what = "vertices")
edge_list <- igraph::as_data_frame(rt_net, what = "edges") %>%
  inner_join(node_list %>% dplyr::select(name, Party, country), by = c("from" = "name")) %>%
  inner_join(node_list %>% dplyr::select(name, Party, country), by = c("to" = "name")) %>%
  dplyr::filter(country.x == country.y) %>%
  mutate(country = country.x) %>%
  dplyr::select(-c(country.x, country.y)) %>%
  mutate(in_group = ifelse(Party.x == Party.y, "Within Party", "Across Parties"))

# Plot the distribution of tie strength by country and type of tie
ggplot(edge_list, aes(x = country, y = log(weight), fill = factor(in_group))) +
  ylab("Strength of connection\n(log nr. re-tweets)") +
  xlab("") +
  scale_fill_brewer(name = "Type of Connection", palette = "Greys") +
  geom_boxplot(outliers = FALSE) +
  theme_bw()

# Filter metadata to include only German MPs from SPD, CDU/CSU, and The Left Party
de_meta_2 <- all_data %>%
  dplyr::filter(country == "Germany") %>%
  ungroup() %>%
  dplyr::select(id, party) %>%
  dplyr::filter(party %in% c("SPD", "CDU/CSU", "The Left Party")) %>%
  distinct()

# Filter retweet dyads between these MPs within Germany
de_df_2 <- full_df %>%
  ungroup() %>%
  dplyr::filter((ctry_sender == "Germany")) %>%
  dplyr::filter((src_screen_name %in% de_meta_2$id) & (rt_screen_name %in% de_meta_2$id))

# Create graph object: directed network with retweet weights
rt_net_de <- graph_from_data_frame(de_df_2 %>% dplyr::select(src_screen_name, rt_screen_name, weight = edge),
  vertices = de_meta_2 %>% dplyr::select(vertex = id, party) %>% distinct(vertex, .keep_all = TRUE)
)

# Define grayscale color mapping by party
vertex.colors <- c("SPD" = gray(0.8), "CDU/CSU" = gray(0.5), "The Left Party" = gray(0.2))

# Set plotting parameters: remove margins, set title space
par(mar = c(.1, .1, 1.9, .1))

# Set seed for reproducible layout and compute Fruchterman-Reingold layout
set.seed(831213)
de_layout <- layout.fruchterman.reingold(rt_net_de)

# Plot the graph
plot(rt_net_de,
  layout = de_layout,
  vertex.shape = "circle", # Simple circular nodes
  vertex.size = 2, # Small node size
  vertex.color = vertex.colors[V(rt_net_de)$party], # Color by party
  edge.color = "lightgray", # Light gray directed edges
  edge.arrow.size = 0.15, # Small arrow heads
  vertex.label = NA, # No text labels
  edge.curved = 0.2, # Slight curvature in edges
  main = "Re-tweet Network\nLargest German Parties", # Title
  asp = FALSE
)


# Add manual legend in top-left corner
legend("topleft",
  legend = c("SPD", "CDU/CSU", "The Left Party"),
  pch = c(19),
  col = vertex.colors,
  bty = "n",
  title = ""
)

## Example for hypothetical parties
set.seed(831213)

# Create a maximally connected party (complete graph of 30 nodes)
full_c <- make_full_graph(30)

# Remove all but one of the edges from node 1, simulating a single disconnected outlier
full_c <- delete_edges(full_c, incident(full_c, 1)[-1])

# Create a maximally divided version: a spanning tree from the same nodes
not_c <- subgraph.edges(full_c, sample_spanning_tree(full_c))

# Convert igraph objects to 'network' class objects
full_c_g <- asNetwork(full_c)
not_c_g <- asNetwork(not_c)

# Combine both graphs into a block-diagonal matrix and estimate 1D latent positions using AME model
net_model_both <- ame(as.matrix(bdiag(as.matrix(full_c_g), as.matrix(not_c_g))),
  family = "bin",
  symmetric = TRUE,
  R = 1,
  burn = 2500, nscan = 500, odens = 1, print = FALSE, plot = FALSE
)

# Create party labels and compute disunity (squared distance from party mean, normalized)
sym_parties <- data.frame(
  pos = c(net_model_both$U),
  party = rep(c("A", "B"), each = 30)
) %>%
  group_by(party) %>%
  mutate(party_mean = mean(pos)) %>%
  ungroup() %>%
  mutate(dist = ((pos - party_mean)^2) / sd(pos))

# Plot layout: 2 rows x 2 columns
par(mfrow = c(2, 2))

# Top left: unified network
par(mar = c(.1, .1, .9, .1))
plot(full_c_g, edge.col = gray(0.6, 0.5), vertex.col = "black", main = "Maximally United")

# Top right: divided network
plot(not_c_g, edge.col = gray(0.6, 0.5), vertex.col = "black", main = "Maximally Divided")

# Bottom left: target plot for united party
par(mar = c(.1, .1, .1, .1))
gplot.target(full_c_g, -c(sym_parties$dist[1:30]),
  main = "",
  circ.lab = FALSE, circ.col = "black", vertex.col = "black",
  usearrows = FALSE, circ.lty = 1, circ.lwd = 0.2,
  edge.col = gray(0.6, 0.7), vertex.border = "white"
)

# Bottom right: target plot for divided party
gplot.target(not_c_g, -c(sym_parties$dist[31:60]),
  main = "", # eigen_centrality(not_c)$vector, main="",
  circ.lab = FALSE, circ.col = "black", vertex.col = "black",
  usearrows = FALSE, circ.lty = 1, circ.lwd = 0.2,
  edge.col = gray(0.6, 0.7), vertex.border = "white"
)

#### Centrality for parties in Germany

# Define a function that extracts retweet networks by party for a given country
party_net_fun <- function(plot_ctr, exclude = NULL) {
  # Filter metadata for the selected country and extract relevant variables
  cty_meta_2 <- all_data %>%
    dplyr::filter(country == plot_ctr) %>%
    ungroup() %>%
    dplyr::select(id, party, AP_avg, Distance, M) %>%
    distinct()
  # Filter retweet dyads where both MPs are from the selected country and present in metadata
  cty_df_2 <- full_df %>%
    ungroup() %>%
    dplyr::filter((ctry_sender == plot_ctr)) %>%
    dplyr::filter((src_screen_name %in% cty_meta_2$id) & (rt_screen_name %in% cty_meta_2$id))

  # Create graph object with party-level metadata
  rt_net_2 <- graph_from_data_frame(
    cty_df_2 %>%
      dplyr::select(src_screen_name, rt_screen_name, weight = edge),
    vertices = cty_meta_2 %>%
      dplyr::select(vertex = id, party, AP_avg, Distance, M) %>%
      distinct(vertex, .keep_all = TRUE) %>%
      mutate(size = log(M + 1))
  )

  # Optionally exclude some minor parties
  if (!is.null(exclude)) {
    party_names <- unique(cty_meta_2$party)[-exclude]
  } else {
    party_names <- unique(cty_meta_2$party)
  }

  # Get party sizes and keep only those with more than 3 members
  sizes <- table(cty_meta_2$party)[party_names]
  sizes <- sizes[order(sizes, decreasing = TRUE)]
  sizes <- sizes[sizes > 3]

  # Create a list of party-specific network graphs
  party_graphs <- lapply(
    names(sizes),
    \(party){
      asNetwork(induced_subgraph(rt_net_2, V(rt_net_2)$party == party))
    }
  )
  names(party_graphs) <- names(sizes)
  return(party_graphs)
}

# Apply function to Germany, excluding very small parties
de_graphs <- party_net_fun("Germany", c(5, 6, 9))

# Set layout and margins for plotting
set.seed(831213)
par(mfrow = c(2, 3))
par(mar = c(.1, .1, .9, .1))

# Loop through each party and create a target plot
for (i in seq_along(de_graphs)) {
  g <- de_graphs[[i]]

  # Plot target plot with:
  # - centrality = inverse of disunity (log-scaled)
  # - node size = scaled AP score
  gplot.target(g, -c(log((g %v% "Distance") + 1e-8)),
    main = names(de_graphs)[i],
    circ.lab = FALSE, circ.col = "black", vertex.col = "black",
    usearrows = FALSE, circ.lty = 1, circ.lwd = 0.2,
    edge.col = gray(0.5, 0.3), vertex.border = "white",
    vertex.cex = pmin(2.5, pmax(0.75, (scale(g %v% "AP_avg") + abs(min(scale(g %v% "AP_avg"))))))
  )
}

## Summary plot of disunity across countries, sorted by average AP score

all_data %>%
  # Reorder countries by their mean AP score
  mutate(country = factor(reorder(country, AP_avg, mean))) %>%
  # Begin ggplot: boxplot of disunity per country
  ggplot(aes(x = factor(reorder(country, AP_avg, mean)), y = log(Distance + 1e-8))) +

  # Draw boxplots for each country (excluding outliers for clarity)
  geom_boxplot(outliers = FALSE) +

  # Add a smoothed median trend line across countries
  geom_smooth(
    data = all_data %>%
      mutate(country = factor(reorder(country, AP_avg, mean))) %>%
      summarize(y = median(log(Distance + 1e-8)), .by = country),
    aes(x = factor(country), y = y, group = 1),
    se = FALSE, color = "black"
  ) +

  # Label axes
  xlab("") +
  ylab("Log Dis(unity)") +
  theme_bw() +

  # Rotate x-axis labels for readability
  theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1))

## Regression ####

# Add weighting variables:
# - peractive: proportional representation of country in dataset (n/totseats)
# - scale_factor: used to weight by representation across countries

all_data <- all_data %>%
  # (!(country %in% c("Netherlands","Slovenia"))) %>% #exclude to satisfy reviewer
  mutate(peractive = n() / totseats, .by = country) %>%
  mutate(scale_factor = n() / sum(peractive))

# Model 1: Linear regression, no random effects, only AP and TDE
mod_list_1 <- list()
mod_list_1 <- lapply(
  1:5,
  \(x){
    tmp <- all_data %>%
      dplyr::filter(Distance > 0)
    tmp[["AP"]] <- tmp[[paste0("AP", x)]]
    tmp[["TDE"]] <- tmp[[paste0("TDE", x)]]
    lm(log(Distance) ~ AP + TDE, weights = 1 / (peractive), data = tmp)
  }
)

# Model 2: Mixed model with country random intercept and polynomial AP
mod_list_2 <- list()
mod_list_2 <- lapply(
  1:5,
  \(x){
    tmp <- all_data %>%
      dplyr::filter(Distance > 0)
    tmp[["AP"]] <- tmp[[paste0("AP", x)]]
    tmp[["TDE"]] <- tmp[[paste0("TDE", x)]]
    lmer(log(Distance) ~ poly(AP, 2, raw = TRUE) + TDE + (1 | country),
      weights = 1 / (peractive), data = tmp
    )
  }
)

# Model 3: Mixed model with country and party random intercepts, full controls
mod_list_3 <- list()
mod_list_3 <- lapply(
  1:5,
  \(x){
    tmp <- all_data %>%
      dplyr::filter(Distance > 0)
    tmp[["AP"]] <- tmp[[paste0("AP", x)]]
    tmp[["TDE"]] <- tmp[[paste0("TDE", x)]]
    lmer(
      log(Distance) ~ poly(AP, 2, raw = TRUE) + TDE + democracy + regionalism +
        party_ideo + presidentialism * gov_party + perseat + (1 | country) + (1 | party),
      weights = 1 / (peractive), data = tmp
    )
  }
)

# Pool coefficients across imputations using Rubin's rules
lm_combined_1 <- pool(mod_list_1)
lm_combined_2 <- pool(mod_list_2)
lm_combined_3 <- pool(mod_list_3)

# Compute approximate R² for each model
r2_1 <- cor(predict(mod_list_1[[2]]), model.frame(mod_list_1[[1]])[, 1])^2
r2_2 <- cor(predict(mod_list_2[[2]]), model.frame(mod_list_2[[1]])[, 1])^2
r2_3 <- cor(predict(mod_list_3[[2]]), model.frame(mod_list_3[[1]])[, 1])^2

# Create regression table using texreg
texreg::htmlreg(list(lm_combined_1, lm_combined_2, lm_combined_3),
  custom.coef.map = list(
    "(Intercept)" = "(Intercept)",
    "AP" = "AP",
    "poly(AP, 2, raw = TRUE)1" = "AP",
    "poly(AP, 2, raw = TRUE)2" = "AP^2",
    "TDE" = "TDE",
    "regionalism" = "Regionalism",
    "presidentialism" = "Presidentialism",
    "gov_party" = "Party In Government?",
    "presidentialism:gov_party" = "Pres. $\\times$ In Government",
    "party_ideo" = "Party Ideology",
    "perseat" = "Party Legislative Size"
  ),
  custom.gof.names = c("N imputations", "N", " ", " "),
  custom.gof.rows = list(
    "Country RI" = c("No", "Yes", "Yes"),
    "Party RI" = c("No", "No", "Yes"),
    "R$^2$" = c(r2_1, r2_2, r2_3)
  )
)

# Create a new data frame with varying AP, holding all other predictors constant at typical values
pred_df <- with(
  all_data,
  data.frame(
    AP = seq(min(AP_avg),
      max(AP_avg),
      length.out = 100
    ), # range of AP
    TDE = mean(TDE_avg), # hold TDE constant
    regionalism = mean(regionalism),
    presidentialism = mean(presidentialism),
    democracy = mean(democracy),
    party_ideo = mean(party_ideo, na.rm = TRUE),
    gov_party = 1,
    perseat = 0.23,
    country = "Colombia", # Force new RI samples
    party = "Verde" # ditto
  )
)

# Simulate predicted values and error bands using bootstrapping
sim_pred <- bootMer(mod_list_3[[5]],
  \(x){
    predict(x, newdata = pred_df, re.form = NULL, allow.new.levels = TRUE)
  },
  nsim = 100,
  seed = 831213,
  use.u = TRUE,
  parallel = "multicore",
  ncpus = 10
)
sd_dist <- sd(all_data$Distance)

# Build a data frame with predictions and ±1 standard deviation error bands
pred_dis <- data.frame(
  AP = pred_df$AP,
  fit = (colMeans(sim_pred$t)),
  UB = (apply(sim_pred$t, 2, \(x)mean(x) + sd(x))),
  LB = (apply(sim_pred$t, 2, \(x)mean(x) - sd(x)))
)

# Plot predicted disunity with confidence ribbon
ggplot(pred_dis, aes(AP, fit)) +
  geom_ribbon(aes(ymin = LB, ymax = UB), fill = "gray") + # shaded confidence area
  geom_line(color = "white", linewidth = 1.1) + # fitted line
  ylab("Log Party Disunity") +
  xlab("Average Personalism") +
  theme_bw()

# Create a grid of hypothetical electoral systems varying by district magnitude and ballot type
pred_M_data <- expand.grid(
  formula = "dhondt",
  ballot_type = c("closed", "open"), # two ballot types: CLPR and OLPR
  new.nvotes = "One",
  pool_level = "party",
  threshold = 0.05,
  M = 2:150
) # simulate magnitude from 2 to 150


# Predict AP based on electoral rules using the previously trained GBM model
pred_AP_data <- predict(optimalGBMIntra,
  newdata = pred_M_data,
  n.trees = optimalGBMIntra$n.trees
)

# Create full prediction dataset for disunity model using predicted AP values
pred_dist_data <- with(
  all_data,
  data.frame(
    AP = pred_AP_data,
    TDE = mean(TDE_avg),
    M = pred_M_data$M,
    ballot_type = pred_M_data$ballot_type,
    regionalism = mean(regionalism),
    presidentialism = mean(presidentialism),
    democracy = mean(democracy),
    party_ideo = mean(party_ideo, na.rm = TRUE),
    gov_party = median(gov_party, na.rm = TRUE), # 1,
    perseat = mean(perseat, na.rm = TRUE), # 0.51,
    country = "Colombia",
    party = "Verde"
  )
)

# Simulate predicted values and uncertainty using bootstrapping
pred_MAP_sim <- bootMer(mod_list_3[[5]],
  \(x){
    predict(x,
      newdata = pred_dist_data,
      re.form = NULL,
      allow.new.levels = TRUE
    )
  },
  nsim = 100,
  seed = 831213,
  use.u = TRUE,
  parallel = "multicore",
  ncpus = 10
)

# Compile predicted values and ±1 SD intervals
pred_MAP <- data.frame(
  AP = pred_dist_data$AP,
  M = pred_dist_data$M,
  Ballot = c("closed" = "CLPR", "open" = "OLPR")[pred_dist_data$ballot_type],
  fit = (colMeans(pred_MAP_sim$t)),
  UB = (apply(pred_MAP_sim$t, 2, \(x)mean(x) + sd(x))),
  LB = (apply(pred_MAP_sim$t, 2, \(x)mean(x) - sd(x)))
)

# Plot predicted disunity as a function of M and ballot type
ggplot(pred_MAP, aes(M, fit, fill = Ballot, color = Ballot)) +
  geom_smooth(se = FALSE) +
  geom_ribbon(aes(ymin = LB, ymax = UB), alpha = 0.5, color = NA) +
  scale_color_manual(values = c("gray80", "black")) +
  scale_fill_manual(values = c("gray80", "black")) +
  ylab("Log Disunity(AP(M))") +
  xlab("Magnitude") +
  theme_bw()