# =============================================================================
# 02_predictive.R
# Group 10 — SMB-A Capstone: Conflict Spillover Analysis
# Team: Karen David, Candace Offman, Mark Oleson
# Advisor: Dr. Elliot Bendoly | OSU Fisher College of Business
# =============================================================================
#
# PURPOSE:
#   Predictive stage of the IIA framework. Builds three models:
#     Model 1 — Negative Binomial GLM (predicts ndeath count)
#     Model 2 — Random Forest classifier (predicts high_sev binary)
#     Model 3 — Logistic Regression (predicts spillover_context)
#
# DEPENDS ON:
#   00_data_load_clean.R  — provides model_df, combined
#   03_spillover.R        — provides spill_hav column on model_df
#
# OUTPUTS:
#   Table 3  — Train/test split summary
#   Figure 8 — NB GLM coefficient plot
#   Figure 9 — ROC curve (Random Forest, AUC annotated)
#   Figure 10 — Variable importance (Gini decrease)
#   Table 5  — Logistic regression odds ratios
#   Table 6  — Model performance summary vs. analytical objective bounds
#   Saved objects: nb_model, rf_model, logit_spill
#
# ANALYTICAL OBJECTIVE BOUNDS (from corrected S-A3 tables):
#   AO1 — NB GLM   McFadden Pseudo-R² >= 0.10
#   AO4 — RF       Accuracy >= 0.80 | Precision >= 0.75 | Recall >= 0.75 | AUC >= 0.75
#   AO3 — Temporal R² >= 0.10  (handled in 03_spillover.R §7)
#
# =============================================================================

if (!exists("model_df")) {
  stop("Object 'model_df' not found. Source 00_data_load_clean.R first.")
}

suppressPackageStartupMessages({
  library(dplyr)
  library(MASS)
  library(pscl)
  library(randomForest)
  library(caret)
  library(pROC)
  library(broom)
  library(ggplot2)
  library(scales)
  library(knitr)
  library(kableExtra)
})

# -----------------------------------------------------------------------------
# 1. PREPARE MODELING DATASET
# -----------------------------------------------------------------------------
# If 03_spillover.R has already been sourced, model_df will have spill_hav.
# Otherwise fall back to the 2.5-degree grid-cell proxy.

if ("spill_hav" %in% names(model_df)) {
  spill_col <- "spill_hav"
} else if ("spillover_context" %in% names(model_df)) {
  spill_col <- "spillover_context"
} else {
  model_df$spillover_context <- 0L
  spill_col <- "spillover_context"
}

# Add country-year event count (ambient conflict density control)
# This was the second-highest RF importance variable but absent from
# baseline NB — resolved here by including it in both.
if (!"cy_events_log" %in% names(model_df)) {
  model_df <- model_df %>%
    group_by(countryname, start_year) %>%
    mutate(country_yr_events = n()) %>%
    ungroup() %>%
    mutate(cy_events_log = log1p(country_yr_events))
}

# Both escalation (continuous dose-response) and esc_high (threshold binary)
# are included. VIF checked in 03_spillover.R: escalation rises from 1.21
# to 6.93 with esc_high added; both below threshold of 10; justified by
# N > 10,000 and distinct theoretical meaning.
model_formula <- as.formula(paste(
  "ndeath ~",
  "factor(etype) + escalation + esc_high +",
  "log_duration + log_npart +",
  "cgovtarget + rgovtarget + repress +",
  "region + styr_scaled + cy_events_log +",
  spill_col
))

rf_formula <- as.formula(paste(
  "high_sev_f ~",
  "factor(etype) + escalation + esc_high +",
  "log_duration + log_npart +",
  "cgovtarget + rgovtarget + repress +",
  "region + styr_scaled + cy_events_log +",
  spill_col
))

model_vars_needed <- c(
  "ndeath", "log_ndeath", "high_sev",
  "etype", "escalation", "esc_high",
  "log_duration", "log_npart",
  "cgovtarget", "rgovtarget", "repress",
  "region", "styr_scaled", "cy_events_log",
  spill_col
)

model_df2 <- model_df %>%
  select(all_of(intersect(model_vars_needed, names(model_df)))) %>%
  drop_na()

model_df2$high_sev_f <- factor(
  model_df2$high_sev,
  levels = c(0, 1),
  labels = c("LowSev", "HighSev")
)

