# Load relevant libraries
library(gbm) # gradient boosting machines: supply the TDE and AP predictions
library(fixest) # fast fixed-effects estimation; supplies feols() below
# plyr is NOT attached. Its ddply() is called as plyr::ddply() below.
# Attaching plyr would mask dplyr's summarise(), mutate(), rename() and
# arrange() -- plyr has functions of the same names that ignore dplyr
# grouping -- and that masking would persist for any other chapter
# rendered afterwards in the same R session.
library(tidyverse) # bundle of data-handling packages (dplyr, ggplot2, ...)
library(mice) # multiple imputation; pools results across the five datasets
library(miscTools) # assorted utilities
library(kableExtra) # formats knitr tables for HTML output
library(tidyverse) # (loaded a second time; harmless, R simply ignores it)

# Function to run bootstrap, district level
boot.predict.se.enp <- function(tdeValue = 0.5, tot_divValue, reps = 500, seed = 123354, allCountries = TRUE, countryData) {
  # require() loads a package from inside a function, so it still works if
  # copied into a fresh session.
  require(fixest)
  require(miscTools)

  # Set Seed
  # Bootstrapping is random; fixing the seed makes the numbers reproducible.
  set.seed(seed)

  # Name the five TDE columns once, and refer to them by name from here on.
  # paste0("totalEff.hat", 1:5) builds the character vector
  #   "totalEff.hat1" "totalEff.hat2" ... "totalEff.hat5"
  # Both places below need to overwrite all five at once. Selecting them by
  # name rather than by position means the code keeps working if the column
  # order of the data ever changes -- and, unlike a numeric range, a wrong name
  # is visible on the page rather than silently shifting which variable is set.
  tde_cols <- paste0("totalEff.hat", 1:5)

  # Run Model
  # The same model fitted to each of the five sets of GBM predictions.
  # Reading the formula:
  #   enp                     effective number of parties -- the outcome
  #   totalEff.hat1 * tot_div TDE interacted with diversity. The * expands to
  #                           three terms: TDE alone, diversity alone, and their
  #                           product. The product is the chapter's argument.
  #   pers.hat1               Average Personalism, held constant as a control
  #   | ctyyear               country-year fixed effects: a separate intercept
  #                           for every country-year, so the estimate rests on
  #                           comparisons BETWEEN districts within the same
  #                           election, not across countries or over time
  #   se = 'hetero'           heteroskedasticity-robust standard errors
  # Note the index on the predictor names: model 1 uses totalEff.hat1 and
  # pers.hat1, model 2 uses ...hat2, and so on through the five predictions.
  temp_list <- list()
  temp_list[[1]] <- feols(enp ~ totalEff.hat1 * tot_div + pers.hat1 | ctyyear, se = "hetero", data = tempData)
  temp_list[[2]] <- feols(enp ~ totalEff.hat2 * tot_div + pers.hat2 | ctyyear, se = "hetero", data = tempData)
  temp_list[[3]] <- feols(enp ~ totalEff.hat3 * tot_div + pers.hat3 | ctyyear, se = "hetero", data = tempData)
  temp_list[[4]] <- feols(enp ~ totalEff.hat4 * tot_div + pers.hat4 | ctyyear, se = "hetero", data = tempData)
  temp_list[[5]] <- feols(enp ~ totalEff.hat5 * tot_div + pers.hat5 | ctyyear, se = "hetero", data = tempData)

  # Point Estimate
  # Choose the population to average over: every district in the data, or only
  # the districts of the system passed in through countryData.
  if (allCountries) {
    newData <- tempData # copy data
  } else {
    newData <- countryData # copy data
  }
  # Overwrite the two variables of interest, leaving everything else untouched.
  # This is the standard "hold all else constant" manoeuvre: whatever difference
  # appears between two calls is attributable to the values imposed here.
  # tde_cols, defined at the top of the function, names all five TDE prediction
  # columns. They are set at once so that each of the five fitted models sees
  # the imposed value rather than the observed one.
  newData[, tde_cols] <- tdeValue # Replace tde value
  newData[, "tot_div"] <- tot_divValue # Replace tot_div value
  # predict() gives one fitted value per district; mean() averages them into a
  # single expected N for this combination of TDE and diversity.
  point.estimate1 <- mean(predict(temp_list[[1]], newdata = newData))
  point.estimate2 <- mean(predict(temp_list[[2]], newdata = newData))
  point.estimate3 <- mean(predict(temp_list[[3]], newdata = newData))
  point.estimate4 <- mean(predict(temp_list[[4]], newdata = newData))
  point.estimate5 <- mean(predict(temp_list[[5]], newdata = newData))

  # Avg
  point.estimate <- (point.estimate1 + point.estimate2 + point.estimate3 + point.estimate4 + point.estimate5) / 5

  # Calculate SE
  # ------------------------------------------------------------------------
  # THE BOOTSTRAP
  # A point estimate says nothing about precision. The bootstrap answers "how
  # much would this number move if we had drawn a different sample?" by
  # repeatedly resampling the data we have, with replacement, and re-running
  # the whole calculation on each resample. The spread of the results estimates
  # the standard error.
  # One row per resample, one column per set of GBM predictions.
  # ------------------------------------------------------------------------
  sterrs <- matrix(NA, nrow = reps, ncol = 5)
  for (i in 1:reps) {
    # Create Data
    # Draw row numbers at random, as many as the data has rows, with
    # replacement -- so some rows appear twice and others not at all.
    index <- sample(1:nrow(tempData), nrow(tempData), replace = TRUE) # sample clusters
    bootdat <- tempData[index, ]
    # Run Model
    temp_list_boot <- list()
    temp_list_boot[[1]] <- feols(enp ~ totalEff.hat1 * tot_div + pers.hat1 | ctyyear, se = "hetero", data = bootdat)
    temp_list_boot[[2]] <- feols(enp ~ totalEff.hat2 * tot_div + pers.hat2 | ctyyear, se = "hetero", data = bootdat)
    temp_list_boot[[3]] <- feols(enp ~ totalEff.hat3 * tot_div + pers.hat3 | ctyyear, se = "hetero", data = bootdat)
    temp_list_boot[[4]] <- feols(enp ~ totalEff.hat4 * tot_div + pers.hat4 | ctyyear, se = "hetero", data = bootdat)
    temp_list_boot[[5]] <- feols(enp ~ totalEff.hat5 * tot_div + pers.hat5 | ctyyear, se = "hetero", data = bootdat)
    # Calculate predict value
    # Point Estimate
    if (allCountries) {
      newDataBootData <- bootdat # copy data
    } else {
      newDataBootData <- countryData # copy data
    }
    # Same substitution as above, on the resampled data. Selected by name via
    # tde_cols so both places cannot drift apart.
    newDataBootData[, tde_cols] <- tdeValue # Replace TDE value
    newDataBootData[, "tot_div"] <- tot_divValue # Replace tot_div value
    sterrs[i, 1] <- mean(predict(temp_list_boot[[1]], newdata = newDataBootData))
    sterrs[i, 2] <- mean(predict(temp_list_boot[[2]], newdata = newDataBootData))
    sterrs[i, 3] <- mean(predict(temp_list_boot[[3]], newdata = newDataBootData))
    sterrs[i, 4] <- mean(predict(temp_list_boot[[4]], newdata = newDataBootData))
    sterrs[i, 5] <- mean(predict(temp_list_boot[[5]], newdata = newDataBootData))
  }
  # prepare to return results
  # ------------------------------------------------------------------------
  # RUBIN'S RULES
  # Total uncertainty has two sources, added together here:
  #   WITHIN  -- mean(apply(sterrs, 2, var)) takes the variance down each of the
  #              five columns and averages them: ordinary sampling uncertainty.
  #   BETWEEN -- the sum-of-squares term measures how far the five column means
  #              sit from the overall mean: the extra uncertainty from not
  #              knowing the true TDE and AP values. The multiplier
  #              (5+1)/(5*(5-1)) is Rubin's correction for having five sets.
  # apply(X, 2, var) means "apply var() to each column" (2 = columns, 1 = rows).
  # ------------------------------------------------------------------------
  sd_pool <- mean(apply(sterrs, 2, var, na.rm = TRUE), na.rm = TRUE) + (5 + 1) / (5 * (5 - 1)) * (sum((colMeans(sterrs, na.rm = TRUE) - mean(sterrs, na.rm = TRUE))^2))

  # CIs
  # qnorm() returns the cut-points of a normal distribution: qnorm(0.025) is
  # -1.96 and qnorm(0.975) is +1.96, which bracket the middle 95%. The 90%
  # interval uses 0.05 and 0.95, giving the narrower inner band in the figures.
  lb95 <- point.estimate + sd_pool * qnorm(0.025)
  ub95 <- point.estimate + sd_pool * qnorm(0.975)
  lb90 <- point.estimate + sd_pool * qnorm(0.05)
  ub90 <- point.estimate + sd_pool * qnorm(0.95)
  # Bundle the six numbers into one named vector and hand it back to the caller.
  val <- c(point.estimate, sd_pool, lb95, ub95, lb90, ub90)
  names(val) <- c("Point Estimate", "Std. Error", "lb95", "ub95", "lb90", "ub90")
  return(val)
}

