# ============================================================================
# Additional File 2
# Reproducible R analysis code for:
# Benchmarking Machine Learning Models for Cardiovascular Disease
# Classification in Population Health Surveys
# Target: BMC Artificial Intelligence, Collection: AI in chronic disease management
# ============================================================================
# Analytical map to manuscript:
# S01  Setup and package checks
# S02  Data import and cohort construction
# S03  Predictor definitions and temporal pairs
# S04  Survey, country-year, and class-imbalance weights (Manuscript Eq. 1)
# S05  Fixed country-grouped internal-external folds
# S06  Fold-specific preprocessing
# S07  Weighted performance metrics (including Manuscript Eq. 5)
# S08  Probability recalibration and threshold selection (Manuscript Eq. 4)
# S09  Model fitting: logistic, elastic net, random forest, XGBoost
# S10  Five-fold country-grouped external benchmarking
# S11  Detection-adjusted benchmark
# S12  Geographic and subgroup summaries
# S13  Earlier-to-later temporal benchmarking
# S14  XGBoost SHAP importance and figures
# S15  Reproducibility outputs
# ============================================================================

# S01. Setup ------------------------------------------------------------------
options(stringsAsFactors = FALSE)
SEED <- 20260903L
set.seed(SEED)

required_packages <- c(
  "data.table", "Matrix", "glmnet", "ranger", "xgboost", "ggplot2"
)
missing_packages <- required_packages[!vapply(required_packages, requireNamespace,
                                               quietly = TRUE, FUN.VALUE = logical(1))]
if (length(missing_packages) > 0L) {
  stop(
    "Install the following R packages before running this script: ",
    paste(missing_packages, collapse = ", ")
  )
}

library(data.table)
library(Matrix)
library(glmnet)
library(ranger)
library(xgboost)
library(ggplot2)

DATA_FILE <- "Additional_File_3_CVD_analysis_dataset.csv.gz"
OUT_DIR <- "analysis_outputs"
dir.create(OUT_DIR, showWarnings = FALSE, recursive = TRUE)
dir.create(file.path(OUT_DIR, "tables"), showWarnings = FALSE, recursive = TRUE)
dir.create(file.path(OUT_DIR, "figures"), showWarnings = FALSE, recursive = TRUE)

# S02. Data import and cohort construction ------------------------------------
dat <- fread(DATA_FILE, na.strings = c("", "NA", "NaN"))

valid_outcome <- c("Yes", "No")
dat <- dat[cvd_history %chin% valid_outcome]
dat[, y := as.integer(cvd_history == "Yes")]

# The analysis dataset distributed with the package has already been restricted
# to countries with both CVD outcome classes. This check protects reproducibility.
elig_check <- dat[, .(
  n = .N,
  cases = sum(y == 1L),
  controls = sum(y == 0L)
), by = .(country, region15)]
if (any(elig_check$cases == 0L | elig_check$controls == 0L)) {
  stop("The analysis dataset contains a country without both CVD classes.")
}

# Normalize survey weights to mean 1 within country-year. Missing or nonpositive
# weights are replaced by the median positive country-year weight.
dat[, wstep1_num := suppressWarnings(as.numeric(wstep1))]
dat[, eval_w := {
  w <- wstep1_num
  positive <- w[is.finite(w) & w > 0]
  replacement <- if (length(positive)) median(positive) else 1
  w[!is.finite(w) | w <= 0] <- replacement
  w / mean(w)
}, by = .(country, survyear)]

# S03. Predictors and temporal pairs ------------------------------------------
numeric_predictors <- c("age", "bmi", "dash_score")
categorical_full <- c(
  "gender", "residence", "edlev", "occupation", "mstatus", "smokehis",
  "alcharm", "physical_activity", "highsalt", "lfi", "lvi",
  "hypertension", "diabetes"
)
full_predictors <- c(numeric_predictors, categorical_full)

categorical_adjusted <- setdiff(categorical_full, c("hypertension", "diabetes"))
adjusted_predictors <- c(numeric_predictors, categorical_adjusted)

