# ==============================================================================
# Supplementary Code
# Development, Internal Validation, External Validation, and SHAP Interpretation
# of Machine-Learning Models for STBs Identification
#
# This script is a cleaned and reordered version of the analysis code used in
# the study. Local file paths and report-generation code have been removed.
# Raw variable names are retained to preserve compatibility with the study data.
#
# IMPORTANT:
# 1. The LASSO feature-selection code was not contained in the uploaded source
#    code file and is therefore not reproduced here.
# 2. The 22 predictors below are the predictors entering model development after
#    feature selection, as reported in the manuscript.
# 3. Data files are not included in this script.
# ==============================================================================

# ------------------------------------------------------------------------------
# 1. Packages and reproducibility
# ------------------------------------------------------------------------------

required_packages <- c(
  "rio", "dplyr", "pROC", "PRROC",
  "randomForest", "xgboost", "lightgbm",
  "ggplot2", "foreach", "doParallel", "tidyr", "patchwork"
)

missing_packages <- required_packages[
  !vapply(required_packages, requireNamespace, logical(1), quietly = TRUE)
]

if (length(missing_packages) > 0) {
  stop(
    "Please install the following packages before running this script: ",
    paste(missing_packages, collapse = ", ")
  )
}

library(rio)
library(dplyr)
library(pROC)
library(PRROC)
library(randomForest)
library(xgboost)
library(lightgbm)
library(ggplot2)
library(foreach)
library(doParallel)
library(tidyr)
library(patchwork)

set.seed(123)


# ------------------------------------------------------------------------------
# 2. Study variables
# ------------------------------------------------------------------------------

# Raw data field used to encode the study outcome (STBs).
outcome_var <- "Child suicidal Ideation"

# Final 22 predictors retained for model development.
predictors <- c(
  "Recent Child Anxiety",
  "Family History of Suicide",
  "Caregiver Education Level",
  "Age Group",
  "Somatization",
  "Parental Anxiety",
  "Family Conflict",
  "Social issues",
  "Weekday Sleep Duration",
  "Mood Disorders",
  "ODD",
  "Bladder Bowel Dysfunction",
  "Daytime sleepiness",
  "Family Cohesion",
  "Night Awakening",
  "Weekend Sleep Duration",
  "Birth Weight",
  "Gestational Age",
  "Daily Secondhand Smoke Exposure",
  "Parental Marital Status",
  "Sleep maintenance",
  "Sleep-Disordered Breathing"
)

model_order <- c("Logistic", "Random Forest", "XGBoost", "LightGBM")


# ------------------------------------------------------------------------------
# 3. Data preprocessing
# ------------------------------------------------------------------------------

prepare_dataset <- function(data, outcome_var, predictors) {

  missing_vars <- setdiff(c(outcome_var, predictors), names(data))
  if (length(missing_vars) > 0) {
    stop(
      "Required variables are missing: ",
      paste(missing_vars, collapse = ", ")
    )
  }

  data[[outcome_var]] <- trimws(as.character(data[[outcome_var]]))

  # Raw database coding:
  # "Risks_exist" = STBs; "Normal" = non-STBs.
  data[[outcome_var]] <- dplyr::case_when(
    data[[outcome_var]] %in%
      c("Risks exist", "Risks_exist", "存在自杀风险", "高风险", "1") ~
      "Risks_exist",
    data[[outcome_var]] %in%
      c("Normal", "未见明显异常", "正常", "0") ~
      "Normal",
    TRUE ~ NA_character_
  )

  data[[outcome_var]] <- factor(
    data[[outcome_var]],
    levels = c("Risks_exist", "Normal")
  )

  for (v in predictors) {
    data[[v]] <- factor(data[[v]])
  }

  model_data <- data[, c(outcome_var, predictors), drop = FALSE]
  model_data <- na.omit(model_data)

  if (length(unique(model_data[[outcome_var]])) < 2) {
    stop("The outcome contains fewer than two classes.")
  }

  model_data
}


# ------------------------------------------------------------------------------
# 4. Development cohort: stratified 70:30 split
# ------------------------------------------------------------------------------