# Function to extract estimates for table 7.4
extract_Estimates <- function(model_name) {
  # Three nested ifelse() calls walk down the significance thresholds, and the
  # first match wins: p < 0.01 gets three stars, otherwise p < 0.05 gets two,
  # otherwise p < 0.10 gets one, otherwise none. paste0() glues the pieces into
  # a single string with the standard error in parentheses.
  out <- ifelse(model_name$p.value < 0.01, paste0(round(model_name$estimate, 3), "***", " (", round(model_name$std.error, 3), ")"),
    ifelse(model_name$p.value < 0.05, paste0(round(model_name$estimate, 3), "**", " (", round(model_name$std.error, 3), ")"),
      ifelse(model_name$p.value < 0.1, paste0(round(model_name$estimate, 3), "*", " (", round(model_name$std.error, 3), ")"),
        paste0(round(model_name$estimate, 3), " (", round(model_name$std.error, 3), ")")
      )
    )
  )
  names(out) <- model_name$term
  out <- data.frame(out)
  out$variable <- row.names(out)
  return(out)
}

# Set the path to Chapter 07 materials

# the short file names below. This is the one line to edit for your machine.


# load Potter's data
# District-level election returns, from Potter's dataset. Supplies the object
# `potter`, whose key columns are enp (effective number of parties), tot_div
# (social diversity), M (district magnitude) and the country-year identifiers.
load(file = "data/ch07/RData/potter_expanded.RData")

# Load district-level GBM objects
# The fitted gradient boosting models from Chapter 6. They are not re-estimated
# here; they are only asked to predict TDE and AP for these real districts.
load(file = "data/shared/AP_district_objects_t285_d13.RData")
load(file = "data/shared/TDE_district_objects_t485_d9.RData")