temporal_pairs <- data.table(
  country = c("Mongolia", "Uganda", "Viet Nam"),
  train_year = c(2013L, 2014L, 2015L),
  test_year = c(2019L, 2023L, 2021L)
)

# S04. Training weights, corresponding to Manuscript Eq. 1 --------------------
# w_i* = w_i^(survey) * b_{country,year} * a_{class}
# a_k = W / (2 W_k)
#
# b is scaled as mean country-year sample size / current country-year sample size.
# This is an inverse sample-size factor with a scale that keeps typical case
# weights close to 1, which matters for tree algorithms with weight-sensitive
# regularization parameters.
make_training_weights <- function(d) {
  size_table <- d[, .N, by = .(country, survyear)]
  mean_n <- mean(size_table$N)
  d2 <- merge(
    copy(d), size_table,
    by = c("country", "survyear"), all.x = TRUE, sort = FALSE
  )
  base <- d2$eval_w * (mean_n / d2$N)
  W <- sum(base)
  W0 <- sum(base[d2$y == 0L])
  W1 <- sum(base[d2$y == 1L])
  class_factor <- ifelse(d2$y == 1L, W / (2 * W1), W / (2 * W0))
  out <- base * class_factor
  # Restore original row order after merge.
  out[match(d$row_id_internal, d2$row_id_internal)]
}

dat[, row_id_internal := .I]

# S05. Fixed country-grouped external folds ----------------------------------
# These assignments reproduce the five external folds used for the manuscript.
country_fold <- c(
  "Afghanistan" = 5, "Algeria" = 5, "Armenia" = 4, "Aruba" = 1,
  "Azerbaijan" = 1, "Belarus" = 1, "Benin" = 1, "Bolivia" = 1,
  "Botswana" = 1, "Brunei Darussalam" = 3, "Burkina Faso" = 2,
  "Cabo Verde" = 4, "Cook Islands" = 3, "Ecuador" = 3,
  "Eswatini" = 5, "Ethiopia" = 1, "Georgia" = 5, "Guyana" = 2,
  "Iraq" = 2, "Jordan" = 4, "Kenya" = 5, "Kiribati" = 2,
  "Kuwait" = 5, "Kyrgyzstan" = 1, "Libya" = 2, "Malawi" = 4,
  "Micronesia (Federated States of)" = 1, "Moldova" = 3,
  "Mongolia" = 2, "Morocco" = 3, "Nauru" = 4, "Pakistan" = 4,
  "Saint Lucia" = 1, "Sao Tome and Principe" = 3,
  "Solomon Islands" = 3, "Tajikistan" = 2, "Tokelau" = 2,
  "Tonga" = 2, "Turkmenistan" = 4, "Tuvalu" = 3, "Uganda" = 5,
  "Ukraine" = 4, "Uzbekistan" = 4, "Viet Nam" = 3,
  "Wallis and Futuna" = 2, "Zambia" = 3
)

dat[, external_fold := unname(country_fold[country])]
if (anyNA(dat$external_fold)) stop("A country is missing from the fixed fold map.")

# S06. Fold-specific preprocessing -------------------------------------------
# Numeric variables: median imputation. Standardization is applied for logistic
# and elastic-net regression. Categorical variables: modal imputation, rare
# development levels with fewer than 10 records collapsed into OTHER, then
# one-hot encoding. Unseen test levels are mapped to OTHER.