development_data <- rio::import("Chongqing_data.xlsx")
model_data <- prepare_dataset(development_data, outcome_var, predictors)

set.seed(123)

positive_data <- model_data %>%
  filter(.data[[outcome_var]] == "Risks_exist")

negative_data <- model_data %>%
  filter(.data[[outcome_var]] == "Normal")

positive_train_id <- sample(
  seq_len(nrow(positive_data)),
  size = floor(0.70 * nrow(positive_data))
)

negative_train_id <- sample(
  seq_len(nrow(negative_data)),
  size = floor(0.70 * nrow(negative_data))
)

train_data <- bind_rows(
  positive_data[positive_train_id, , drop = FALSE],
  negative_data[negative_train_id, , drop = FALSE]
)

valid_data <- bind_rows(
  positive_data[-positive_train_id, , drop = FALSE],
  negative_data[-negative_train_id, , drop = FALSE]
)

train_data <- train_data[sample(seq_len(nrow(train_data))), , drop = FALSE]
valid_data <- valid_data[sample(seq_len(nrow(valid_data))), , drop = FALSE]


# ------------------------------------------------------------------------------
# 5. Design matrices
# ------------------------------------------------------------------------------

x_train <- model.matrix(
  as.formula(paste0("`", outcome_var, "` ~ .")),
  data = train_data
)[, -1, drop = FALSE]

x_valid_raw <- model.matrix(
  as.formula(paste0("`", outcome_var, "` ~ .")),
  data = valid_data
)[, -1, drop = FALSE]

colnames(x_train) <- make.names(colnames(x_train), unique = TRUE)
colnames(x_valid_raw) <- make.names(colnames(x_valid_raw), unique = TRUE)

train_cols <- colnames(x_train)

align_matrix <- function(x_raw, reference_cols) {

  x_aligned <- matrix(
    0,
    nrow = nrow(x_raw),
    ncol = length(reference_cols),
    dimnames = list(NULL, reference_cols)
  )

  common_cols <- intersect(colnames(x_raw), reference_cols)

  x_aligned[, common_cols] <-
    x_raw[, common_cols, drop = FALSE]

  x_aligned
}

x_valid <- align_matrix(x_valid_raw, train_cols)

y_train_factor <- train_data[[outcome_var]]
y_valid_factor <- valid_data[[outcome_var]]

y_train_num <- ifelse(y_train_factor == "Risks_exist", 1, 0)
y_valid_num <- ifelse(y_valid_factor == "Risks_exist", 1, 0)


# ------------------------------------------------------------------------------
# 6. Performance metrics
# ------------------------------------------------------------------------------

get_metrics <- function(y_true, prob) {

  roc_obj <- pROC::roc(y_true, prob, quiet = TRUE)
  auc_value <- as.numeric(pROC::auc(roc_obj))

  pr_obj <- PRROC::pr.curve(
    scores.class0 = prob[y_true == 1],
    scores.class1 = prob[y_true == 0],
    curve = FALSE
  )

  best_cut <- as.numeric(
    pROC::coords(
      roc_obj,
      x = "best",
      best.method = "youden",
      ret = "threshold"
    )[1]
  )

  pred <- ifelse(prob >= best_cut, 1, 0)

  TP <- sum(pred == 1 & y_true == 1)
  FP <- sum(pred == 1 & y_true == 0)
  TN <- sum(pred == 0 & y_true == 0)
  FN <- sum(pred == 0 & y_true == 1)

  data.frame(
    AUROC = auc_value,
    AUPRC = pr_obj$auc.integral,
    Best_Threshold = best_cut,
    Accuracy = (TP + TN) / (TP + FP + TN + FN),
    Sensitivity = TP / (TP + FN),
    Specificity = TN / (TN + FP),
    PPV = TP / (TP + FP),
    NPV = TN / (TN + FN),
    F1 = 2 * TP / (2 * TP + FP + FN),
    TP = TP,
    FP = FP,
    TN = TN,
    FN = FN
  )
}