# create objects with optimal GBM
# $optimalGBM holds the best-performing fitted model from each of the five
# training runs, so each of these is a list of five models.
optimalGBMIntra <- totalAP.objects$optimalGBM
optimalGBMInter <- totalEffENP.objects$optimalGBM

# Load country level scores
# Country-level TDE and AP, merged in further down for the national-level
# comparison in Figure 7.1. Note the vintage in the filename: chapters use
# different versions of this file, and they are not interchangeable.
load(file = "data/shared/RealSystems_Scores_GBM_Aug_2024.RData")

### We need to run predictions five times
# Predict Interparty and Intraparty

potterData <- list() # Initialize a list to store predicted data

# One pass per fitted model: predict TDE and AP for every real district, then
# store that version of the dataset. Each pass overwrites the same two columns
# of `potter` before saving a copy, so the five stored datasets differ only in
# their predicted scores.
for (i in 1:5) {
  # Interparty incentive prediction: Total Duvergerian Effect (TDE)
  # n.trees tells the boosting model how many of its trees to use -- here all
  # of those selected as optimal during training.
  potter$totalEff.hat <- predict(optimalGBMInter[[i]], potter, n.trees = optimalGBMInter[[i]]$n.trees)
  # Intraparty incentive prediction: Average Personalism (AP)
  potter$pers.hat <- predict(optimalGBMIntra[[i]], potter, n.trees = optimalGBMIntra[[i]]$n.trees)
  # Store updated dataset in the list
  potterData[[i]] <- potter
}

# Combine all five versions into a single dataset for aggregation
# Get TDE and AP, then average the values across datasets
# sapply(1:5, function(i) ...) pulls the same columns out of each of the five
# datasets; simplify = FALSE keeps them as a list, and do.call('rbind', ...)
# stacks them into one tall table five times the original height.
all_TDE_AP_dist <- do.call("rbind", sapply(1:5, function(i) potterData[[i]][, c("cnty", "distname", "id", "year", "ctyyear", "distname", "M", "totalEff.hat", "pers.hat", "enp")], simplify = FALSE))

# Average predictions across imputations for each district-year
# ddply() splits the stacked table into groups, applies a function to each, and
# recombines. The grouping variables identify a single district in a single
# election, so this collapses the five stacked copies back into one row per
# district carrying the averaged TDE and AP. M and enp are identical across the
# five copies; taking their mean simply carries them through unchanged.
avg_TDE_AP_dist <- plyr::ddply(
  .data = all_TDE_AP_dist, .variables = c("cnty", "year", "ctyyear", "distname", "id"), .fun = plyr::summarise,
  M = mean(M),
  totalEff.hat = mean(totalEff.hat),
  pers.hat = mean(pers.hat),
  enp = mean(enp)
)


#############################################

# # Correlation between predicted TDE and observed party system size (enp)
# A first look before any modelling: cor.test() reports the correlation between
# predicted TDE and the observed number of parties, with a p-value. The
# chapter's argument predicts a NEGATIVE correlation -- more constraining
# systems, fewer parties. This is a raw association with nothing held constant,
# so it is a sanity check rather than evidence.
cor.test(all_TDE_AP_dist$enp, all_TDE_AP_dist$totalEff.hat)

#############################################

# Match country-level TDE data to Potter dataset by country-year
## Keep TDE (Country level) for countries and years in potter
data2export$ctyyear <- paste0(data2export$country, "-", data2export$year)
tdeCountryLevel <- subset(data2export, ctyyear %in% unique(potter$ctyyear))
# Keep only the first election per country-year
tdeCountryLevel <- subset(tdeCountryLevel, election_count == 1)
# Merge with district-level IDs for consistent identification
tdeCountryLevel <- merge(tdeCountryLevel, unique(avg_TDE_AP_dist[, c("ctyyear", "id")]), by = "ctyyear")

# Run Models

# Initialize lists to store regression models across the 5 imputations
# Eight empty containers, one per model in the grid above. Each will end up
# holding five fitted models -- one per set of GBM predictions.
m1_l <- list()
m1_lb <- list()
m2_l <- list()
m2_lb <- list()
m3_l <- list()
m3_lb <- list()
m4_l <- list()
m4_lb <- list()

# Model Estimation Loop
# One pass fits all eight models to prediction set i.
for (i in 1:5) {
  # Very basic model.
  # Raw association: does TDE alone track the number of parties?
  m1_l[[i]] <- feols(enp ~ totalEff.hat, se = "hetero", data = potterData[[i]])
  m1_lb[[i]] <- feols(enp ~ totalEff.hat | ctyyear, se = "hetero", data = potterData[[i]])
  # add AP
  # Check that TDE is not standing in for the intraparty dimension.
  m2_l[[i]] <- feols(enp ~ totalEff.hat + pers.hat, se = "hetero", data = potterData[[i]])
  m2_lb[[i]] <- feols(enp ~ totalEff.hat + pers.hat | ctyyear, se = "hetero", data = potterData[[i]])
  # Add total (constituency) diversity
  # Diversity enters as a separate cause: more heterogeneous electorates may
  # support more parties whatever the rules.
  m3_l[[i]] <- feols(enp ~ totalEff.hat + tot_div + pers.hat, se = "hetero", data = potterData[[i]])
  m3_lb[[i]] <- feols(enp ~ totalEff.hat + tot_div + pers.hat | ctyyear, se = "hetero", data = potterData[[i]])
  # With Interaction (between TDE and diversity)
  # The chapter's actual hypothesis. The * expands to three terms -- TDE,
  # diversity, and their product -- and it is the product that says the effect
  # of diversity DEPENDS ON how constraining the system is.
  m4_l[[i]] <- feols(enp ~ totalEff.hat * tot_div + pers.hat, se = "hetero", data = potterData[[i]])
  m4_lb[[i]] <- feols(enp ~ totalEff.hat * tot_div + pers.hat | ctyyear, se = "hetero", data = potterData[[i]])
}

