# ==============================================================================
# Supplementary Code S1: 
# Machine Learning Modeling and SHAP Visualization for Predicting 
# Oxaliplatin-Induced Dose-Limiting Toxicity (DLT)
#
# ==============================================================================

# ------------------------------------------------------------------------------
# 1. Environment Setup & Package Loading
# ------------------------------------------------------------------------------
# List of required packages
required_packages <- c("caret", "xgboost", "randomForest", "e1071", "pROC", 
                       "PRROC", "shapviz", "ggplot2", "dplyr", "tidyr")

# Install missing packages automatically
for(pkg in required_packages) {
  if(!requireNamespace(pkg, quietly = TRUE)) install.packages(pkg)
}

library(caret)
library(xgboost)
library(randomForest)
library(e1071)
library(pROC)
library(PRROC)
library(shapviz)
library(ggplot2)
library(dplyr)
library(tidyr)

# Set working directory to the location of your dataset
# setwd("path/to/your/directory") 
data_path <- "crc_dpn_comorbidity_data.csv" # Ensure this file is in your working directory

# ------------------------------------------------------------------------------
# 2. Data Loading & Initial Formatting
# ------------------------------------------------------------------------------
cat("Loading dataset...\n")
df <- read.csv(data_path, stringsAsFactors = FALSE)

# Format target variable as valid R factor names for caret (0 -> NonDLT, 1 -> DLT)
df$DLT <- factor(df$DLT, levels = c(0, 1), labels = c("NonDLT", "DLT"))

# Remove variables with >10% missing values (Simulated rule as per Methods)
missing_rates <- colMeans(is.na(df))
df <- df[, missing_rates <= 0.10]

# ------------------------------------------------------------------------------
# 3. Reproducible Data Partitioning
# ------------------------------------------------------------------------------
cat("Partitioning dataset (75% Training, 25% Testing)...\n")
# Strict random seed for reproducible split
set.seed(202405) 
train_idx <- caret::createDataPartition(df$DLT, p = 0.75, list = FALSE)

train_data <- df[train_idx, ]
test_data  <- df[-train_idx, ]

# ------------------------------------------------------------------------------
# 4. Data Preprocessing (Strictly Post-Split to Prevent Data Leakage)
# ------------------------------------------------------------------------------
cat("Handling missing values and feature scaling (post-split)...\n")

# 4.1 Missing Value Imputation using Training Set Medians
preProc_impute <- caret::preProcess(train_data, method = "medianImpute")
train_data_imp <- predict(preProc_impute, train_data)
test_data_imp  <- predict(preProc_impute, test_data) # Apply train-derived medians to test

# 4.2 Feature Scaling (Required for SVM, KNN, LR; XGBoost & RF use unscaled)
# We keep an unscaled version for tree models to maintain SHAP interpretability
preProc_scale <- caret::preProcess(train_data_imp, method = c("center", "scale"))
train_data_scaled <- predict(preProc_scale, train_data_imp)
test_data_scaled  <- predict(preProc_scale, test_data_imp)

# Prepare DMatrix for XGBoost
X_train_unscaled <- model.matrix(DLT ~ . - 1, data = train_data_imp)
X_test_unscaled  <- model.matrix(DLT ~ . - 1, data = test_data_imp)
Y_train_num <- ifelse(train_data_imp$DLT == "DLT", 1, 0)
Y_test_num  <- ifelse(test_data_imp$DLT == "DLT", 1, 0)

# ------------------------------------------------------------------------------
# 5. Hyperparameter Tuning via Grid Search & 5-Fold Cross-Validation
# ------------------------------------------------------------------------------
cat("Initiating 5-fold cross-validation and hyperparameter grid search...\n")

# Define robust CV control
cv_control <- trainControl(method = "cv", number = 5, 
                           classProbs = TRUE, summaryFunction = twoClassSummary)

# --- A. XGBoost ---
cat("Training XGBoost...\n")
xgb_grid <- expand.grid(nrounds = c(50, 100, 150),
                        max_depth = c(3, 4, 5),
                        eta = c(0.01, 0.05, 0.1),
                        gamma = 0,
                        colsample_bytree = 0.8,
                        min_child_weight = 1,
                        subsample = 0.8)

set.seed(202405)
xgb_tuned <- train(DLT ~ ., data = train_data_imp, method = "xgbTree",
                   trControl = cv_control, tuneGrid = xgb_grid, metric = "ROC")

# --- B. Random Forest ---
cat("Training Random Forest...\n")
rf_grid <- expand.grid(mtry = c(2, 4, 6, 8))
set.seed(202405)
rf_tuned <- train(DLT ~ ., data = train_data_imp, method = "rf",
                  trControl = cv_control, tuneGrid = rf_grid, metric = "ROC", ntree = 300)

