# ===================== Package Loading =====================
library(caret)
library(DALEX)
library(ggplot2)
library(randomForest)
library(kernlab)
library(kernelshap)
library(pROC)
library(shapviz)
library(xgboost)
library(klaR)
library(RColorBrewer)
library(pheatmap)

# ===================== Parameter Settings =====================
work_dir <- "D:\\Tutorial2025\\12.5points.113MachineLearning+SHAP\\35.SHAPInterpretationGeneAnalysis"
result_dir <- "results"  # Output folder for results
random_seed <- 12345
num_colors <- 20  # Pastel palette supports up to 9 colors; beyond that, it will automatically extend with gradients
train_ratio <- 0.7
cv_folds <- 5
roc_file <- file.path(result_dir, "model_roc_curve.pdf")
barplot_file <- file.path(result_dir, "shap_importance_barplot.pdf")
bee_file <- file.path(result_dir, "shap_importance_beeswarm.pdf")
dependence_file <- file.path(result_dir, "shap_dependence_all.pdf")
waterfall_file <- file.path(result_dir, "shap_waterfall_sample1.pdf")
force_file <- file.path(result_dir, "shap_force_sample1.pdf")
heatmap_file <- file.path(result_dir, "shap_value_heatmap.pdf")
top20_file <- file.path(result_dir, "shap_top_genes.csv")
data_file <- "geneexp.csv"
pdf_width <- 6
pdf_height <- 6
progress_style <- 3

ml_methods <- data.frame(
  ModelName = c("RF", "SVM", "XGB", "GBM", "KNN"),
  MethodID = c("rf", "svmRadial", "xgbTree", "gbm", "knn")
)

# ===================== Step 1: Environment Preparation =====================
cat("Step 1/8: Setting the working directory and random seed...\n")
if (!dir.exists(work_dir)) stop("The working directory does not exist. Please check the path!")
setwd(work_dir)
set.seed(random_seed)

# Create output folder for results
if (!dir.exists(result_dir)) dir.create(result_dir)

# ===================== Step 2: Color Palette =====================
cat("Step 2/8: Generating a cute color palette...\n")
if (num_colors > 9) {
  color_palette <- colorRampPalette(brewer.pal(9, "Pastel1"))(num_colors)
} else {
  color_palette <- brewer.pal(num_colors, "Pastel1")
}

# ===================== Step 3: Data Reading and Preprocessing =====================
cat("Step 3/8: Reading and preprocessing expression data...\n")
if (!file.exists(data_file)) stop("The data file does not exist!")
raw_data <- read.csv(data_file, header=TRUE, check.names=FALSE, row.names=1)
if (nrow(raw_data) == 0) stop("The data file is empty!")
transposed_data <- t(raw_data)
sample_ids <- rownames(transposed_data)
if (!all(grepl("(_con$|_tra$)", sample_ids))) stop("Sample names must end with _con or _tra!")
group_labels <- ifelse(grepl("_con$", sample_ids), "Control",
                       ifelse(grepl("_tra$", sample_ids), "Treatment", NA))
if (any(is.na(group_labels))) stop("Failed to generate group labels!")
expr_data <- as.data.frame(transposed_data)
expr_data$Group <- as.factor(group_labels)

# ===================== Step 4: Train/Test Split =====================
cat("Step 4/8: Splitting the data into training and testing sets...\n")
if (length(unique(expr_data$Group)) != 2) stop("The number of groups is not 2, binary classification is not possible!")
split_idx <- createDataPartition(y=expr_data$Group, p=train_ratio, list=FALSE)
train_set <- expr_data[split_idx, ]
test_set <- expr_data[-split_idx, ]
test_labels <- test_set$Group
test_features <- test_set[, -ncol(test_set)]

# ===================== Step 5: Model Training and Evaluation =====================
cat("Step 5/8: Training multiple models and evaluating ROC...\n")
auc_results <- c()
model_auc_map <- list()
roc_colors <- color_palette[1:nrow(ml_methods)]
roc_ci_list <- list()
pb <- txtProgressBar(min=0, max=nrow(ml_methods), style=progress_style)