# Combine Estimates Across Imputations
# Use mice::pool to combine results
# pool() applies Rubin's rules to the five fitted versions of each model,
# returning one combined set of coefficients and standard errors; summary()
# then presents it as the tidy table extract_Estimates() expects.
out1 <- summary(pool(m1_l))
out1b <- summary(pool(m1_lb))
out2 <- summary(pool(m2_l))
out2b <- summary(pool(m2_lb))
out3 <- summary(pool(m3_l))
out3b <- summary(pool(m3_lb))
out4 <- summary(pool(m4_l))
out4b <- summary(pool(m4_lb))

# We need to build a table for the book
# Gather model outputs
# The eight pooled summaries, in the column order they will appear.
model_list <- list(out1, out1b, out2, out2b, out3, out3b, out4, out4b)
out_estimates <- sapply(model_list, extract_Estimates, simplify = FALSE)

# Create table
# Merge coefficients across models by variable name
# all = TRUE keeps rows present in only some models -- essential here, because
# the simpler specifications have no diversity or interaction terms and those
# cells must come out blank rather than dropping the row entirely.
mytable <- merge(out_estimates[[1]], out_estimates[[2]], by = "variable", all = TRUE)
for (i in c(3:8)) mytable <- merge(mytable, out_estimates[[i]], by = "variable", all = TRUE)

# Reorder rows and rename them
# Reorder variables
mytable <- mytable[c(4, 2, 3, 5, 1), ]
# Rename rows
row.names(mytable) <- mytable$variable
mytable$variable <- NULL

# Calculate TSS: Total sum of squares for ENP (based on first imputation)
TSS <- sum((potterData[[1]]$enp - mean(potterData[[1]]$enp))^2)

# Compute R2 and Adj. R2 for all models
r2_m1 <- round(Reduce("+", lapply(m1_l, function(x) (1 - (sum(x$residuals^2) / TSS)))) / length(m1_l), 3)
r2_m1b <- round(Reduce("+", lapply(m1_lb, function(x) (1 - (sum(x$residuals^2) / TSS)))) / length(m1_lb), 3)
r2_m2 <- round(Reduce("+", lapply(m2_l, function(x) (1 - (sum(x$residuals^2) / TSS)))) / length(m2_l), 3)
r2_m2b <- round(Reduce("+", lapply(m2_lb, function(x) (1 - (sum(x$residuals^2) / TSS)))) / length(m2_lb), 3)
r2_m3 <- round(Reduce("+", lapply(m3_l, function(x) (1 - (sum(x$residuals^2) / TSS)))) / length(m3_l), 3)
r2_m3b <- round(Reduce("+", lapply(m3_lb, function(x) (1 - (sum(x$residuals^2) / TSS)))) / length(m3_lb), 3)
r2_m4 <- round(Reduce("+", lapply(m4_l, function(x) (1 - (sum(x$residuals^2) / TSS)))) / length(m4_l), 3)
r2_m4b <- round(Reduce("+", lapply(m4_lb, function(x) (1 - (sum(x$residuals^2) / TSS)))) / length(m4_lb), 3)

# Calculate Adj. Pseudo R2
adjr2_m1 <- round(Reduce("+", lapply(m1_l, function(x) (1 - (((sum(x$residuals^2) / TSS) * (x$nobs - 1)) / (x$nobs - x$nparams - 1))))) / length(m1_l), 3)
adjr2_m1b <- round(Reduce("+", lapply(m1_lb, function(x) (1 - (((sum(x$residuals^2) / TSS) * (x$nobs - 1)) / (x$nobs - x$nparams - 1))))) / length(m1_lb), 3)
adjr2_m2 <- round(Reduce("+", lapply(m2_l, function(x) (1 - (((sum(x$residuals^2) / TSS) * (x$nobs - 1)) / (x$nobs - x$nparams - 1))))) / length(m2_l), 3)
adjr2_m2b <- round(Reduce("+", lapply(m2_lb, function(x) (1 - (((sum(x$residuals^2) / TSS) * (x$nobs - 1)) / (x$nobs - x$nparams - 1))))) / length(m2_lb), 3)
adjr2_m3 <- round(Reduce("+", lapply(m3_l, function(x) (1 - (((sum(x$residuals^2) / TSS) * (x$nobs - 1)) / (x$nobs - x$nparams - 1))))) / length(m3_l), 3)
adjr2_m3b <- round(Reduce("+", lapply(m3_lb, function(x) (1 - (((sum(x$residuals^2) / TSS) * (x$nobs - 1)) / (x$nobs - x$nparams - 1))))) / length(m3_lb), 3)
adjr2_m4 <- round(Reduce("+", lapply(m4_l, function(x) (1 - (((sum(x$residuals^2) / TSS) * (x$nobs - 1)) / (x$nobs - x$nparams - 1))))) / length(m4_l), 3)
adjr2_m4b <- round(Reduce("+", lapply(m4_lb, function(x) (1 - (((sum(x$residuals^2) / TSS) * (x$nobs - 1)) / (x$nobs - x$nparams - 1))))) / length(m4_lb), 3)