fit_preprocessor <- function(train, numeric_vars, categorical_vars) {
  medians <- vapply(train[, ..numeric_vars], function(x) median(x, na.rm = TRUE), numeric(1))
  medians[!is.finite(medians)] <- 0
  means <- numeric(length(numeric_vars)); sds <- numeric(length(numeric_vars))
  names(means) <- names(sds) <- numeric_vars

  cat_info <- list()
  for (v in categorical_vars) {
    x <- as.character(train[[v]])
    x[is.na(x) | x == ""] <- "__MISSING__"
    tab <- sort(table(x), decreasing = TRUE)
    mode_value <- if (length(tab)) names(tab)[1] else "__MISSING__"
    rare <- names(tab)[tab < 10L]
    kept <- setdiff(names(tab), rare)
    levels_v <- unique(c(kept, "OTHER"))
    cat_info[[v]] <- list(mode = mode_value, rare = rare, levels = levels_v)
  }

  temp_num <- copy(train[, ..numeric_vars])
  for (v in numeric_vars) {
    x <- as.numeric(temp_num[[v]])
    x[!is.finite(x)] <- medians[[v]]
    means[[v]] <- mean(x)
    sds[[v]] <- sd(x)
    if (!is.finite(sds[[v]]) || sds[[v]] == 0) sds[[v]] <- 1
  }

  list(
    numeric_vars = numeric_vars,
    categorical_vars = categorical_vars,
    medians = medians,
    means = means,
    sds = sds,
    cat_info = cat_info
  )
}

apply_preprocessor <- function(d, pp, standardize_numeric = FALSE) {
  parts <- list()

  if (length(pp$numeric_vars)) {
    num <- as.data.frame(d[, ..pp$numeric_vars])
    for (v in pp$numeric_vars) {
      x <- as.numeric(num[[v]])
      x[!is.finite(x)] <- pp$medians[[v]]
      if (standardize_numeric) x <- (x - pp$means[[v]]) / pp$sds[[v]]
      num[[v]] <- x
    }
    parts[["numeric"]] <- Matrix(as.matrix(num), sparse = TRUE)
  }

  if (length(pp$categorical_vars)) {
    cat_frames <- list()
    for (v in pp$categorical_vars) {
      info <- pp$cat_info[[v]]
      x <- as.character(d[[v]])
      x[is.na(x) | x == ""] <- info$mode
      x[x %in% info$rare | !(x %in% info$levels)] <- "OTHER"
      f <- factor(x, levels = info$levels)
      mm <- sparse.model.matrix(~ f - 1)
      colnames(mm) <- paste0(v, "__", sub("^f", "", colnames(mm)))
      cat_frames[[v]] <- mm
    }
    parts[["categorical"]] <- do.call(cbind, cat_frames)
  }

  do.call(cbind, parts)
}

# S07. Weighted metrics --------------------------------------------------------
weighted_roc_auc <- function(y, p, w) {
  o <- order(p, decreasing = TRUE)
  y <- y[o]; w <- w[o]
  pos <- sum(w[y == 1L]); neg <- sum(w[y == 0L])
  if (pos <= 0 || neg <= 0) return(NA_real_)
  tpr <- c(0, cumsum(w * (y == 1L)) / pos)
  fpr <- c(0, cumsum(w * (y == 0L)) / neg)
  sum(diff(fpr) * (head(tpr, -1) + tail(tpr, -1)) / 2)
}

weighted_pr_auc <- function(y, p, w) {
  o <- order(p, decreasing = TRUE)
  y <- y[o]; w <- w[o]
  pos_total <- sum(w[y == 1L])
  if (pos_total <= 0) return(NA_real_)
  tp <- cumsum(w * (y == 1L))
  fp <- cumsum(w * (y == 0L))
  precision <- tp / pmax(tp + fp, .Machine$double.eps)
  delta_recall <- (w * (y == 1L)) / pos_total
  sum(precision * delta_recall)
}

weighted_brier <- function(y, p, w) {
  # Manuscript Eq. 5
  sum(w * (y - p)^2) / sum(w)
}

weighted_classification <- function(y, p, w, threshold) {
  pred <- as.integer(p >= threshold)
  tp <- sum(w[y == 1L & pred == 1L]); fn <- sum(w[y == 1L & pred == 0L])
  tn <- sum(w[y == 0L & pred == 0L]); fp <- sum(w[y == 0L & pred == 1L])
  data.table(
    sensitivity = tp / (tp + fn),
    specificity = tn / (tn + fp),
    ppv = tp / (tp + fp),
    npv = tn / (tn + fn),
    f1 = 2 * tp / (2 * tp + fp + fn)
  )
}