# --- C. SVM (Uses scaled data) ---
cat("Training SVM...\n")
svm_grid <- expand.grid(sigma = c(0.01, 0.05, 0.1), C = c(0.1, 1, 10))
set.seed(202405)
svm_tuned <- train(DLT ~ ., data = train_data_scaled, method = "svmRadial",
                   trControl = cv_control, tuneGrid = svm_grid, metric = "ROC")

# --- D. KNN (Uses scaled data) ---
cat("Training KNN...\n")
knn_grid <- expand.grid(k = c(3, 5, 7, 9))
set.seed(202405)
knn_tuned <- train(DLT ~ ., data = train_data_scaled, method = "knn",
                   trControl = cv_control, tuneGrid = knn_grid, metric = "ROC")

# --- E. Logistic Regression (Uses scaled data) ---
cat("Training Logistic Regression...\n")
set.seed(202405)
lr_tuned <- train(DLT ~ ., data = train_data_scaled, method = "glm", family = "binomial",
                  trControl = cv_control, metric = "ROC")

# ------------------------------------------------------------------------------
# 6. Model Evaluation on Independent Testing Cohort
# ------------------------------------------------------------------------------
cat("Evaluating models on the testing cohort...\n")

preds_prob <- data.frame(Actual = test_data_imp$DLT, Actual_Num = Y_test_num)

preds_prob$XGBoost <- predict(xgb_tuned, test_data_imp, type = "prob")[,"DLT"]
preds_prob$RF      <- predict(rf_tuned, test_data_imp, type = "prob")[,"DLT"]
preds_prob$SVM     <- predict(svm_tuned, test_data_scaled, type = "prob")[,"DLT"]
preds_prob$KNN     <- predict(knn_tuned, test_data_scaled, type = "prob")[,"DLT"]
preds_prob$LR      <- predict(lr_tuned, test_data_scaled, type = "prob")[,"DLT"]

# Print Test Set AUCs
models <- c("XGBoost", "RF", "SVM", "KNN", "LR")
for(m in models) {
  roc_obj <- roc(preds_prob$Actual_Num, preds_prob[[m]], quiet = TRUE)
  cat(sprintf("%s Test AUC: %.3f\n", m, roc_obj$auc))
}

# ------------------------------------------------------------------------------
# 7. Decision Threshold Determination via Youden's Index (Training Set)
# ------------------------------------------------------------------------------
# To maintain rigor, the cutoff for the confusion matrix must be derived from the training set
train_xgb_preds <- predict(xgb_tuned, train_data_imp, type = "prob")[,"DLT"]
roc_train_xgb <- roc(Y_train_num, train_xgb_preds, quiet = TRUE)
best_coords <- coords(roc_train_xgb, "best", ret = "threshold", transpose = FALSE)
best_threshold <- best_coords[1, 1]

cat(sprintf("\nOptimal Decision Threshold (from Training ROC Youden's Index): %.3f\n", best_threshold))

# Apply threshold to Testing set for Confusion Matrix
test_pred_class <- ifelse(preds_prob$XGBoost >= best_threshold, 1, 0)
cm <- table(Predicted = factor(test_pred_class, levels=c(1,0)), 
            Actual = factor(Y_test_num, levels=c(1,0)))
print("Testing Cohort Confusion Matrix (XGBoost):")
print(cm)

# ------------------------------------------------------------------------------
# 8. SHAP Interpretability Analysis (Using the optimal XGBoost model)
# ------------------------------------------------------------------------------
cat("\nGenerating SHAP Global and Local Explanations...\n")

# Extract the final xgboost model object from caret
final_xgb <- xgb_tuned$finalModel

# Create SHAP object (using unscaled test matrix for realistic clinical interpretation)
# X_pred is the numerical matrix required for XGBoost calculation
# X is the data frame used to retain original labels (e.g., TNM stage characters)
X_test_display <- test_data_imp %>% select(-DLT)
shp <- shapviz(final_xgb, X_pred = X_test_unscaled, X = X_test_display)

# 8.1 Global SHAP Summary Plot
p_summary <- sv_importance(shp, kind = "beeswarm", max_display = 10) +
  theme_classic() + ggtitle("SHAP Summary Plot")

# 8.2 SHAP Dependence Plot (Example: HbA1c colored by Oxaliplatin)
p_dependence <- sv_dependence(shp, v = "HbA1c", color_var = "Cumulative_Oxaliplatin") +
  theme_classic() + ggtitle("SHAP Dependence: HbA1c")

# 8.3 Local SHAP Explanation (Waterfall plot for the patient with the highest risk)
highest_risk_idx <- which.max(preds_prob$XGBoost)
p_waterfall <- sv_waterfall(shp, row_id = highest_risk_idx, max_display = 8) +
  theme_classic() + ggtitle(sprintf("Local SHAP Explanation (Patient ID: %d)", highest_risk_idx))

# Display plots (In actual use, save these using ggsave)
print(p_summary)
print(p_dependence)
print(p_waterfall)

cat("\nScript execution completed successfully.\n")
# ============================== END OF SCRIPT =================================