# Add R2, adj.R2, and N rows to the regression table
mytable <- rbind(mytable, "Fixed Effects by Country-Year" = c("No", "Yes", "No", "Yes", "No", "Yes", "No", "Yes"))
mytable <- rbind(mytable, "R2" = c(r2_m1, r2_m1b, r2_m2, r2_m2b, r2_m3, r2_m3b, r2_m4, r2_m4b))
mytable <- rbind(mytable, "Adj R2" = c(adjr2_m1, adjr2_m1b, adjr2_m2, adjr2_m2b, adjr2_m3, adjr2_m3b, adjr2_m4, adjr2_m4b))
mytable <- rbind(mytable, "Observations" = c(
  m1_l[[1]]$nobs, m1_lb[[1]]$nobs,
  m2_l[[1]]$nobs, m2_lb[[1]]$nobs,
  m3_l[[1]]$nobs, m3_lb[[1]]$nobs,
  m4_l[[1]]$nobs, m4_lb[[1]]$nobs
))

# Rename rows
row.names(mytable) <- c(
  "TDE",
  "AP",
  "Constituency Diversity",
  "TDE x Constituency Diversity",
  "Intercept",
  "Fixed Effects by Country-Election",
  "R2",
  "Adj. R2",
  "Observations"
)

# Replace NA for '' (empty strings) and print table
mytable[is.na(mytable)] <- ""

# Create table
mytable %>%
  kbl(
    caption = ": Association Between tde, ap, and the Effective Number of Electoral Parties",
    col.names = c("(1)", "(2)", "(3)", "(4)", "(5)", "(6)", "(7)", "(8)"),
    format = "html",
    linesep = "",
  )

# Create unified dataset for prediction

# Rearange data to run predict() and calculate SEs
# Start from the first prediction set, keeping the outcome, the two predictors,
# the control and the fixed-effect identifier.
tempData <- potterData[[1]][, c("enp", "totalEff.hat", "tot_div", "pers.hat", "ctyyear")]
# Rename columns 2 and 4 so they carry the "1" suffix the function expects.
names(tempData)[c(2, 4)] <- c("totalEff.hat1", "pers.hat1")
# Adds TDE and AP predictions from the other four imputations.
tempData <- cbind(
  tempData, potterData[[2]][, c("totalEff.hat", "pers.hat")],
  potterData[[3]][, c("totalEff.hat", "pers.hat")],
  potterData[[4]][, c("totalEff.hat", "pers.hat")],
  potterData[[5]][, c("totalEff.hat", "pers.hat")]
)
names(tempData)[6:13] <- c(
  "totalEff.hat2", "pers.hat2", "totalEff.hat3", "pers.hat3",
  "totalEff.hat4", "pers.hat4", "totalEff.hat5", "pers.hat5"
)
# Reorder variables for clarity
tempData <- tempData[, c(
  "enp", "tot_div", "ctyyear",
  "pers.hat1", "pers.hat2", "pers.hat3", "pers.hat4", "pers.hat5",
  "totalEff.hat1", "totalEff.hat2", "totalEff.hat3", "totalEff.hat4", "totalEff.hat5"
)]

# Fix tde at minimum, median, and max, then vary tot_div (Constituency Diversity) from min to max
# seq(from, to, length.out = 15) builds a ladder of 15 evenly spaced diversity
# values spanning the observed range. Walking up that ladder twice -- once with
# TDE pinned at its minimum, once at its maximum -- traces the two lines of
# Figure 7.2. If the interaction is real, the two lines will not be parallel.
tot_divValues <- seq(min(tempData$tot_div), max(tempData$tot_div), length.out = 15)
outTDEMin <- list()
outTDEMax <- list()

# 15 rungs x 2 lines x 500 bootstrap resamples x 5 models: this loop is the
# slowest part of the chapter.
for (i in 1:length(tot_divValues)) {
  outTDEMin[[i]] <- boot.predict.se.enp(tdeValue = min(all_TDE_AP_dist$totalEff.hat), tot_divValue = tot_divValues[i], reps = 500, seed = 123354)
  outTDEMax[[i]] <- boot.predict.se.enp(tdeValue = max(all_TDE_AP_dist$totalEff.hat), tot_divValue = tot_divValues[i], reps = 500, seed = 123354)
}

# Fix tot_div at minimum, median, and max, then vary TDE from min to max
# The mirror image of the block above: now TDE varies along the ladder and
# diversity is pinned at its extremes. Same interaction viewed from the other
# side, which is a useful check that the finding is not an artefact of which
# variable is placed on the horizontal axis.
tdeValues <- seq(min(all_TDE_AP_dist$totalEff.hat), max(all_TDE_AP_dist$totalEff.hat), length.out = 15)
outTotDivMin <- list()
outTotDivMax <- list()
for (i in 1:length(tdeValues)) {
  outTotDivMin[[i]] <- boot.predict.se.enp(tdeValue = tdeValues[i], tot_divValue = min(tempData$tot_div), reps = 500, seed = 123354)
  outTotDivMax[[i]] <- boot.predict.se.enp(tdeValue = tdeValues[i], tot_divValue = max(tempData$tot_div), reps = 500, seed = 123354)
}

# Unlist results
# Combine all predictions into data frames
# Each list element is one six-number vector; rbind() stacks them into a table
# with 15 rows and columns for the estimate, its error, and the four CI bounds.
outTDEMin <- do.call("rbind", outTDEMin)
outTDEMax <- do.call("rbind", outTDEMax)
outTotDivMin <- do.call("rbind", outTotDivMin)
outTotDivMax <- do.call("rbind", outTotDivMax)