# ------------------------------------------------------------------------------
# 7. Model development
# ------------------------------------------------------------------------------

# 7.1 Logistic regression
x_train_logit <- cbind(Intercept = 1, x_train)
x_valid_logit <- cbind(Intercept = 1, x_valid)

logit_model <- glm.fit(
  x = x_train_logit,
  y = y_train_num,
  family = binomial()
)

logit_prob_train <- as.numeric(
  plogis(x_train_logit %*% coef(logit_model))
)

logit_prob_valid <- as.numeric(
  plogis(x_valid_logit %*% coef(logit_model))
)


# 7.2 Random Forest
rf_model <- randomForest(
  x = x_train,
  y = y_train_factor,
  ntree = 500,
  importance = TRUE
)

rf_prob_train <- predict(
  rf_model,
  newdata = x_train,
  type = "prob"
)[, "Risks_exist"]

rf_prob_valid <- predict(
  rf_model,
  newdata = x_valid,
  type = "prob"
)[, "Risks_exist"]


# 7.3 XGBoost
dtrain <- xgboost::xgb.DMatrix(
  data = x_train,
  label = y_train_num
)

dvalid <- xgboost::xgb.DMatrix(
  data = x_valid,
  label = y_valid_num
)

xgb_params <- list(
  objective = "binary:logistic",
  eval_metric = "auc",
  max_depth = 3,
  eta = 0.05,
  subsample = 0.8,
  colsample_bytree = 0.8
)

xgb_model <- xgboost::xgb.train(
  params = xgb_params,
  data = dtrain,
  nrounds = 200,
  verbose = 0
)

xgb_prob_train <- predict(xgb_model, newdata = dtrain)
xgb_prob_valid <- predict(xgb_model, newdata = dvalid)


# 7.4 LightGBM
lgb_train <- lightgbm::lgb.Dataset(
  data = x_train,
  label = y_train_num
)

lgb_params <- list(
  objective = "binary",
  metric = "auc",
  learning_rate = 0.05,
  num_leaves = 15,
  feature_fraction = 0.8,
  bagging_fraction = 0.8,
  bagging_freq = 1,
  verbosity = -1
)

lgb_model <- lightgbm::lgb.train(
  params = lgb_params,
  data = lgb_train,
  nrounds = 200,
  verbose = -1
)

lgb_prob_train <- predict(lgb_model, newdata = x_train)
lgb_prob_valid <- predict(lgb_model, newdata = x_valid)


# ------------------------------------------------------------------------------
# 8. Training and internal-validation performance
# ------------------------------------------------------------------------------

train_summary <- bind_rows(
  cbind(Model = "Logistic",
        get_metrics(y_train_num, logit_prob_train)),
  cbind(Model = "Random Forest",
        get_metrics(y_train_num, rf_prob_train)),
  cbind(Model = "XGBoost",
        get_metrics(y_train_num, xgb_prob_train)),
  cbind(Model = "LightGBM",
        get_metrics(y_train_num, lgb_prob_train))
)

valid_summary <- bind_rows(
  cbind(Model = "Logistic",
        get_metrics(y_valid_num, logit_prob_valid)),
  cbind(Model = "Random Forest",
        get_metrics(y_valid_num, rf_prob_valid)),
  cbind(Model = "XGBoost",
        get_metrics(y_valid_num, xgb_prob_valid)),
  cbind(Model = "LightGBM",
        get_metrics(y_valid_num, lgb_prob_valid))
)

print(train_summary)
print(valid_summary)


# ------------------------------------------------------------------------------
# 9. ROC, precision-recall, calibration, and decision-curve functions
# ------------------------------------------------------------------------------

make_roc_data <- function(y, probs) {

  bind_rows(lapply(names(probs), function(m) {

    roc_obj <- pROC::roc(y, probs[[m]], quiet = TRUE)

    data.frame(
      FPR = 1 - roc_obj$specificities,
      TPR = roc_obj$sensitivities,
      Model = m
    )
  }))
}