calibration_stats <- function(y, p, w) {
  z <- qlogis(pmin(pmax(p, 1e-6), 1 - 1e-6))
  fit <- glm(y ~ z, family = binomial(), weights = w)
  data.table(
    calibration_intercept = unname(coef(fit)[1]),
    calibration_slope = unname(coef(fit)[2])
  )
}

all_metrics <- function(d) {
  basic <- data.table(
    n = nrow(d), cases = sum(d$y),
    weighted_prevalence = weighted.mean(d$y, d$eval_w),
    roc_auc = weighted_roc_auc(d$y, d$prob, d$eval_w),
    pr_auc = weighted_pr_auc(d$y, d$prob, d$eval_w),
    brier = weighted_brier(d$y, d$prob, d$eval_w)
  )
  cbind(
    basic,
    weighted_classification(d$y, d$prob, d$eval_w, d$threshold),
    calibration_stats(d$y, d$prob, d$eval_w)
  )
}

# S08. Recalibration and threshold selection ---------------------------------
fit_platt <- function(y, p_raw, w) {
  z <- qlogis(pmin(pmax(p_raw, 1e-6), 1 - 1e-6))
  glm(y ~ z, family = binomial(), weights = w)
}

apply_platt <- function(fit, p_raw) {
  z <- qlogis(pmin(pmax(p_raw, 1e-6), 1 - 1e-6))
  plogis(coef(fit)[1] + coef(fit)[2] * z)
}

select_threshold <- function(y, p, w) {
  candidates <- seq(0.02, 0.70, by = 0.01)
  score <- vapply(candidates, function(t) {
    weighted_classification(y, p, w, t)$f1
  }, numeric(1))
  candidates[which.max(score)]
}

# S09. Candidate model fitting ------------------------------------------------
fit_candidate_model <- function(model_name, x_train, y_train, w_train) {
  if (model_name == "Logistic") {
    fit <- glmnet(
      x_train, y_train, family = "binomial", alpha = 0,
      lambda = 1 / nrow(x_train), weights = w_train, standardize = FALSE
    )
    return(list(name = model_name, object = fit, s = 1 / nrow(x_train)))
  }

  if (model_name == "ElasticNet") {
    fit <- glmnet(
      x_train, y_train, family = "binomial", alpha = 0.5,
      lambda = 2 / nrow(x_train), weights = w_train, standardize = FALSE
    )
    return(list(name = model_name, object = fit, s = 2 / nrow(x_train)))
  }

  if (model_name == "RandomForest") {
    rf_dat <- data.frame(y = factor(y_train), as.matrix(x_train), check.names = TRUE)
    fit <- ranger(
      y ~ ., data = rf_dat, probability = TRUE,
      num.trees = 160, max.depth = 14, min.node.size = 10,
      mtry = max(1L, floor(sqrt(ncol(x_train)))),
      case.weights = w_train, seed = SEED,
      num.threads = max(1L, parallel::detectCores() - 1L)
    )
    return(list(name = model_name, object = fit))
  }

  if (model_name == "XGBoost") {
    dtrain <- xgb.DMatrix(x_train, label = y_train, weight = w_train)
    fit <- xgb.train(
      params = list(
        objective = "binary:logistic", eval_metric = "logloss",
        max_depth = 4, eta = 0.05, subsample = 0.8,
        colsample_bytree = 0.8, min_child_weight = 5,
        lambda = 1.0, nthread = max(1L, parallel::detectCores() - 1L)
      ),
      data = dtrain, nrounds = 200, verbose = 0
    )
    return(list(name = model_name, object = fit))
  }

  stop("Unknown model: ", model_name)
}

predict_candidate_model <- function(fit, x_new) {
  if (fit$name %in% c("Logistic", "ElasticNet")) {
    return(as.numeric(predict(fit$object, newx = x_new, s = fit$s, type = "response")))
  }
  if (fit$name == "RandomForest") {
    nd <- data.frame(as.matrix(x_new), check.names = TRUE)
    pr <- predict(fit$object, data = nd)$predictions
    if (is.matrix(pr)) return(as.numeric(pr[, "1"]))
    return(as.numeric(pr))
  }
  if (fit$name == "XGBoost") {
    return(as.numeric(predict(fit$object, xgb.DMatrix(x_new))))
  }
  stop("Unknown fitted model type.")
}