# pdf('~/Downloads/Plots/tde_district_national.pdf', width = 14)
# The commented pdf() and dev.off() lines would write the figure to a file
# instead of displaying it. Both are disabled so it appears inline.
# Margins: bottom 10, to leave room for the rotated country-year labels.
par(mar = c(10, 5, 2, 2))

# Order the data to match plotting order
# Both datasets must be sorted the same way, because the country-level points
# and the axis labels are positioned by the shared `id` column.
all_TDE_AP_dist <- all_TDE_AP_dist[order(all_TDE_AP_dist$ctyyear), ]
tdeCountryLevel <- tdeCountryLevel[order(tdeCountryLevel$ctyyear), ]

# Plot district-level TDE scores
# One grey dot per district. pch = 20 is a small filled circle, cex = 2 doubles
# its size, and axes = FALSE suppresses the default axes so custom ones can be
# drawn below.
plot(
  x = (all_TDE_AP_dist$id), (all_TDE_AP_dist$totalEff.hat), pch = 20, # filled circles
  ylab = "", xlab = "", axes = FALSE, cex = 2, col = "grey", ylim = c(0, 4.1)
)
# Add country-level TDE scores as black triangles
points(x = tdeCountryLevel$id, y = tdeCountryLevel$totalEff.hat, pch = 17)
# Draw x-axis with country-year labels
axis(1, at = tdeCountryLevel$id, labels = tdeCountryLevel$ctyyear, las = 2, cex.axis = 1.25)
# Draw y-axis for TDE
axis(2, las = 2, cex.axis = 1.25, at = seq(0, 4, 0.5))
mtext("TDE", 2, line = 3.5, cex = 1.25)

# dev.off()

# For text that goes with Figure 7.1
# Get summary stats for Switzerland (used in text)
# The square brackets subset the column to rows matching one country-year, and
# summary() reports min, quartiles, median, mean and max of what remains.
# The first line describes the spread ACROSS Swiss districts; the second gives
# the single national score. The gap between them is the figures quoted in the
# chapter text. Nothing is stored -- these print and are read off.
summary(avg_TDE_AP_dist$totalEff.hat[avg_TDE_AP_dist$ctyyear == "Switzerland-2007"])

summary(tdeCountryLevel$totalEff.hat[tdeCountryLevel$ctyyear == "Switzerland-2007"])

# pdf('~/Downloads/Plots/enp_hat.pdf', width = 8)

# Sets the layout to a single panel and adjusts plot margins for axis labels and clarity.
# layout() divides the drawing device into panels; a 1x1 matrix means one panel.
layout(matrix(c(1), ncol = 1))
par(mar = c(6, 6, 0, 1))

# Panel (a): Low TDE (Permissive System)
# Minimun
plot(
  x = tot_divValues,
  y = outTDEMin[, 1], type = "n", ylim = c(0, 10), axes = FALSE, xlim = c(0, 1.05),
  xlab = "", ylab = "", cex.lab = 2
)
# Adds axis labels for the Y and X axes, respectively.
mtext("Expected Effective Number of Parties", side = 2, line = 3, cex = 2)
mtext("Constituency Diversity", side = 1, line = 4, cex = 2)
# Draws the X and Y axes with enlarged text and rotated Y-axis labels.
axis(1, cex.axis = 2.5, at = round(seq(min(tempData$tot_div), max(tempData$tot_div), length.out = 5), 3))
axis(2, cex.axis = 2.5, las = 2)
# Fills in the 95% confidence interval around the predicted ENP curve
polygon(
  x = c(tot_divValues, rev(tot_divValues)),
  y = c(outTDEMin[, "ub95"], rev(outTDEMin[, "lb95"])), col = adjustcolor("grey", alpha.f = 0.95), border = NA
)
# Overlays the predicted ENP line in white for visual contrast over the shaded area.
lines(x = tot_divValues, y = outTDEMin[, 1], lty = 1, col = "white", lwd = 1)

# Panel (b): High TDE (Constraining System)
# Maximum

plot(
  x = tot_divValues,
  y = outTDEMax[, 1], type = "n", ylim = c(0, 10), axes = FALSE, xlim = c(0, 1.05),
  xlab = "", ylab = "", cex.lab = 2
)
# Repeats axis labeling.
mtext("Expected Effective Number of Parties", side = 2, line = 3, cex = 2)
mtext("Constituency Diversity", side = 1, line = 4, cex = 2)
# Draws axes with consistent formatting.
axis(1, cex.axis = 2.5, at = round(seq(min(tempData$tot_div), max(tempData$tot_div), length.out = 5), 3))
axis(2, cex.axis = 2.5, las = 2)
# Fills the confidence interval ribbon for the high-TDE prediction line.
polygon(
  x = c(tot_divValues, rev(tot_divValues)),
  y = c(outTDEMax[, "ub95"], rev(outTDEMax[, "lb95"])), col = adjustcolor("grey", alpha.f = 0.95), border = NA
)
# Adds the predicted ENP line for the high-TDE case.
lines(x = tot_divValues, y = outTDEMax[, 1], lty = 1, col = "white", lwd = 1)

# dev.off()

# pdf('~/Downloads/Plots/enp_hat_varyingTDE.pdf', width = 8)