pdf(file=roc_file, width=pdf_width, height=pdf_height)
par(mar=c(5, 5, 4, 2) + 0.1)
for (i in 1:nrow(ml_methods)) {
  setTxtProgressBar(pb, i)
  mdl_name <- ml_methods$ModelName[i]
  mdl_id <- ml_methods$MethodID[i]
  if (nrow(train_set) < 10) stop("Too few samples in the training set!")
  if (mdl_id == "svmRadial") {
    mdl <- train(Group ~ ., data=train_set, method=mdl_id, prob.model=TRUE,
                 trControl=trainControl(method="repeatedcv", number=cv_folds, savePredictions=TRUE))
  } else {
    mdl <- train(Group ~ ., data=train_set, method=mdl_id,
                 trControl=trainControl(method="repeatedcv", number=cv_folds, savePredictions=TRUE))
  }
  pred_prob <- predict(mdl, newdata=test_features, type="prob")
  if (!"Treatment" %in% colnames(pred_prob)) stop("The predicted probability column does not contain Treatment!")
  roc_obj <- roc(as.numeric(test_labels)-1, as.numeric(pred_prob[,"Treatment"]))
  auc_val <- as.numeric(roc_obj$auc)
  auc_ci <- ci.auc(roc_obj, conf.level=0.95)
  roc_ci_list[[mdl_id]] <- auc_ci
  auc_results <- c(auc_results, 
                   sprintf("%s: %.03f [%.03f, %.03f]", mdl_name, auc_val, auc_ci[1], auc_ci[3]))
  model_auc_map[[mdl_id]] <- auc_val
  if (i == 1) {
    plot(roc_obj, print.auc=FALSE, legacy.axes=TRUE, main="", col=roc_colors[i], lwd=3)
  } else {
    plot(roc_obj, print.auc=FALSE, legacy.axes=TRUE, main="", col=roc_colors[i], lwd=3, add=TRUE)
  }
}
close(pb)
legend('bottomright', auc_results, col=roc_colors, lwd=3, bty='n', cex=0.9)
dev.off()

# ===================== Step 6: Optimal Model Selection and SHAP Calculation =====================
cat("Step 6/8: Selecting the optimal model and calculating SHAP values...\n")
auc_vec <- unlist(model_auc_map)
if (length(auc_vec) == 0) stop("AUC results are empty!")
best_method <- names(which.max(auc_vec))
cat("The model with the highest AUC is:", best_method, "\n")
final_model <- train(Group ~ ., data=train_set, method=best_method,
                     trControl=trainControl(method="repeatedcv", number=cv_folds, savePredictions=TRUE))
shap_fit <- kernelshap(
  object = final_model,
  X = train_set[, -ncol(train_set)],
  pred_fun = function(obj, newdata) {
    pred <- predict(obj, newdata, type="prob")
    if (!"Treatment" %in% colnames(pred)) stop("The SHAP predicted probability column does not contain Treatment!")
    pred[,"Treatment"]
  }
)

# ===================== Step 7: SHAP Visualization =====================
cat("Step 7/8: SHAP Visualization...\n")
shap_obj <- shapviz(shap_fit, X_pred = train_set[, -ncol(train_set)], X = train_set[, -ncol(train_set)], interactions=TRUE)
shap_importance <- sort(colMeans(abs(shap_obj$S)), decreasing=TRUE)
top_features <- names(shap_importance)

# Unified theme
custom_theme <- theme_bw() + 
  theme(plot.title = element_text(hjust = 0.5),
        axis.title = element_text(size = 12),
        axis.text = element_text(size = 10),
        legend.position = "right")

options(shapviz.colors = color_palette)

# SHAP importance bar plot
pdf(file=barplot_file, width=pdf_width, height=pdf_height)
print(sv_importance(shap_obj, kind="bar", show_numbers=TRUE) + custom_theme)
dev.off()

# SHAP importance beeswarm plot
pdf(file=bee_file, width=pdf_width, height=pdf_height)
print(sv_importance(shap_obj, kind="bee", show_numbers=TRUE) + custom_theme)
dev.off()