# S10. Country-grouped benchmark ---------------------------------------------
run_external_benchmark <- function(predictor_set = c("Full", "Detection-adjusted")) {
  predictor_set <- match.arg(predictor_set)
  cat_vars <- if (predictor_set == "Full") categorical_full else categorical_adjusted
  predictors <- c(numeric_predictors, cat_vars)
  model_names <- c("Logistic", "ElasticNet", "RandomForest", "XGBoost")
  all_predictions <- list()
  meta <- list()

  for (fold in 1:5) {
    message("External fold ", fold, " | ", predictor_set)
    dev <- dat[external_fold != fold]
    test <- dat[external_fold == fold]

    set.seed(SEED + fold)
    case_ids <- which(dev$y == 1L); ctrl_ids <- which(dev$y == 0L)
    cal_case <- sample(case_ids, size = floor(0.15 * length(case_ids)))
    cal_ctrl <- sample(ctrl_ids, size = floor(0.15 * length(ctrl_ids)))
    cal_idx <- sort(c(cal_case, cal_ctrl))
    calibration <- dev[cal_idx]
    train <- dev[-cal_idx]

    pp <- fit_preprocessor(train, numeric_predictors, cat_vars)

    for (model_name in model_names) {
      scale_numeric <- model_name %in% c("Logistic", "ElasticNet")
      x_train <- apply_preprocessor(train, pp, standardize_numeric = scale_numeric)
      x_cal <- apply_preprocessor(calibration, pp, standardize_numeric = scale_numeric)
      x_test <- apply_preprocessor(test, pp, standardize_numeric = scale_numeric)

      w_train <- make_training_weights(train)
      model <- fit_candidate_model(model_name, x_train, train$y, w_train)
      raw_cal <- predict_candidate_model(model, x_cal)
      raw_test <- predict_candidate_model(model, x_test)

      platt <- fit_platt(calibration$y, raw_cal, calibration$eval_w)
      p_cal <- apply_platt(platt, raw_cal)
      p_test <- apply_platt(platt, raw_test)
      threshold <- select_threshold(calibration$y, p_cal, calibration$eval_w)

      pred <- test[, .(
        record_id, country, region15, survyear, gender, age_group,
        residence, y, eval_w, external_fold
      )]
      pred[, `:=`(
        predictor_set = predictor_set,
        model = model_name,
        prob = p_test,
        threshold = threshold
      )]
      all_predictions[[length(all_predictions) + 1L]] <- pred

      meta[[length(meta) + 1L]] <- data.table(
        predictor_set = predictor_set, model = model_name, fold = fold,
        n_train = nrow(train), n_calibration = nrow(calibration),
        n_external = nrow(test), threshold = threshold,
        platt_intercept = coef(platt)[1], platt_slope = coef(platt)[2]
      )
    }
  }

  list(predictions = rbindlist(all_predictions), metadata = rbindlist(meta))
}

# Execute both predictor sets.
benchmark_full <- run_external_benchmark("Full")
benchmark_adjusted <- run_external_benchmark("Detection-adjusted")

fwrite(benchmark_full$predictions,
       file.path(OUT_DIR, "tables", "external_predictions_full.csv.gz"))
fwrite(benchmark_adjusted$predictions,
       file.path(OUT_DIR, "tables", "external_predictions_detection_adjusted.csv.gz"))
fwrite(rbindlist(list(benchmark_full$metadata, benchmark_adjusted$metadata)),
       file.path(OUT_DIR, "tables", "fold_model_metadata.csv"))

# Overall performance.
overall_full <- benchmark_full$predictions[, all_metrics(.SD), by = .(predictor_set, model)]
overall_adjusted <- benchmark_adjusted$predictions[, all_metrics(.SD), by = .(predictor_set, model)]
overall <- rbindlist(list(overall_full, overall_adjusted))
fwrite(overall, file.path(OUT_DIR, "tables", "overall_model_performance.csv"))