# Sets a one-panel layout and defines margins for plot clarity.
layout(matrix(c(1), ncol = 1))
par(mar = c(6, 6, 0, 1))

# Panel (a): Low Diversity – Varying TDE
# Minimum
plot(
  x = tdeValues,
  y = outTotDivMin[, 1], type = "n", ylim = c(0, 10), axes = FALSE, xlim = c(min(all_TDE_AP_dist$totalEff.hat), max(all_TDE_AP_dist$totalEff.hat)),
  xlab = "", ylab = "", cex.lab = 2
)
# Adds axis labels
mtext("Expected Effective Number of Parties", side = 2, line = 3, cex = 2)
mtext("TDE", side = 1, line = 4, cex = 2)
# Draws x- and y-axes with large tick labels
axis(1, cex.axis = 2.5, at = round(seq(min(all_TDE_AP_dist$totalEff.hat), max(all_TDE_AP_dist$totalEff.hat), length.out = 5), 3))
axis(2, cex.axis = 2.5, las = 2)
# Fills the 95% confidence interval ribbon for low diversity
polygon(
  x = c(tdeValues, rev(tdeValues)),
  y = c(outTotDivMin[, "ub95"], rev(outTotDivMin[, "lb95"])), col = adjustcolor("grey", alpha.f = 0.95), border = NA
)
# Draws the predicted ENP line (white) over the shaded area for visual clarity
lines(x = tdeValues, y = outTotDivMin[, 1], lty = 1, col = "white", lwd = 1)

# Panel (b): High Diversity – Varying TDE
# Maximum
plot(
  x = tdeValues,
  y = outTotDivMax[, 1], type = "n", ylim = c(0, 10), axes = FALSE, xlim = c(min(all_TDE_AP_dist$totalEff.hat), max(all_TDE_AP_dist$totalEff.hat)),
  xlab = "", ylab = "", cex.lab = 2
)
# Labels the axes again
mtext("Expected Effective Number of Parties", side = 2, line = 3, cex = 2)
mtext("TDE", side = 1, line = 4, cex = 2)
# Draws tick marks
axis(1, cex.axis = 2.5, at = round(seq(min(all_TDE_AP_dist$totalEff.hat), max(all_TDE_AP_dist$totalEff.hat), length.out = 5), 3))
axis(2, cex.axis = 2.5, las = 2)
# Plots confidence interval and the prediction line for high diversity
polygon(
  x = c(tdeValues, rev(tdeValues)),
  y = c(outTotDivMax[, "ub95"], rev(outTotDivMax[, "lb95"])), col = adjustcolor("grey", alpha.f = 0.95), border = NA
)
lines(x = tdeValues, y = outTotDivMax[, 1], lty = 1, col = "white", lwd = 1)

# dev.off()

# Name the two blocks of prediction columns, and use these names everywhere
# below instead of column numbers. Each block holds five values -- one per
# imputation -- so rowSums(...)/5 averages across them.
#   ap_cols   pers.hat1     ... pers.hat5      (Average Personalism)
#   tde_cols  totalEff.hat1 ... totalEff.hat5  (Total Duvergerian Effect)
# Selecting by name rather than by position means the code cannot quietly
# operate on the wrong variables if the column order of the data ever changes.
ap_cols <- paste0("pers.hat", 1:5)
tde_cols <- paste0("totalEff.hat", 1:5)

# Find the range in Diversity and TDE

# Calculates the range of Constituency Diversity and TDE for each country-year using district-level data
div <- by(tempData$tot_div, tempData$ctyyear, function(x) max(x) - min(x), simplify = TRUE)
tde <- by(rowSums(tempData[, tde_cols]) / 5, tempData$ctyyear, function(x) max(x) - min(x), simplify = TRUE)
div <- array2DF(div, simplify = TRUE)
# The TDE value is averaged across the five imputations before computing the range.
tde <- array2DF(tde, simplify = TRUE)
# Converts the summary statistics into data frames and merges them for joint inspection.
div_tde <- merge(div, tde, by = "tempData$ctyyear")

# plot to find the observation at the top of the right-top corner (most illustrative cases)
plot(div_tde$Value.x, div_tde$Value.y, type = "n")
text(x = div_tde$Value.x, y = div_tde$Value.y, labels = div_tde$"tempData$ctyyear")

### We will use Switzerland-2007

# Vary diversity and fix TDE

# Filters the data to only include districts from Switzerland 2007.
swissData <- tempData[tempData$ctyyear == "Switzerland-2007", ]

# Identifies the minimum and maximum diversity values.
swissDataMin <- min(swissData$tot_div) # Min diversity
swissDataMax <- max(swissData$tot_div) # Max diversity

# Sets Average Personalism (AP) to its mean value for prediction consistency.
swissData[1, ap_cols] <- mean(rowSums(swissData[, ap_cols]) / 5) # Replace AP with its mean value

# Calculates the average TDE across imputations.
avgTDESwitzerland <- mean(rowSums(swissData[, tde_cols]) / 5) # Avg TDE for Switzerland

# Computes expected ENP for the least and most diverse Swiss districts, holding TDE and AP constant.
switzerlandDivMin <- boot.predict.se.enp(
  tdeValue = avgTDESwitzerland, tot_divValue =
    swissDataMin, reps = 500, seed = 123354,
  allCountries = FALSE, countryData = swissData[1, ]
)
switzerlandDivMax <- boot.predict.se.enp(
  tdeValue = avgTDESwitzerland, tot_divValue =
    swissDataMax, reps = 500, seed = 123354,
  allCountries = FALSE, countryData = swissData[1, ]
)
switzerlandDivMin