# -----------------------------------------------------------------------------
# 2. TRAIN / TEST SPLIT (stratified 75/25)
# -----------------------------------------------------------------------------

set.seed(2025)
train_idx <- createDataPartition(model_df2$high_sev_f, p = 0.75, list = FALSE)
train_df  <- model_df2[ train_idx, ]
test_df   <- model_df2[-train_idx, ]

split_summary <- data.frame(
  Split       = c("Training", "Test"),
  N           = c(nrow(train_df), nrow(test_df)),
  HighSev_N   = c(sum(train_df$high_sev), sum(test_df$high_sev)),
  HighSev_Pct = c(
    round(mean(train_df$high_sev) * 100, 1),
    round(mean(test_df$high_sev)  * 100, 1)
  )
)

kable(split_summary,
      col.names = c("Split", "N", "High-Sev N", "High-Sev %"),
      caption   = "Table 3: Stratified Train/Test Split (75/25)")

assign("train_df",  train_df,  envir = .GlobalEnv)
assign("test_df",   test_df,   envir = .GlobalEnv)
assign("model_df2", model_df2, envir = .GlobalEnv)

# -----------------------------------------------------------------------------
# 3. MODEL 1 — NEGATIVE BINOMIAL GLM
# -----------------------------------------------------------------------------
# ndeath is overdispersed count data (mean >> variance, median = 0).
# Negative Binomial GLM is the IIA framework-recommended model for
# overdispersed count outcomes.

nb_model <- tryCatch(
  MASS::glm.nb(model_formula, data = train_df),
  error = function(e) {
    message("NB GLM failed: ", conditionMessage(e))
    NULL
  }
)

if (!is.null(nb_model)) {

  null_model <- MASS::glm.nb(ndeath ~ 1, data = train_df)
  pseudo_r2  <- round(1 - as.numeric(logLik(nb_model)) /
                          as.numeric(logLik(null_model)), 3)

  nb_pred_log <- predict(nb_model, newdata = test_df, type = "response")
  nb_corr     <- round(cor(log1p(test_df$ndeath), log1p(nb_pred_log),
                           use = "complete.obs"), 3)
  nb_rmse     <- round(sqrt(mean((test_df$ndeath - nb_pred_log)^2,
                                  na.rm = TRUE)), 2)
  nb_aic      <- round(AIC(nb_model), 1)

  assign("nb_pseudo_r2", pseudo_r2, envir = .GlobalEnv)
  assign("nb_corr",      nb_corr,   envir = .GlobalEnv)
  assign("nb_rmse",      nb_rmse,   envir = .GlobalEnv)
  assign("nb_aic",       nb_aic,    envir = .GlobalEnv)
  assign("nb_model",     nb_model,  envir = .GlobalEnv)

  # Figure 8: NB GLM Coefficient Plot
  nb_tidy <- broom::tidy(nb_model, conf.int = TRUE, exponentiate = FALSE) %>%
    filter(term != "(Intercept)") %>%
    mutate(
      significant = p.value < 0.05,
      term_clean  = gsub("factor\\(etype\\)", "etype=", term),
      term_clean  = gsub("regionLatin America", "Region: Latin America", term_clean),
      term_clean  = gsub("_", " ", term_clean)
    )

  fig8 <- ggplot(nb_tidy,
                 aes(x = reorder(term_clean, estimate),
                     y = estimate,
                     color = significant)) +
    geom_point(size = 2.5) +
    geom_errorbar(aes(ymin = conf.low, ymax = conf.high), width = 0.3) +
    geom_hline(yintercept = 0, linetype = "dashed", color = "gray50") +
    coord_flip() +
    scale_color_manual(
      values = c("TRUE" = "#C0392B", "FALSE" = "gray60"),
      labels = c("TRUE" = "p < 0.05", "FALSE" = "p >= 0.05")
    ) +
    labs(
      title    = "Figure 8: Negative Binomial GLM — Coefficient Plot",
      subtitle = paste0("Outcome: ndeath (count) | McFadden Pseudo-R² = ", pseudo_r2),
      x        = NULL,
      y        = "Coefficient Estimate (log scale, NB link)",
      color    = NULL
    ) +
    theme_minimal(base_size = 11) +
    theme(plot.title = element_text(face = "bold"),
          legend.position = "bottom",
          panel.grid.minor = element_blank())

  print(fig8)
  if (!dir.exists("figures")) dir.create("figures")
  ggsave("figures/fig8_nb_coef_plot.png", fig8, width = 8, height = 6, dpi = 150)

} else {
  nb_pseudo_r2 <- NA; nb_corr <- NA; nb_rmse <- NA; nb_aic <- NA
}