# S11. Detection-adjusted changes --------------------------------------------
full_key <- overall_full[, .(model, roc_auc, pr_auc, brier, sensitivity,
                             specificity, ppv, f1, calibration_slope)]
adj_key <- overall_adjusted[, .(model, roc_auc, pr_auc, brier, sensitivity,
                                specificity, ppv, f1, calibration_slope)]
setnames(full_key, setdiff(names(full_key), "model"),
         paste0(setdiff(names(full_key), "model"), "_full"))
setnames(adj_key, setdiff(names(adj_key), "model"),
         paste0(setdiff(names(adj_key), "model"), "_adjusted"))
delta <- merge(full_key, adj_key, by = "model")
for (m in c("roc_auc", "pr_auc", "brier", "sensitivity", "specificity",
            "ppv", "f1", "calibration_slope")) {
  delta[, (paste0("delta_", m)) := get(paste0(m, "_adjusted")) - get(paste0(m, "_full"))]
}
fwrite(delta, file.path(OUT_DIR, "tables", "detection_adjustment_changes.csv"))

# S12. Geographic and subgroup summaries ------------------------------------
country_perf <- benchmark_full$predictions[, all_metrics(.SD),
  by = .(country, region15, model, external_fold)]
fwrite(country_perf, file.path(OUT_DIR, "tables", "country_external_performance.csv"))

subregion_perf <- benchmark_full$predictions[, all_metrics(.SD),
  by = .(region15, model)]
fwrite(subregion_perf, file.path(OUT_DIR, "tables", "subregion_external_performance.csv"))

subgroups <- rbindlist(list(
  benchmark_full$predictions[, all_metrics(.SD), by = .(model, subgroup_variable = "gender", subgroup = gender)],
  benchmark_full$predictions[, all_metrics(.SD), by = .(model, subgroup_variable = "age_group", subgroup = age_group)],
  benchmark_full$predictions[, all_metrics(.SD), by = .(model, subgroup_variable = "residence", subgroup = residence)]
), fill = TRUE)
fwrite(subgroups, file.path(OUT_DIR, "tables", "subgroup_external_performance.csv"))

# S13. Earlier-to-later temporal benchmarking --------------------------------
run_temporal_pair <- function(country_name, train_year, test_year,
                              predictor_set = "Full") {
  cat_vars <- if (predictor_set == "Full") categorical_full else categorical_adjusted
  d_country <- dat[country == country_name]
  dev <- d_country[survyear == train_year]
  test <- d_country[survyear == test_year]
  if (!nrow(dev) || !nrow(test)) stop("Temporal pair not found: ", country_name)

  set.seed(SEED + train_year)
  case_ids <- which(dev$y == 1L); ctrl_ids <- which(dev$y == 0L)
  cal_idx <- sort(c(
    sample(case_ids, size = max(1L, floor(0.15 * length(case_ids)))),
    sample(ctrl_ids, size = max(1L, floor(0.15 * length(ctrl_ids))))
  ))
  calibration <- dev[cal_idx]
  train <- dev[-cal_idx]
  pp <- fit_preprocessor(train, numeric_predictors, cat_vars)

  out <- list()
  for (model_name in c("Logistic", "ElasticNet", "RandomForest", "XGBoost")) {
    scale_numeric <- model_name %in% c("Logistic", "ElasticNet")
    x_train <- apply_preprocessor(train, pp, scale_numeric)
    x_cal <- apply_preprocessor(calibration, pp, scale_numeric)
    x_test <- apply_preprocessor(test, pp, scale_numeric)
    model <- fit_candidate_model(model_name, x_train, train$y, make_training_weights(train))
    raw_cal <- predict_candidate_model(model, x_cal)
    raw_test <- predict_candidate_model(model, x_test)
    platt <- fit_platt(calibration$y, raw_cal, calibration$eval_w)
    p_cal <- apply_platt(platt, raw_cal)
    p_test <- apply_platt(platt, raw_test)
    threshold <- select_threshold(calibration$y, p_cal, calibration$eval_w)
    temp <- test[, .(record_id, y, eval_w)]
    temp[, `:=`(prob = p_test, threshold = threshold)]
    m <- all_metrics(temp)
    m[, `:=`(
      country = country_name, train_year = train_year, test_year = test_year,
      predictor_set = predictor_set, model = model_name
    )]
    out[[model_name]] <- m
  }
  rbindlist(out, use.names = TRUE, fill = TRUE)
}