make_pr_data <- function(y, probs) {

  bind_rows(lapply(names(probs), function(m) {

    pr_obj <- PRROC::pr.curve(
      scores.class0 = probs[[m]][y == 1],
      scores.class1 = probs[[m]][y == 0],
      curve = TRUE
    )

    data.frame(
      Recall = pr_obj$curve[, 1],
      Precision = pr_obj$curve[, 2],
      Model = m
    )
  }))
}

make_calibration_data <- function(y, probs) {

  bind_rows(lapply(names(probs), function(m) {

    temp <- data.frame(y = y, prob = probs[[m]])

    loess_fit <- loess(y ~ prob, data = temp, span = 0.75)

    pred_grid <- seq(
      min(temp$prob, na.rm = TRUE),
      max(temp$prob, na.rm = TRUE),
      length.out = 100
    )

    pred_res <- predict(
      loess_fit,
      newdata = data.frame(prob = pred_grid),
      se = TRUE
    )

    data.frame(
      Predicted = pred_grid,
      Observed = pmax(pmin(pred_res$fit, 1), 0),
      Lower = pmax(pmin(pred_res$fit - 1.96 * pred_res$se.fit, 1), 0),
      Upper = pmax(pmin(pred_res$fit + 1.96 * pred_res$se.fit, 1), 0),
      Model = m
    )
  }))
}

calculate_net_benefit <- function(pred, y, threshold) {

  tp <- sum(pred >= threshold & y == 1)
  fp <- sum(pred >= threshold & y == 0)
  n <- length(y)

  tp / n - fp / n * threshold / (1 - threshold)
}

make_dca_data <- function(y, probs,
                          thresholds = seq(0.001, 0.20, by = 0.002)) {

  model_data <- bind_rows(lapply(names(probs), function(m) {

    data.frame(
      Threshold = thresholds,
      NetBenefit = vapply(
        thresholds,
        function(t) calculate_net_benefit(probs[[m]], y, t),
        numeric(1)
      ),
      Model = m
    )
  }))

  prevalence <- mean(y == 1)

  treat_all <- data.frame(
    Threshold = thresholds,
    NetBenefit =
      prevalence -
      thresholds * (1 - prevalence) / (1 - thresholds),
    Model = "Treat All"
  )

  treat_none <- data.frame(
    Threshold = thresholds,
    NetBenefit = 0,
    Model = "Treat None"
  )

  bind_rows(model_data, treat_all, treat_none)
}

train_probs <- list(
  Logistic = logit_prob_train,
  `Random Forest` = rf_prob_train,
  XGBoost = xgb_prob_train,
  LightGBM = lgb_prob_train
)

valid_probs <- list(
  Logistic = logit_prob_valid,
  `Random Forest` = rf_prob_valid,
  XGBoost = xgb_prob_valid,
  LightGBM = lgb_prob_valid
)

train_roc <- make_roc_data(y_train_num, train_probs)
train_pr <- make_pr_data(y_train_num, train_probs)
train_calibration <- make_calibration_data(y_train_num, train_probs)
train_dca <- make_dca_data(y_train_num, train_probs)

valid_roc <- make_roc_data(y_valid_num, valid_probs)
valid_pr <- make_pr_data(y_valid_num, valid_probs)
valid_calibration <- make_calibration_data(y_valid_num, valid_probs)
valid_dca <- make_dca_data(y_valid_num, valid_probs)


# ------------------------------------------------------------------------------
# 10. Independent external validation
# ------------------------------------------------------------------------------