# Dependency plots for all top features
pdf(file=dependence_file, width=12, height=pdf_height)
print(sv_dependence(shap_obj, v=top_features) + custom_theme)
dev.off()

# For each top gene, generate dependency plots, density plots, and scatter plots
cat("Generating dependency plots, density plots, and scatter plots for each top gene...\n")
for (gene in top_features) {
  safe_gene <- tolower(gsub("[^A-Za-z0-9_]", "_", gene))
  # Dependency plot
  pdf(file = file.path(result_dir, paste0("shap_dependence_", safe_gene, ".pdf")), width = pdf_width, height = pdf_height)
  print(sv_dependence(shap_obj, v = gene) + custom_theme + ggtitle(gene))
  dev.off()
  # Density plot
  shap_df <- data.frame(
    SHAP = shap_obj$S[, gene],
    Group = train_set$Group
  )
  pdf(file = file.path(result_dir, paste0("shap_density_", safe_gene, ".pdf")), width=pdf_width, height=pdf_height)
  print(ggplot(shap_df, aes(x=SHAP, fill=Group)) +
          geom_density(alpha=0.5) +
          custom_theme +
          ggtitle(paste("SHAP Density for", gene)))
  dev.off()
  # SHAP-expression scatter plot
  shap_df <- data.frame(
    SHAP = shap_obj$S[, gene],
    Expr = train_set[, gene],
    Group = train_set$Group
  )
  pdf(file = file.path(result_dir, paste0("shap_scatter_", safe_gene, ".pdf")), width=pdf_width, height=pdf_height)
  print(ggplot(shap_df, aes(x=Expr, y=SHAP, color=Group)) +
          geom_point(alpha=0.7) +
          custom_theme +
          ggtitle(paste("SHAP vs Expression for", gene)))
  dev.off()
}

# Interaction scatter plot for the top 2 genes
gene1 <- top_features[1]
gene2 <- top_features[2]
safe_gene1 <- tolower(gsub("[^A-Za-z0-9_]", "_", gene1))
safe_gene2 <- tolower(gsub("[^A-Za-z0-9_]", "_", gene2))
shap_df <- data.frame(
  SHAP = shap_obj$S[, gene1],
  Expr1 = train_set[, gene1],
  Expr2 = train_set[, gene2],
  Group = train_set$Group
)
pdf(file.path(result_dir, paste0("shap_interaction_scatter_", safe_gene1, "_vs_", safe_gene2, ".pdf")), width=pdf_width, height=pdf_height)
print(
  ggplot(shap_df, aes(x=Expr1, y=Expr2, color=SHAP)) +
    geom_point(size=2, alpha=0.8) +
    scale_color_gradient2(low="blue", mid="white", high="red", midpoint=0) +
    labs(title=paste("Interaction-like plot:", gene1, "vs", gene2),
         x=gene1, y=gene2, color="SHAP") +
    custom_theme
)
dev.off()

# SHAP clustering heatmap
pdf(heatmap_file, width=pdf_width, height=pdf_height)
pheatmap(shap_obj$S, 
         cluster_rows=TRUE, cluster_cols=TRUE, 
         show_rownames=FALSE, show_colnames=TRUE, 
         color = colorRampPalette(rev(brewer.pal(9, "RdBu")))(100))
dev.off()

# ===================== Step 8: Output Genes =====================
cat("Step 8/8: Outputting genes...\n")
top20_genes <- head(shap_importance, 20)
write.csv(data.frame(Gene=names(top20_genes), Importance=as.numeric(top20_genes)), 
          top20_file, row.names=FALSE)

# SHAP feature cumulative contribution curve
cum_contrib <- cumsum(abs(shap_importance)) / sum(abs(shap_importance))
pdf(file.path(result_dir, "shap_cumulative_contribution_curve.pdf"), width = pdf_width, height = pdf_height)
plot(seq_along(cum_contrib), cum_contrib, type = "s", lwd = 2,
     xlab = "Number of top features", ylab = "Cumulative SHAP Contribution",
     main = "Cumulative SHAP Contribution Curve")
abline(h = 0.8, col = "red", lty = 2)  # 80% reference line
dev.off()