# -----------------------------------------------------------------------------
# 4. MODEL 2 — RANDOM FOREST CLASSIFIER
# -----------------------------------------------------------------------------
# Predicts high_sev (ndeath > 10) as a binary outcome.
# Class imbalance (~4.4% HighSev) addressed via upsampling of minority class.

set.seed(2025)
train_up <- upSample(
  x     = train_df %>% select(-high_sev_f, -high_sev),
  y     = train_df$high_sev_f,
  yname = "high_sev_f"
)

rf_predictors <- c(
  "etype", "escalation", "esc_high",
  "log_duration", "log_npart",
  "cgovtarget", "rgovtarget", "repress",
  "styr_scaled", "cy_events_log",
  spill_col
)
train_up$region_n <- as.integer(train_up$region)
test_df$region_n  <- as.integer(factor(test_df$region,
                                        levels = c("Africa", "Latin America")))
rf_predictors <- c(rf_predictors, "region_n")
rf_predictors_present <- intersect(rf_predictors, names(train_up))

set.seed(2025)
rf_model <- randomForest(
  x         = train_up[, rf_predictors_present],
  y         = train_up$high_sev_f,
  ntree     = 300,
  mtry      = 3,
  importance = TRUE
)

assign("rf_model", rf_model, envir = .GlobalEnv)

rf_pred_class <- predict(rf_model,
                          newdata = test_df[, rf_predictors_present],
                          type = "class")
rf_pred_prob  <- predict(rf_model,
                          newdata = test_df[, rf_predictors_present],
                          type = "prob")[, "HighSev"]

cm <- confusionMatrix(rf_pred_class, test_df$high_sev_f, positive = "HighSev")

rf_accuracy  <- round(cm$overall["Accuracy"], 3)
rf_precision <- round(cm$byClass["Precision"], 3)
rf_recall    <- round(cm$byClass["Recall"], 3)
rf_f1        <- round(cm$byClass["F1"], 3)

roc_obj <- roc(
  response  = test_df$high_sev_f,
  predictor = rf_pred_prob,
  levels    = c("LowSev", "HighSev"),
  quiet     = TRUE
)
rf_auc <- round(as.numeric(auc(roc_obj)), 3)

assign("rf_accuracy",  rf_accuracy,  envir = .GlobalEnv)
assign("rf_precision", rf_precision, envir = .GlobalEnv)
assign("rf_recall",    rf_recall,    envir = .GlobalEnv)
assign("rf_f1",        rf_f1,        envir = .GlobalEnv)
assign("rf_auc",       rf_auc,       envir = .GlobalEnv)

# Figure 9: ROC Curve
roc_df <- data.frame(
  fpr = 1 - roc_obj$specificities,
  tpr = roc_obj$sensitivities
)

fig9 <- ggplot(roc_df, aes(x = fpr, y = tpr)) +
  geom_line(color = "#C0392B", linewidth = 1.2) +
  geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "gray60") +
  annotate("text", x = 0.65, y = 0.20,
           label = paste0("AUC = ", rf_auc),
           size = 5, fontface = "bold", color = "#C0392B") +
  scale_x_continuous(labels = percent_format()) +
  scale_y_continuous(labels = percent_format()) +
  labs(
    title    = "Figure 9: ROC Curve — Random Forest Classifier",
    subtitle = "Outcome: High-Severity Event (ndeath > 10) | Upsampled training set",
    x        = "False Positive Rate (1 - Specificity)",
    y        = "True Positive Rate (Sensitivity)"
  ) +
  theme_minimal(base_size = 12) +
  theme(plot.title = element_text(face = "bold"),
        panel.grid.minor = element_blank())

print(fig9)
ggsave("figures/fig9_roc_curve.png", fig9, width = 6, height = 5, dpi = 150)