switzerlandDivMax

# Varying TDE with Fixed Diversity
swissData <- tempData[tempData$ctyyear == "Switzerland-2007", ]
# Identifies minimum and maximum TDE across Swiss districts.
swissDataMin <- min(rowSums(swissData[, tde_cols]) / 5) # Min TDE
swissDataMax <- max(rowSums(swissData[, tde_cols]) / 5) # Max TDE

# Fixes AP again at the mean.
swissData[1, ap_cols] <- mean(rowSums(swissData[, ap_cols]) / 5) # Replace AP with its mean value
# Fixes diversity at its average value across districts.
avgDivSwiss <- mean(swissData$tot_div[swissData$ctyyear == "Switzerland-2007"]) # Avg Diversity for Switzerland

# Computes expected ENP for districts with weakest and strongest electoral incentives, holding diversity and AP constant.
swissDivMin <- boot.predict.se.enp(
  tdeValue = swissDataMin, tot_divValue = avgDivSwiss,
  reps = 500, seed = 123354,
  allCountries = FALSE, countryData = swissData[1, ]
)
swissDivMax <- boot.predict.se.enp(
  tdeValue = swissDataMax, tot_divValue = avgDivSwiss,
  reps = 500, seed = 123354,
  allCountries = FALSE, countryData = swissData[1, ]
)

swissDivMin

swissDivMax

# Creates a data frame representing two versions of a “typical” closed-list PR system: one with low district magnitude (M = 3) and one with moderate magnitude (M = 8)
xTDE <- data.frame(
  formula = "dhondt",
  ballot_type = "closed",
  new.nvotes = "One",
  pool_level = "party",
  threshold = 0.025,
  M = c(3, 8)
)

# Ensure variables are treated as categorical
xTDE$formula <- as.factor(xTDE$formula)
xTDE$ballot_type <- as.factor(xTDE$ballot_type)
xTDE$new.nvotes <- as.factor(xTDE$new.nvotes)
xTDE$pool_level <- as.factor(xTDE$pool_level)

# # Get TDE predictions for each of the five imputed GBM models, these create lists
tde_hat1 <- predict(optimalGBMInter[[1]], xTDE, n.trees = optimalGBMInter[[1]]$n.trees)
tde_hat2 <- predict(optimalGBMInter[[2]], xTDE, n.trees = optimalGBMInter[[2]]$n.trees)
tde_hat3 <- predict(optimalGBMInter[[3]], xTDE, n.trees = optimalGBMInter[[3]]$n.trees)
tde_hat4 <- predict(optimalGBMInter[[4]], xTDE, n.trees = optimalGBMInter[[4]]$n.trees)
tde_hat5 <- predict(optimalGBMInter[[5]], xTDE, n.trees = optimalGBMInter[[5]]$n.trees)

# Take the average TDE across the five models
xTDE$tde <- colMeans(rbind(tde_hat1, tde_hat2, tde_hat3, tde_hat4, tde_hat5))

# Simulate ENP for average TDE, upper bound, and lower bound of “sweet spot”
# Calculate expected values
outAvg <- list()
outUB <- list() # outside the sweet spot, upper bound
outLB <- list() # outside the sweet spot, lower bound
for (i in 1:length(tot_divValues)) {
  # Average TDE from M = 3 & 8 systems
  outAvg[[i]] <- boot.predict.se.enp(tdeValue = mean(xTDE$tde), tot_divValue = tot_divValues[i], reps = 500, seed = 123354, allCountries = TRUE)
  # Upper bound TDE = 1.56 (M = 3, more constraining)
  outUB[[i]] <- boot.predict.se.enp(tdeValue = 1.56, tot_divValue = tot_divValues[i], reps = 500, seed = 123354, allCountries = TRUE)
  # Lower bound TDE = 1.21 (M = 8, more permissive)
  outLB[[i]] <- boot.predict.se.enp(tdeValue = 1.21, tot_divValue = tot_divValues[i], reps = 500, seed = 123354, allCountries = TRUE)
}

# Combine the results into data frames
outAvg <- do.call("rbind", outAvg)
outUB <- do.call("rbind", outUB)
outLB <- do.call("rbind", outLB)

# Start the plot
# pdf('~/Downloads//Plots/TDE_M3_8.pdf', width = 8)
layout(matrix(c(1), ncol = 1))
par(mar = c(6, 6, 0, 0))

# Creates an empty plot frame with the X-axis representing Constituency Diversity, and the Y-axis showing expected ENP (ranging from 2 to 6).
plot(
  x = tot_divValues,
  y = outAvg[, 1], type = "n", ylim = c(2, 6), axes = FALSE,
  xlab = "", ylab = "", cex.lab = 2
)
# Adds axis labels and tick marks.
mtext("Expected Effective Number of Parties", side = 2, line = 3, cex = 2)
mtext("Constituency Diversity", side = 1, line = 4, cex = 2)
axis(1, cex.axis = 2.5)
axis(2, cex.axis = 2.5, las = 2, at = c(2, 3, 4, 5, 6))
# Plots the expected ENP line for the “average” TDE case and shades the confidence interval ribbon.
polygon(
  x = c(tot_divValues, rev(tot_divValues)),
  y = c(outAvg[, 3], rev(outAvg[, 4])), col = adjustcolor("grey", alpha.f = 0.95), border = NA
)
lines(x = tot_divValues, y = outAvg[, 1], lty = 1, col = "white", lwd = 1)

# dev.off()