external_validate <- function(file_name) {

  ext_data <- rio::import(file_name)
  ext_model_data <- prepare_dataset(ext_data, outcome_var, predictors)

  y_ext_factor <- ext_model_data[[outcome_var]]
  y_ext_num <- ifelse(y_ext_factor == "Risks_exist", 1, 0)

  x_ext_raw <- model.matrix(
    as.formula(paste0("`", outcome_var, "` ~ .")),
    data = ext_model_data
  )[, -1, drop = FALSE]

  colnames(x_ext_raw) <-
    make.names(colnames(x_ext_raw), unique = TRUE)

  x_ext <- align_matrix(x_ext_raw, train_cols)

  x_ext_logit <- cbind(Intercept = 1, x_ext)

  logit_prob_ext <- as.numeric(
    plogis(x_ext_logit %*% coef(logit_model))
  )

  rf_prob_ext <- predict(
    rf_model,
    newdata = x_ext,
    type = "prob"
  )[, "Risks_exist"]

  xgb_prob_ext <- predict(
    xgb_model,
    newdata = xgboost::xgb.DMatrix(data = x_ext)
  )

  lgb_prob_ext <- predict(
    lgb_model,
    newdata = x_ext
  )

  ext_probs <- list(
    Logistic = logit_prob_ext,
    `Random Forest` = rf_prob_ext,
    XGBoost = xgb_prob_ext,
    LightGBM = lgb_prob_ext
  )

  performance <- bind_rows(
    cbind(Model = "Logistic",
          get_metrics(y_ext_num, logit_prob_ext)),
    cbind(Model = "Random Forest",
          get_metrics(y_ext_num, rf_prob_ext)),
    cbind(Model = "XGBoost",
          get_metrics(y_ext_num, xgb_prob_ext)),
    cbind(Model = "LightGBM",
          get_metrics(y_ext_num, lgb_prob_ext))
  )

  list(
    performance = performance,
    roc = make_roc_data(y_ext_num, ext_probs),
    pr = make_pr_data(y_ext_num, ext_probs),
    calibration = make_calibration_data(y_ext_num, ext_probs),
    dca = make_dca_data(y_ext_num, ext_probs),
    probabilities = ext_probs,
    outcome = y_ext_num
  )
}

shaanxi_validation <- external_validate("Shaanxi_data.xlsx")
yunnan_validation <- external_validate("Yunnan_data.xlsx")

print(shaanxi_validation$performance)
print(yunnan_validation$performance)


# ------------------------------------------------------------------------------
# 11. SHAP analysis: sampling strategy
# ------------------------------------------------------------------------------

# SHAP calculations were based on a random sample of 300 internal-validation
# participants and 100 training participants as the background sample.

set.seed(123)

x_val_mat <- as.matrix(x_valid)
x_train_mat <- as.matrix(x_train)

if (nrow(x_val_mat) > 300) {
  shap_id <- sample(seq_len(nrow(x_val_mat)), 300)
  x_shap <- x_val_mat[shap_id, , drop = FALSE]
} else {
  shap_id <- seq_len(nrow(x_val_mat))
  x_shap <- x_val_mat
}

if (nrow(x_train_mat) > 100) {
  bg_id <- sample(seq_len(nrow(x_train_mat)), 100)
  x_bg <- x_train_mat[bg_id, , drop = FALSE]
} else {
  bg_id <- seq_len(nrow(x_train_mat))
  x_bg <- x_train_mat
}


# ------------------------------------------------------------------------------
# 12. SHAP analysis: Logistic Regression and Random Forest
# ------------------------------------------------------------------------------

pred_logit <- function(newdata, model) {

  beta <- coef(model)
  x <- as.matrix(newdata)

  common <- intersect(
    colnames(x),
    names(beta)[-1]
  )

  lp <- beta[1] +
    x[, common, drop = FALSE] %*%
    beta[common]

  as.numeric(plogis(lp))
}

pred_rf <- function(newdata, model) {

  predict(
    model,
    newdata = as.matrix(newdata),
    type = "prob"
  )[, "Risks_exist"]
}

parallel_shap <- function(
    X,
    background,
    model,
    pred_fun,
    nsim = 30) {

  X <- as.matrix(X)
  background <- as.matrix(background)

  p <- ncol(X)
  n <- nrow(X)
  feature_names <- colnames(X)

  shap_list <- foreach(
    i = seq_len(n),
    .combine = rbind,
    .packages = c("randomForest")
  ) %dopar% {

    set.seed(123 + i)

    x_i <- X[i, , drop = FALSE]
    shap_i <- rep(0, p)

    for (s in seq_len(nsim)) {

      bg <- background[
        sample(seq_len(nrow(background)), 1),
        ,
        drop = FALSE
      ]

      order_j <- sample(seq_len(p))
      x_old <- bg
      pred_old <- pred_fun(x_old, model)

      for (j in order_j) {

        x_new <- x_old
        x_new[, j] <- x_i[, j]

        pred_new <- pred_fun(x_new, model)

        shap_i[j] <-
          shap_i[j] + as.numeric(pred_new - pred_old)

        x_old <- x_new
        pred_old <- pred_new
      }
    }

    shap_i / nsim
  }

  shap_mat <- as.matrix(shap_list)
  colnames(shap_mat) <- feature_names

  shap_mat
}