# Figure 10: Variable Importance
importance_df <- as.data.frame(importance(rf_model)) %>%
  tibble::rownames_to_column("variable") %>%
  arrange(desc(MeanDecreaseGini)) %>%
  mutate(
    variable_clean = gsub("factor\\(etype\\)", "etype=", variable),
    variable_clean = gsub("_", " ", variable_clean),
    variable_clean = factor(variable_clean, levels = rev(variable_clean))
  )

fig10 <- ggplot(importance_df,
                aes(x = variable_clean, y = MeanDecreaseGini)) +
  geom_col(fill = "#C0392B", alpha = 0.8, width = 0.7) +
  coord_flip() +
  labs(
    title    = "Figure 10: Random Forest Variable Importance",
    subtitle = "Mean Decrease in Gini — escalation is the dominant predictor",
    x        = NULL,
    y        = "Mean Decrease in Gini Impurity"
  ) +
  theme_minimal(base_size = 11) +
  theme(plot.title = element_text(face = "bold"),
        panel.grid.minor = element_blank())

print(fig10)
ggsave("figures/fig10_variable_importance.png", fig10, width = 7, height = 5, dpi = 150)

# -----------------------------------------------------------------------------
# 5. MODEL 3 — LOGISTIC REGRESSION (SPILLOVER CONTEXT)
# -----------------------------------------------------------------------------
# Provides odds ratios for each predictor — interpretable for the
# Connections table (which utilities drive spillover risk?).

if (spill_col %in% names(model_df2) && var(model_df2[[spill_col]], na.rm = TRUE) > 0) {

  logit_formula <- as.formula(paste(
    spill_col, "~",
    "factor(etype) + escalation + esc_high +",
    "log_duration + log_npart +",
    "cgovtarget + rgovtarget + repress +",
    "region + styr_scaled + cy_events_log"
  ))

  logit_spill <- glm(logit_formula, data = train_df, family = binomial(link = "logit"))
  assign("logit_spill", logit_spill, envir = .GlobalEnv)

  logit_tidy <- broom::tidy(logit_spill, conf.int = TRUE, exponentiate = TRUE) %>%
    filter(term != "(Intercept)") %>%
    mutate(
      OR         = round(estimate, 3),
      CI_low     = round(conf.low, 3),
      CI_high    = round(conf.high, 3),
      p_value    = round(p.value, 4),
      significant = p.value < 0.05,
      term_clean = gsub("factor\\(etype\\)", "etype=", term),
      term_clean = gsub("regionLatin America", "Region: Latin America", term_clean),
      term_clean = gsub("_", " ", term_clean)
    ) %>%
    select(term_clean, OR, CI_low, CI_high, p_value, significant) %>%
    arrange(desc(abs(log(OR))))

  kable(logit_tidy,
        col.names = c("Predictor", "Odds Ratio", "95% CI Low",
                      "95% CI High", "p-value", "Significant"),
        caption   = "Table 5: Spillover Context — Logistic Regression Odds Ratios")

} else {
  logit_spill <- NULL
}

# -----------------------------------------------------------------------------
# 6. TABLE 6 — MODEL PERFORMANCE SUMMARY
# -----------------------------------------------------------------------------

table6 <- data.frame(
  Model = c(
    "NB GLM (ndeath count)",
    "Random Forest (high_sev)",
    "Logistic Reg. (spillover)"
  ),
  Metric = c(
    paste0("Pseudo-R² = ", ifelse(exists("nb_pseudo_r2"), nb_pseudo_r2, "N/A")),
    paste0("AUC = ", rf_auc, " | Accuracy = ", rf_accuracy),
    ifelse(!is.null(logit_spill), "Fit (see Table 5)", "Not run")
  ),
  Bound = c(
    "Pseudo-R² >= 0.10",
    "AUC >= 0.75 | Accuracy >= 0.80",
    "N/A"
  ),
  Meets_Target = c(
    ifelse(exists("nb_pseudo_r2") && !is.na(nb_pseudo_r2),
           ifelse(nb_pseudo_r2 >= 0.10, "YES", "NO"), "N/A"),
    ifelse(rf_auc >= 0.75 && rf_accuracy >= 0.80, "YES", "NO"),
    "—"
  )
)

kable(table6,
      col.names = c("Model", "Performance", "Analytical Obj. Bound", "Meets Target"),
      caption   = "Table 6: Model Performance Summary vs. Analytical Objective Bounds")