# Group-wise SHAP importance bar plot
group_levels <- levels(train_set$Group)
shap_group <- sapply(group_levels, function(grp) 
  colMeans(abs(shap_obj$S[train_set$Group == grp, , drop = FALSE]))
)

# Group-wise SHAP importance bar plot
shap_group <- as.data.frame(shap_group)
shap_group$Gene <- rownames(shap_group)
library(reshape2)
shap_group_long <- reshape2::melt(shap_group, id.vars = "Gene", variable.name = "Group", value.name = "SHAP_Importance")

pdf(file.path(result_dir, "shap_group_importance_barplot.pdf"), width = max(8, pdf_width*1.5), height = pdf_height)
print(
  ggplot(shap_group_long, aes(x = reorder(Gene, SHAP_Importance), y = SHAP_Importance, fill = Group)) +
    geom_bar(stat='identity', position = "dodge") +
    coord_flip() +
    labs(title="Group-wise SHAP Feature Importance", x="Gene", y="Mean(|SHAP|)") +
    custom_theme
)
dev.off()

# Save SHAP feature importance
feature_importance_df <- data.frame(
  Gene = names(shap_importance),
  MeanAbsSHAP = as.numeric(shap_importance)
)
write.csv(feature_importance_df, file.path(result_dir, "shap_feature_importance.csv"), row.names = FALSE)
group_levels <- levels(train_set$Group)
shap_group <- sapply(group_levels, function(grp) 
  colMeans(abs(shap_obj$S[train_set$Group == grp, , drop = FALSE]))
)
shap_group_df <- data.frame(Gene = rownames(shap_group), shap_group)
write.csv(shap_group_df, file.path(result_dir, "shap_groupwise_importance.csv"), row.names = FALSE)

# Save model AUC summary
auc_table <- data.frame(
  Model = ml_methods$ModelName,
  AUC = sapply(ml_methods$MethodID, function(id) as.numeric(model_auc_map[[id]])),
  CI_lower = sapply(ml_methods$MethodID, function(id) roc_ci_list[[id]][1]),
  CI_upper = sapply(ml_methods$MethodID, function(id) roc_ci_list[[id]][3])
)
write.csv(auc_table, file.path(result_dir, "model_auc_summary.csv"), row.names = FALSE)

# Save top influential samples
sample_influence <- rowSums(abs(shap_obj$S))
influential_samples <- order(sample_influence, decreasing=TRUE)[1:10]
influence_df <- data.frame(
  Sample = rownames(train_set)[influential_samples],
  InfluenceScore = sample_influence[influential_samples]
)
write.csv(influence_df, file.path(result_dir, "influential_samples.csv"), row.names=FALSE)

# Convert y to 0/1
y_numeric <- ifelse(train_set$Group == "Treatment", 1, 0)
library(DALEX)
explainer <- explain(
  final_model,
  data = train_set[, -ncol(train_set)],
  y = y_numeric,
  label = "BestModel"
)
vi <- model_parts(explainer, loss_function = loss_one_minus_auc)
vi_df <- as.data.frame(vi)
pdf(file.path(result_dir, "feature_permutation_importance.pdf"), width=pdf_width*1.2, height=pdf_height)
print(
  ggplot(vi_df[vi_df$variable != "_full_model_", ], aes(x=reorder(variable, -dropout_loss), y=dropout_loss)) +
    geom_bar(stat="identity", fill="tomato") +
    coord_flip() +
    labs(title="Permutation Feature Importance (AUC Drop)", x="Feature", y="AUC Drop") +
    custom_theme
)
dev.off()
write.csv(vi_df, file.path(result_dir, "feature_permutation_importance.csv"), row.names=FALSE)

# Waterfall plot (for a specific sample)
pdf(file=waterfall_file, width=pdf_width, height=pdf_height)
print(sv_waterfall(shap_obj, row_id=30) + custom_theme)
dev.off()

# Force plot (for a specific sample)
pdf(file=force_file, width=pdf_width, height=pdf_height)
print(sv_force(shap_obj, row_id=30) + custom_theme)
dev.off()

cat("All steps are completed!\n")
rownames(train_set)[30]