n_cores <- max(1, parallel::detectCores() - 1)
cl <- parallel::makeCluster(n_cores)
doParallel::registerDoParallel(cl)

shap_logit <- parallel_shap(
  X = x_shap,
  background = x_bg,
  model = logit_model,
  pred_fun = pred_logit,
  nsim = 30
)

shap_rf <- parallel_shap(
  X = x_shap,
  background = x_bg,
  model = rf_model,
  pred_fun = pred_rf,
  nsim = 30
)

parallel::stopCluster(cl)


# ------------------------------------------------------------------------------
# 13. SHAP analysis: XGBoost and LightGBM
# ------------------------------------------------------------------------------

shap_xgb <- predict(
  xgb_model,
  newdata = xgboost::xgb.DMatrix(x_shap),
  predcontrib = TRUE
)

shap_xgb <- as.matrix(shap_xgb)
shap_xgb <- shap_xgb[, -ncol(shap_xgb), drop = FALSE]
colnames(shap_xgb) <- colnames(x_shap)

shap_lgb <- predict(
  lgb_model,
  newdata = x_shap,
  type = "contrib"
)

shap_lgb <- as.matrix(shap_lgb)
shap_lgb <- shap_lgb[, -ncol(shap_lgb), drop = FALSE]
colnames(shap_lgb) <- colnames(x_shap)


# ------------------------------------------------------------------------------
# 14. Aggregate dummy-coded SHAP values to the 22 original predictors
# ------------------------------------------------------------------------------

aggregate_SHAP_matrix <- function(shap_mat, original_vars) {

  shap_mat <- as.matrix(shap_mat)
  feature_names <- colnames(shap_mat)

  result <- matrix(
    0,
    nrow = nrow(shap_mat),
    ncol = length(original_vars),
    dimnames = list(NULL, original_vars)
  )

  for (v in original_vars) {

    cols <- feature_names[
      grepl(make.names(v), feature_names, fixed = TRUE)
    ]

    if (length(cols) == 0) {
      cols <- feature_names[
        grepl(v, feature_names, fixed = TRUE)
      ]
    }

    if (length(cols) > 0) {
      result[, v] <-
        rowSums(shap_mat[, cols, drop = FALSE], na.rm = TRUE)
    }
  }

  result
}

shap22_list <- list(
  Logistic = aggregate_SHAP_matrix(shap_logit, predictors),
  `Random Forest` = aggregate_SHAP_matrix(shap_rf, predictors),
  XGBoost = aggregate_SHAP_matrix(shap_xgb, predictors),
  LightGBM = aggregate_SHAP_matrix(shap_lgb, predictors)
)

get_shap_importance <- function(shap_mat, model_name) {

  data.frame(
    Variable = colnames(shap_mat),
    Mean_ABS_SHAP =
      colMeans(abs(as.matrix(shap_mat)), na.rm = TRUE),
    Model = model_name
  ) %>%
    arrange(desc(Mean_ABS_SHAP))
}

shap_importance <- bind_rows(
  get_shap_importance(shap22_list$Logistic, "Logistic"),
  get_shap_importance(shap22_list$`Random Forest`, "Random Forest"),
  get_shap_importance(shap22_list$XGBoost, "XGBoost"),
  get_shap_importance(shap22_list$LightGBM, "LightGBM")
)

print(shap_importance)


# ------------------------------------------------------------------------------
# 15. SHAP summary/beeswarm data
# ------------------------------------------------------------------------------