temporal_results <- rbindlist(lapply(seq_len(nrow(temporal_pairs)), function(i) {
  x <- temporal_pairs[i]
  run_temporal_pair(x$country, x$train_year, x$test_year, "Full")
}))
fwrite(temporal_results, file.path(OUT_DIR, "tables", "temporal_benchmarking.csv"))

# S14. XGBoost SHAP importance and selected figures --------------------------
# Fit one descriptive global XGBoost model on the full 46-country cohort.
# This model is used only for global feature attribution, not for external
# performance estimation.
pp_global <- fit_preprocessor(dat, numeric_predictors, categorical_full)
x_global <- apply_preprocessor(dat, pp_global, standardize_numeric = FALSE)
model_global <- fit_candidate_model(
  "XGBoost", x_global, dat$y, make_training_weights(dat)
)

shap_matrix <- predict(
  model_global$object, xgb.DMatrix(x_global), predcontrib = TRUE
)
shap_matrix <- shap_matrix[, colnames(shap_matrix) != "BIAS", drop = FALSE]
shap_importance <- data.table(
  encoded_feature = colnames(shap_matrix),
  mean_abs_shap = colMeans(abs(shap_matrix))
)
shap_importance[, source_feature := sub("__.*$", "", encoded_feature)]
shap_global <- shap_importance[, .(mean_abs_shap = sum(mean_abs_shap)), by = source_feature]
setorder(shap_global, -mean_abs_shap)
fwrite(shap_global, file.path(OUT_DIR, "tables", "xgboost_shap_global_importance.csv"))

p_shap <- ggplot(head(shap_global, 15), aes(x = reorder(source_feature, mean_abs_shap), y = mean_abs_shap)) +
  geom_col() + coord_flip() +
  labs(x = NULL, y = "Mean absolute SHAP value", title = "Global XGBoost feature importance") +
  theme_minimal(base_size = 11)
ggsave(file.path(OUT_DIR, "figures", "xgboost_shap_importance.png"), p_shap,
       width = 8, height = 6, dpi = 600)

# Subregional ROC AUC heatmap.
heat <- subregion_perf[, .(region15, model, roc_auc)]
p_heat <- ggplot(heat, aes(x = model, y = region15, fill = roc_auc)) +
  geom_tile() + geom_text(aes(label = sprintf("%.2f", roc_auc)), size = 3) +
  labs(x = NULL, y = NULL, fill = "ROC AUC",
       title = "Subregional transportability of externally evaluated models") +
  theme_minimal(base_size = 10) +
  theme(axis.text.x = element_text(angle = 30, hjust = 1))
ggsave(file.path(OUT_DIR, "figures", "subregion_transportability_roc_auc.png"), p_heat,
       width = 10, height = 7, dpi = 600)

# S15. Reproducibility outputs -------------------------------------------------
writeLines(capture.output(sessionInfo()), file.path(OUT_DIR, "R_sessionInfo.txt"))
writeLines(
  c(
    paste0("Random seed: ", SEED),
    paste0("Run date: ", Sys.Date()),
    paste0("Analysis rows: ", nrow(dat)),
    paste0("Countries: ", uniqueN(dat$country)),
    paste0("Subregions: ", uniqueN(dat$region15)),
    paste0("CVD cases: ", sum(dat$y)),
    "Temporal pairs: Mongolia 2013 to 2019; Uganda 2014 to 2023; Viet Nam 2015 to 2021"
  ),
  file.path(OUT_DIR, "run_summary.txt")
)

message("Analysis complete. Outputs are in: ", normalizePath(OUT_DIR))