# Use the exact validation participants sampled above; do not re-sample here.
x_valid_shap <- x_valid[shap_id, , drop = FALSE]

make_feature22 <- function(x, vars) {

  result <- matrix(
    0,
    nrow = nrow(x),
    ncol = length(vars),
    dimnames = list(NULL, vars)
  )

  features <- colnames(x)

  for (v in vars) {

    cols <- features[
      grepl(make.names(v), features, fixed = TRUE)
    ]

    if (length(cols) == 0) {
      cols <- features[
        grepl(v, features, fixed = TRUE)
      ]
    }

    if (length(cols) > 0) {
      result[, v] <-
        rowSums(x[, cols, drop = FALSE], na.rm = TRUE)
    }
  }

  as.data.frame(result)
}

feature22 <- make_feature22(x_valid_shap, predictors)

make_beeswarm_data <- function(shap22, model_name) {

  bind_rows(lapply(colnames(shap22), function(v) {

    data.frame(
      Variable = v,
      SHAP = shap22[, v],
      Feature = feature22[, v],
      Model = model_name
    )
  }))
}

beeswarm_data <- bind_rows(
  make_beeswarm_data(shap22_list$Logistic, "Logistic"),
  make_beeswarm_data(shap22_list$`Random Forest`, "Random Forest"),
  make_beeswarm_data(shap22_list$XGBoost, "XGBoost"),
  make_beeswarm_data(shap22_list$LightGBM, "LightGBM")
)


# ------------------------------------------------------------------------------
# 16. Logistic-regression SHAP heatmap data
# ------------------------------------------------------------------------------

logit_order <- shap_importance %>%
  filter(Model == "Logistic") %>%
  arrange(desc(Mean_ABS_SHAP)) %>%
  pull(Variable)

shap_logit22 <- shap22_list$Logistic[, logit_order, drop = FALSE]

# Order participants by total absolute SHAP contribution for visualization.
individual_order <- order(
  rowSums(abs(shap_logit22), na.rm = TRUE)
)

shap_logit22 <- shap_logit22[individual_order, , drop = FALSE]

heatmap_data <- as.data.frame(shap_logit22)
heatmap_data$Individual <- seq_len(nrow(heatmap_data))

heatmap_data <- heatmap_data %>%
  tidyr::pivot_longer(
    cols = -Individual,
    names_to = "Variable",
    values_to = "SHAP"
  )

heatmap_data$Variable <- factor(
  heatmap_data$Variable,
  levels = rev(logit_order)
)


# ------------------------------------------------------------------------------
# 17. Optional plotting examples for supplementary figures
# ------------------------------------------------------------------------------

# Global SHAP importance
p_shap_importance <- ggplot(
  shap_importance,
  aes(
    x = Mean_ABS_SHAP,
    y = reorder(Variable, Mean_ABS_SHAP)
  )
) +
  geom_col() +
  coord_flip() +
  facet_wrap(~Model, scales = "free_y") +
  theme_bw() +
  labs(
    x = "Mean |SHAP value|",
    y = NULL
  )

# SHAP summary plot
p_beeswarm <- ggplot(
  beeswarm_data,
  aes(
    x = SHAP,
    y = Variable,
    color = Feature
  )
) +
  geom_jitter(height = 0.18, size = 1.2, alpha = 0.55) +
  scale_color_gradient(
    low = "#2166AC",
    high = "#B2182B",
    name = "Feature value"
  ) +
  facet_wrap(~Model, ncol = 2) +
  theme_classic() +
  labs(
    x = "SHAP value",
    y = NULL
  )

# Logistic-regression individual-level SHAP heatmap
p_heatmap <- ggplot(
  heatmap_data,
  aes(
    x = Individual,
    y = Variable,
    fill = SHAP
  )
) +
  geom_tile() +
  scale_fill_gradient2(
    low = "#2166AC",
    mid = "white",
    high = "#B2182B",
    midpoint = 0,
    name = "SHAP"
  ) +
  theme_bw() +
  labs(
    x = "Individuals",
    y = NULL
  )


# ==============================================================================
# End of supplementary code
# ==============================================================================