# 1. Packages ---------------------------------------------------------------
packages <- c(
  "readxl", "dplyr", "caret", "pROC", "ggplot2", "tidyr", "purrr",
  "xgboost", "lightgbm", "ranger", "gbm", "glmnet",
  "kernlab", "nnet", "openxlsx", "Cairo", "doParallel", "data.table",
  "PRROC"
)

missing_packages <- packages[!sapply(packages, requireNamespace, quietly = TRUE)]
if (length(missing_packages) > 0) {
  stop("Please install required R packages first:\n", paste(missing_packages, collapse = ", "))
}

suppressPackageStartupMessages({
  library(readxl)
  library(dplyr)
  library(caret)
  library(pROC)
  library(ggplot2)
  library(tidyr)
  library(purrr)
  library(xgboost)
  library(lightgbm)
  library(ranger)
  library(gbm)
  library(glmnet)
  library(kernlab)
  library(nnet)
  library(openxlsx)
  library(Cairo)
  library(doParallel)
  library(data.table)
  library(PRROC)
})

options(scipen = 999, digits = 4)
set.seed(123)

# 2. User settings ----------------------------------------------------------
train_file <- "MIMICtrain_data_winsorized.csv"
test_file  <- "MIMICtest_data_winsorized.csv"
external_file <- "eICU_external_validation_data.xlsx"

cmd_args <- commandArgs(trailingOnly = TRUE)
if (length(cmd_args) >= 3) {
  train_file <- cmd_args[1]
  test_file <- cmd_args[2]
  external_file <- cmd_args[3]
}

current_date <- format(Sys.Date(), "%Y%m%d")
output_dir <- file.path(dirname(train_file), paste0("ML_results_10fold_NestedCV_LASSO_external_", current_date))
model_dir  <- file.path(output_dir, "models")
plot_dir   <- file.path(output_dir, "plots")
table_dir  <- file.path(output_dir, "tables")

input_files <- c(train_file, test_file, external_file)
missing_input_files <- input_files[!file.exists(input_files)]
if (length(missing_input_files) > 0) {
  stop("Input file(s) not found:\n", paste(missing_input_files, collapse = "\n"))
}

for (d in c(output_dir, model_dir, plot_dir, table_dir)) {
  if (!dir.exists(d)) dir.create(d, recursive = TRUE)
}

# 3. Parallel ---------------------------------------------------------------
n_cores <- max(1, parallel::detectCores() - 1)
cl <- parallel::makeCluster(n_cores)
doParallel::registerDoParallel(cl)
Sys.setenv(OMP_NUM_THREADS = n_cores)

on.exit({
  try(parallel::stopCluster(cl), silent = TRUE)
  try(registerDoSEQ(), silent = TRUE)
}, add = TRUE)

# 4. Models and candidate predictors ---------------------------------------
models_to_train <- c(
  "Logistic Regression",
  "Elastic Net",
  "Random Forest",
  "XGBoost",
  "GBM",
  "SVM",
  "kNN",
  "Neural Network",
  "LightGBM"
)

color_palette <- c(
  "Logistic Regression" = "#009E73",
  "Elastic Net" = "#E69F00",
  "Random Forest" = "#D55E00",
  "XGBoost" = "#CC79A7",
  "GBM" = "#AA3377",
  "SVM" = "#882255",
  "kNN" = "#999933",
  "Neural Network" = "#0072B2",
  "LightGBM" = "#56B4E9",
  "SOFA" = "#444444",
  "SAPS II" = "#777777"
)

outcome_var <- "new_mods_after_24h_true"

predictors_master <- c(
  "Vasoactive.Drug", "Mechvent", "SOFA", "BUN", "Creatinine",
  "Hematocrit", "CAD", "RR", "Bicarbonate", "Hemoglobin",
  "Platelets", "Heart.Failure", "Hypertension", "Potassium",
  "Age", "MAP", "Chloride", "HR", "CKD", "RRT",
  "GCS", "Gender", "PH", "PTT", "Diabetes", "COPD",
  "Lactate", "Temperature", "Sodium", "INR", "SPO2"
)


# 5. Data import and harmonisation -----------------------------------------
df_train_raw <- fread(train_file, data.table = FALSE)
df_test_raw  <- fread(test_file, data.table = FALSE)
df_external_raw <- read_excel(external_file)

canonicalize_dataset <- function(df) {
  df <- as.data.frame(df)
  alias_map <- list(
    HR = c("hr"),
    RR = c("rr"),
    MAP = c("map"),
    PH = c("ph", "pH"),
    Temperature = c("temperature", "Temp"),
    SPO2 = c("spo2", "SpO2"),
    Hemoglobin = c("hemoglobin"),
    Hematocrit = c("hematocrit"),
    Platelets = c("platelets"),
    Sodium = c("sodium"),
    Potassium = c("potassium"),
    Chloride = c("chloride"),
    Bicarbonate = c("bicarbonate"),
    Creatinine = c("creatinine"),
    BUN = c("bun"),
    Lactate = c("lactate"),
    PTT = c("ptt"),
    INR = c("inr"),
    GCS = c("gcs"),
    SOFA = c("sofa"),
    Mechvent = c("mechvent", "Mechanical.Ventilation", "mechanical_ventilation"),
    Heart.Failure = c("heart_failure"),
    Hypertension = c("hypertension"),
    Diabetes = c("diabetes"),
    COPD = c("copd"),
    CKD = c("ckd"),
    CAD = c("cad"),
    RRT = c("rrt", "CRRT", "crrt"),
    Vasoactive.Drug = c("vasoactive_drug", "Vasoactive.Drug.1"),
    new_mods_after_24h_true = c("MODS.Flag", "MODS.Flag.1", "MODS Flag", "MODS_Flag")
  )
  
  for (canonical in names(alias_map)) {
    aliases <- alias_map[[canonical]]
    existing_aliases <- aliases[aliases %in% names(df)]
    if (!(canonical %in% names(df)) && length(existing_aliases) > 0) {
      names(df)[names(df) == existing_aliases[1]] <- canonical
    } else if ((canonical %in% names(df)) && length(existing_aliases) > 0) {
      for (al in existing_aliases) {
        miss_idx <- which(is.na(df[[canonical]]) & !is.na(df[[al]]))
        if (length(miss_idx) > 0) df[[canonical]][miss_idx] <- df[[al]][miss_idx]
      }
    }
  }
  df
}

df_train_raw <- canonicalize_dataset(df_train_raw)
df_test_raw <- canonicalize_dataset(df_test_raw)
df_external_raw <- canonicalize_dataset(df_external_raw)

if (!(outcome_var %in% names(df_train_raw)) ||
    !(outcome_var %in% names(df_test_raw)) ||
    !(outcome_var %in% names(df_external_raw))) {
  stop("Outcome variable is missing in at least one dataset: ", outcome_var)
}

candidate_predictors <- Reduce(intersect, list(
  intersect(predictors_master, names(df_train_raw)),
  intersect(predictors_master, names(df_test_raw)),
  intersect(predictors_master, names(df_external_raw))
))

if (length(candidate_predictors) < 2) {
  stop("Too few common predictors are available across train/test/external datasets.")
}

safe_gender_to_numeric <- function(x) {
  x0 <- as.character(x)
  x1 <- trimws(tolower(x0))
  out <- rep(NA_real_, length(x1))
  out[x1 %in% c("male", "m", "1")] <- 1
  out[x1 %in% c("female", "f", "0")] <- 0
  suppressWarnings(num_try <- as.numeric(x0))
  fill_idx <- is.na(out) & !is.na(num_try)
  out[fill_idx] <- num_try[fill_idx]
  out
}

prepare_dataset <- function(df, outcome_var, predictors) {
  need_cols <- unique(c(outcome_var, predictors))
  for (cc in need_cols) {
    if (!(cc %in% names(df))) df[[cc]] <- NA
  }
  df <- df[, need_cols, drop = FALSE]
  if ("Gender" %in% names(df)) df$Gender <- safe_gender_to_numeric(df$Gender)
  for (nm in setdiff(names(df), outcome_var)) {
    df[[nm]] <- suppressWarnings(as.numeric(as.character(df[[nm]])))
  }
  df[[outcome_var]] <- factor(
    ifelse(as.numeric(as.character(df[[outcome_var]])) == 1, "MODS", "No_MODS"),
    levels = c("No_MODS", "MODS")
  )
  df
}

df_train <- prepare_dataset(df_train_raw, outcome_var, candidate_predictors)
df_test <- prepare_dataset(df_test_raw, outcome_var, candidate_predictors)
df_external <- prepare_dataset(df_external_raw, outcome_var, candidate_predictors)

y_train <- df_train[[outcome_var]]
y_test <- df_test[[outcome_var]]
y_external <- df_external[[outcome_var]]

# 6. LASSO feature selection ------------------------------------------------
fit_lasso_selection <- function(df, outcome_var, predictors, seed = 123) {
  set.seed(seed)
  x <- as.matrix(df[, predictors, drop = FALSE])
  y <- as.numeric(df[[outcome_var]] == "MODS")
  
  pp <- caret::preProcess(as.data.frame(x), method = c("center", "scale"))
  x_scaled <- as.matrix(predict(pp, as.data.frame(x)))
  
  cvfit <- glmnet::cv.glmnet(
    x = x_scaled,
    y = y,
    family = "binomial",
    alpha = 1,
    nfolds = 10,
    type.measure = "auc"
  )
  
  coef_min <- coef(cvfit, s = "lambda.min")
  nonzero <- rownames(coef_min)[as.numeric(coef_min) != 0]
  selected <- setdiff(nonzero, "(Intercept)")
  
  if (length(selected) == 0) {
    coef_1se <- coef(cvfit, s = "lambda.1se")
    nonzero <- rownames(coef_1se)[as.numeric(coef_1se) != 0]
    selected <- setdiff(nonzero, "(Intercept)")
  }
  
  list(cvfit = cvfit, preprocess = pp, selected = selected)
}

lasso_fit <- fit_lasso_selection(df_train, outcome_var, candidate_predictors, seed = 123)
selected_predictors <- lasso_fit$selected

if (length(selected_predictors) == 0) {
  stop("LASSO selected no predictors. Please inspect candidate variables and outcome coding.")
}

saveRDS(lasso_fit$cvfit, file.path(model_dir, "lasso_feature_selection_cvglmnet.rds"))
write.xlsx(
  data.frame(Selected_Predictor = selected_predictors),
  file.path(table_dir, "LASSO_Selected_Predictors.xlsx"),
  overwrite = TRUE
)
writeLines(selected_predictors, file.path(model_dir, "selected_predictors_lasso.txt"))

predictors <- selected_predictors

# 7. Preprocessing ----------------------------------------------------------
fit_preprocess_on_train <- function(df_train_fold, predictors) {
  medians <- sapply(df_train_fold[, predictors, drop = FALSE], function(x) {
    med <- median(x, na.rm = TRUE)
    if (is.na(med)) med <- 0
    med
  })
  
  impute_with_medians <- function(df) {
    out <- df
    for (nm in predictors) {
      idx <- which(is.na(out[[nm]]))
      if (length(idx) > 0) out[[nm]][idx] <- medians[[nm]]
    }
    out
  }
  
  df_train_imp <- impute_with_medians(df_train_fold)
  preprocess <- preProcess(df_train_imp[, predictors, drop = FALSE], method = c("center", "scale"))
  
  transform_data <- function(df) {
    df_imp <- impute_with_medians(df)
    x_scaled <- predict(preprocess, df_imp[, predictors, drop = FALSE])
    x_scaled[is.na(x_scaled)] <- 0
    x_scaled
  }
  
  list(medians = medians, preprocess = preprocess, transform = transform_data)
}

make_lgb_safe_names <- function(x) {
  x <- gsub("\\[|\\]|\\(|\\)|%|#|\\^|/|-|\\s+|\\.", "_", x)
  x <- gsub("__+", "_", x)
  x
}

predictors_lgb <- make_lgb_safe_names(predictors)

as_lgb_matrix <- function(x_df) {
  x_mat <- as.matrix(x_df)
  colnames(x_mat) <- predictors_lgb
  x_mat
}

# 8. Hyperparameter grids ---------------------------------------------------
sigma_est <- tryCatch({
  as.numeric(kernlab::sigest(as.matrix(df_train[, predictors, drop = FALSE])))
}, error = function(e) c(0.001, 0.01, 0.1))

sigma_center <- ifelse(length(sigma_est) >= 2, sigma_est[2], median(sigma_est, na.rm = TRUE))
if (is.na(sigma_center) || sigma_center <= 0) sigma_center <- 0.01

svm_grid <- expand.grid(
  sigma = unique(signif(c(sigma_center / 4, sigma_center / 2, sigma_center, sigma_center * 2), 3)),
  C = c(0.25, 0.5, 1, 2, 4)
) %>% filter(sigma > 0)

elastic_grid <- expand.grid(alpha = seq(0.1, 1, 0.2), lambda = 10^seq(-4, 0, length.out = 20))

rf_grid <- expand.grid(
  mtry = unique(pmax(1, round(c(sqrt(length(predictors)), length(predictors) / 4, length(predictors) / 3, length(predictors) / 2)))),
  splitrule = "gini",
  min.node.size = c(5, 10, 20, 30, 50)
)

xgb_grid <- expand.grid(
  nrounds = c(100, 200, 300),
  max_depth = c(2, 3, 4),
  eta = c(0.01, 0.03, 0.05),
  gamma = c(0, 0.1),
  colsample_bytree = c(0.7, 0.9),
  min_child_weight = c(1, 5),
  subsample = c(0.7, 0.9)
) %>% slice(1:48)

gbm_grid <- expand.grid(
  n.trees = c(100, 200, 300),
  interaction.depth = c(1, 2, 3, 4),
  shrinkage = c(0.01, 0.03, 0.05),
  n.minobsinnode = c(10, 20)
) %>% slice(1:48)

knn_grid <- expand.grid(k = seq(5, 41, by = 2))
nnet_grid <- expand.grid(size = c(3, 5, 8, 12), decay = c(0.001, 0.01, 0.1, 0.5, 1))

lgb_grid <- expand.grid(
  learning_rate = c(0.01, 0.03, 0.05),
  num_leaves = c(15, 31),
  max_depth = c(3, 4, 6),
  min_data_in_leaf = c(10, 20, 40),
  feature_fraction = c(0.7, 0.9),
  bagging_fraction = c(0.7, 0.9),
  lambda_l1 = c(0, 0.5, 1),
  lambda_l2 = c(0, 0.5, 1)
) %>% slice(1:48)

write.xlsx(
  list(
    Elastic_Net = elastic_grid,
    Random_Forest = rf_grid,
    XGBoost = xgb_grid,
    GBM = gbm_grid,
    SVM = svm_grid,
    kNN = knn_grid,
    Neural_Network = nnet_grid,
    LightGBM = lgb_grid
  ),
  file.path(table_dir, "Hyperparameter_Grids.xlsx"),
  overwrite = TRUE
)

# 9. Training and prediction helpers ---------------------------------------
make_inner_ctrl <- function() {
  trainControl(
    method = "cv",
    number = 10,
    classProbs = TRUE,
    summaryFunction = twoClassSummary,
    savePredictions = "final",
    verboseIter = FALSE,
    allowParallel = TRUE
  )
}

train_caret_model <- function(model_name, x, y, seed = 123) {
  set.seed(seed)
  ctrl <- make_inner_ctrl()
  
  if (model_name == "Logistic Regression") {
    return(train(x = x, y = y, method = "glm", family = binomial, trControl = ctrl, metric = "ROC"))
  }
  if (model_name == "Elastic Net") {
    return(train(x = x, y = y, method = "glmnet", family = "binomial", trControl = ctrl, metric = "ROC", tuneGrid = elastic_grid))
  }
  if (model_name == "Random Forest") {
    return(train(x = x, y = y, method = "ranger", trControl = ctrl, metric = "ROC", tuneGrid = rf_grid, num.trees = 500, importance = "none"))
  }
  if (model_name == "XGBoost") {
    return(train(x = x, y = y, method = "xgbTree", trControl = ctrl, metric = "ROC", tuneGrid = xgb_grid, verbose = FALSE))
  }
  if (model_name == "GBM") {
    return(train(x = x, y = y, method = "gbm", distribution = "bernoulli", verbose = FALSE, trControl = ctrl, metric = "ROC", tuneGrid = gbm_grid))
  }
  if (model_name == "SVM") {
    return(train(x = x, y = y, method = "svmRadial", trControl = ctrl, metric = "ROC", tuneGrid = svm_grid))
  }
  if (model_name == "kNN") {
    return(train(x = x, y = y, method = "knn", trControl = ctrl, metric = "ROC", tuneGrid = knn_grid))
  }
  if (model_name == "Neural Network") {
    return(train(x = x, y = y, method = "nnet", trControl = ctrl, metric = "ROC", tuneGrid = nnet_grid, trace = FALSE, maxit = 500, linout = FALSE))
  }
  stop("Unknown caret model: ", model_name)
}

train_lightgbm_inner_cv <- function(x, y, seed = 123, nfold = 10) {
  set.seed(seed)
  y_bin <- as.numeric(y == "MODS")
  lgb_train <- lgb.Dataset(as_lgb_matrix(x), label = y_bin)
  
  best_auc <- -Inf
  best_iter <- NA_integer_
  best_params <- NULL
  best_row <- NULL
  
  for (ii in seq_len(nrow(lgb_grid))) {
    g <- lgb_grid[ii, ]
    params <- list(
      objective = "binary",
      metric = "auc",
      learning_rate = g$learning_rate,
      num_leaves = g$num_leaves,
      max_depth = g$max_depth,
      min_data_in_leaf = g$min_data_in_leaf,
      feature_fraction = g$feature_fraction,
      bagging_fraction = g$bagging_fraction,
      bagging_freq = 5,
      lambda_l1 = g$lambda_l1,
      lambda_l2 = g$lambda_l2,
      verbosity = -1,
      num_threads = n_cores
    )
    
    cv_fit <- tryCatch(
      lgb.cv(
        params = params,
        data = lgb_train,
        nrounds = 700,
        nfold = nfold,
        early_stopping_rounds = 40,
        verbose = -1,
        stratified = TRUE
      ),
      error = function(e) NULL
    )
    
    if (is.null(cv_fit)) next
    cv_auc <- max(unlist(cv_fit$record_evals$valid$auc$eval), na.rm = TRUE)
    if (!is.na(cv_auc) && cv_auc > best_auc) {
      best_auc <- cv_auc
      best_iter <- cv_fit$best_iter
      best_params <- params
      best_row <- g
    }
  }
  
  final_model <- lightgbm(params = best_params, data = lgb_train, nrounds = best_iter, verbose = -1)
  attr(final_model, "cv_auc") <- best_auc
  attr(final_model, "best_iter") <- best_iter
  attr(final_model, "best_params") <- best_params
  attr(final_model, "best_row") <- best_row
  final_model
}

train_model <- function(model_name, x, y, seed = 123) {
  if (model_name == "LightGBM") train_lightgbm_inner_cv(x, y, seed = seed, nfold = 10) else train_caret_model(model_name, x, y, seed = seed)
}

predict_model_prob <- function(model, model_name, newx) {
  if (model_name == "LightGBM") return(as.numeric(predict(model, as_lgb_matrix(newx))))
  pred <- predict(model, newdata = newx, type = "prob")
  if (is.data.frame(pred) || is.matrix(pred)) {
    if ("MODS" %in% colnames(pred)) return(as.numeric(pred[, "MODS"]))
    if (ncol(pred) >= 2) return(as.numeric(pred[, 2]))
  }
  rep(0.5, nrow(newx))
}

get_model_best_auc <- function(model, model_name) {
  if (model_name == "LightGBM") return(as.numeric(attr(model, "cv_auc")))
  res <- model$results
  best <- model$bestTune
  if (is.null(res) || is.null(best)) return(NA_real_)
  tmp <- res
  for (nm in names(best)) tmp <- tmp[tmp[[nm]] == best[[nm]], , drop = FALSE]
  if ("ROC" %in% names(tmp) && nrow(tmp) >= 1) return(as.numeric(tmp$ROC[1]))
  NA_real_
}

extract_best_tune <- function(model, model_name, fold_label = "Final") {
  if (model_name == "LightGBM") {
    row <- as.data.frame(attr(model, "best_row"))
    row$best_iter <- attr(model, "best_iter")
    row$CV_AUC <- attr(model, "cv_auc")
  } else {
    row <- as.data.frame(model$bestTune)
    row$CV_AUC <- get_model_best_auc(model, model_name)
  }
  row$Model <- model_name
  row$Fold <- fold_label
  row
}

# 10. Metrics ---------------------------------------------------------------
calc_pr_auc <- function(y_bin, p_vec) {
  pos <- p_vec[y_bin == 1]
  neg <- p_vec[y_bin == 0]
  if (length(pos) == 0 || length(neg) == 0) return(NA_real_)
  pr <- PRROC::pr.curve(scores.class0 = pos, scores.class1 = neg, curve = FALSE)
  as.numeric(pr$auc.integral)
}

calc_calibration <- function(y_bin, p_vec) {
  p_clip <- pmin(pmax(p_vec, 1e-6), 1 - 1e-6)
  lp <- qlogis(p_clip)
  intercept <- tryCatch({
    fit_int <- glm(y_bin ~ 1, offset = lp, family = binomial)
    as.numeric(coef(fit_int)[1])
  }, error = function(e) NA_real_)
  slope <- tryCatch({
    fit_slope <- glm(y_bin ~ lp, family = binomial)
    as.numeric(coef(fit_slope)[2])
  }, error = function(e) NA_real_)
  list(intercept = intercept, slope = slope)
}

calc_metrics <- function(y_true, p_vec, positive_class = "MODS") {
  y_true <- factor(y_true, levels = c("No_MODS", "MODS"))
  y_bin <- ifelse(y_true == positive_class, 1, 0)
  
  if (length(unique(y_true)) < 2) {
    return(list(
      AUC = NA, AUC_LCL = NA, AUC_UCL = NA, PR_AUC = NA, Brier = NA,
      Calibration_Intercept = NA, Calibration_Slope = NA,
      Accuracy = NA, Sensitivity = NA, Specificity = NA,
      PPV = NA, NPV = NA, F1 = NA, Best_Threshold = NA
    ))
  }
  
  roc_obj <- roc(y_true, p_vec, levels = c("No_MODS", "MODS"), direction = "<", quiet = TRUE)
  auc_val <- as.numeric(auc(roc_obj))
  auc_ci <- tryCatch(as.numeric(ci.auc(roc_obj)), error = function(e) c(NA, NA, NA))
  best_thresh <- tryCatch(as.numeric(coords(roc_obj, "best", ret = "threshold", best.method = "youden")), error = function(e) 0.5)
  if (is.na(best_thresh)) best_thresh <- 0.5
  
  pred_class <- factor(ifelse(p_vec >= best_thresh, "MODS", "No_MODS"), levels = c("No_MODS", "MODS"))
  cm <- confusionMatrix(pred_class, y_true, positive = "MODS")
  calib <- calc_calibration(y_bin, p_vec)
  
  list(
    AUC = auc_val,
    AUC_LCL = auc_ci[1],
    AUC_UCL = auc_ci[3],
    PR_AUC = calc_pr_auc(y_bin, p_vec),
    Brier = mean((y_bin - p_vec)^2),
    Calibration_Intercept = calib$intercept,
    Calibration_Slope = calib$slope,
    Accuracy = unname(cm$overall["Accuracy"]),
    Sensitivity = unname(cm$byClass["Sensitivity"]),
    Specificity = unname(cm$byClass["Specificity"]),
    PPV = unname(cm$byClass["Pos Pred Value"]),
    NPV = unname(cm$byClass["Neg Pred Value"]),
    F1 = unname(cm$byClass["F1"]),
    Best_Threshold = best_thresh
  )
}

make_perf_df <- function(res_list, dataset_name) {
  do.call(rbind, lapply(names(res_list), function(nm) {
    x <- res_list[[nm]]
    data.frame(
      Dataset = dataset_name,
      Model = nm,
      AUC = round(x$AUC, 3),
      AUC_95CI = paste0(sprintf("%.3f", x$AUC_LCL), " - ", sprintf("%.3f", x$AUC_UCL)),
      PR_AUC = round(x$PR_AUC, 3),
      Brier = round(x$Brier, 4),
      Calibration_Intercept = round(x$Calibration_Intercept, 3),
      Calibration_Slope = round(x$Calibration_Slope, 3),
      Accuracy = round(x$Accuracy, 3),
      Sensitivity = round(x$Sensitivity, 3),
      Specificity = round(x$Specificity, 3),
      PPV = round(x$PPV, 3),
      NPV = round(x$NPV, 3),
      F1 = round(x$F1, 3),
      Best_Threshold = round(x$Best_Threshold, 4),
      stringsAsFactors = FALSE
    )
  }))
}

# 11. Software versions -----------------------------------------------------
package_versions <- data.frame(
  Package = c("R", packages),
  Version = c(paste(R.version$major, R.version$minor, sep = "."), sapply(packages, function(pkg) as.character(packageVersion(pkg)))),
  row.names = NULL
)
write.xlsx(package_versions, file.path(table_dir, "Software_Versions.xlsx"), overwrite = TRUE)
capture.output(sessionInfo(), file = file.path(table_dir, "sessionInfo.txt"))

# 12. Nested CV -------------------------------------------------------------
cat("\nStarting 10-fold nested CV...\n")
outer_folds <- createFolds(y_train, k = 10, returnTrain = FALSE)
outer_predictions <- data.frame()
outer_fold_metrics <- data.frame()
outer_best_tunes <- list()

for (fold_id in seq_along(outer_folds)) {
  cat("\nOuter fold ", fold_id, "/10\n", sep = "")
  valid_idx <- outer_folds[[fold_id]]
  train_idx <- setdiff(seq_len(nrow(df_train)), valid_idx)
  
  fold_train <- df_train[train_idx, , drop = FALSE]
  fold_valid <- df_train[valid_idx, , drop = FALSE]
  y_fold_train <- fold_train[[outcome_var]]
  y_fold_valid <- fold_valid[[outcome_var]]
  
  pp <- fit_preprocess_on_train(fold_train, predictors)
  x_fold_train <- pp$transform(fold_train)
  x_fold_valid <- pp$transform(fold_valid)
  
  for (model_name in models_to_train) {
    cat("  Training ", model_name, "\n", sep = "")
    model_fit <- tryCatch(
      train_model(model_name, x_fold_train, y_fold_train, seed = 1000 + fold_id),
      error = function(e) {
        cat("    Failed: ", e$message, "\n", sep = "")
        NULL
      }
    )
    if (is.null(model_fit)) next
    
    p_valid <- predict_model_prob(model_fit, model_name, x_fold_valid)
    met <- calc_metrics(y_fold_valid, p_valid)
    
    outer_predictions <- rbind(
      outer_predictions,
      data.frame(
        Outer_Fold = fold_id,
        Row_Index = valid_idx,
        Model = model_name,
        Outcome = y_fold_valid,
        Predicted_Probability = p_valid
      )
    )
    
    outer_fold_metrics <- rbind(
      outer_fold_metrics,
      data.frame(
        Outer_Fold = fold_id,
        Model = model_name,
        AUC = met$AUC,
        PR_AUC = met$PR_AUC,
        Brier = met$Brier,
        Calibration_Intercept = met$Calibration_Intercept,
        Calibration_Slope = met$Calibration_Slope,
        Accuracy = met$Accuracy,
        Sensitivity = met$Sensitivity,
        Specificity = met$Specificity,
        PPV = met$PPV,
        NPV = met$NPV,
        F1 = met$F1
      )
    )
    
    outer_best_tunes[[paste0("Fold", fold_id, "_", model_name)]] <- extract_best_tune(model_fit, model_name, paste0("OuterFold_", fold_id))
  }
}

nested_oof_metrics <- do.call(rbind, lapply(models_to_train, function(model_name) {
  tmp <- outer_predictions %>% filter(Model == model_name)
  if (nrow(tmp) == 0) return(NULL)
  met <- calc_metrics(tmp$Outcome, tmp$Predicted_Probability)
  make_perf_df(setNames(list(met), model_name), "NestedCV_OutOfFold")
}))

nested_fold_summary <- outer_fold_metrics %>%
  group_by(Model) %>%
  summarise(
    Mean_AUC = round(mean(AUC, na.rm = TRUE), 3),
    SD_AUC = round(sd(AUC, na.rm = TRUE), 3),
    Mean_PR_AUC = round(mean(PR_AUC, na.rm = TRUE), 3),
    Mean_Brier = round(mean(Brier, na.rm = TRUE), 4),
    Mean_Calibration_Intercept = round(mean(Calibration_Intercept, na.rm = TRUE), 3),
    Mean_Calibration_Slope = round(mean(Calibration_Slope, na.rm = TRUE), 3),
    .groups = "drop"
  ) %>%
  arrange(desc(Mean_AUC))

write.xlsx(
  list(
    NestedCV_OutOfFold_Performance = nested_oof_metrics,
    NestedCV_Fold_Metrics = outer_fold_metrics,
    NestedCV_Fold_Summary = nested_fold_summary,
    NestedCV_OutOfFold_Predictions = outer_predictions,
    NestedCV_Best_Hyperparameters = bind_rows(outer_best_tunes)
  ),
  file.path(table_dir, "NestedCV_10fold_Results.xlsx"),
  overwrite = TRUE
)

# 13. Final models ----------------------------------------------------------
cat("\nTraining final models on full training set...\n")
pp_final <- fit_preprocess_on_train(df_train, predictors)
x_train <- pp_final$transform(df_train)
x_test <- pp_final$transform(df_test)
x_external <- pp_final$transform(df_external)

final_models <- list()
final_cv_auc <- list()
final_best_tunes <- list()

for (model_name in models_to_train) {
  cat("Final model: ", model_name, "\n", sep = "")
  fit <- tryCatch(
    train_model(model_name, x_train, y_train, seed = 123),
    error = function(e) {
      cat("  Failed: ", e$message, "\n", sep = "")
      NULL
    }
  )
  if (is.null(fit)) next
  final_models[[model_name]] <- fit
  final_cv_auc[[model_name]] <- get_model_best_auc(fit, model_name)
  final_best_tunes[[model_name]] <- extract_best_tune(fit, model_name, "FinalFullTrain")
}

model_names <- names(final_models)
write.xlsx(bind_rows(final_best_tunes), file.path(table_dir, "Final_Best_Hyperparameters_10foldCV.xlsx"), overwrite = TRUE)

if ("Random Forest" %in% model_names) {
  rf_best <- final_models[["Random Forest"]]$bestTune
  rf_final <- ranger::ranger(
    x = x_train,
    y = y_train,
    num.trees = 500,
    mtry = rf_best$mtry,
    min.node.size = rf_best$min.node.size,
    splitrule = as.character(rf_best$splitrule),
    probability = TRUE,
    importance = "permutation",
    seed = 123
  )
  final_models[["Random Forest_final_with_importance"]] <- rf_final
  rf_imp <- data.frame(Feature = names(rf_final$variable.importance), Importance = as.numeric(rf_final$variable.importance)) %>%
    arrange(desc(Importance))
  write.xlsx(rf_imp, file.path(table_dir, "RandomForest_Variable_Importance.xlsx"), overwrite = TRUE)
}

probs_train <- list()
probs_test <- list()
probs_external <- list()

for (model_name in model_names) {
  if (model_name == "Random Forest" && "Random Forest_final_with_importance" %in% names(final_models)) {
    pred_train <- predict(final_models[["Random Forest_final_with_importance"]], data = x_train)$predictions
    pred_test <- predict(final_models[["Random Forest_final_with_importance"]], data = x_test)$predictions
    pred_external <- predict(final_models[["Random Forest_final_with_importance"]], data = x_external)$predictions
    idx <- if ("MODS" %in% colnames(pred_train)) which(colnames(pred_train) == "MODS") else 2
    probs_train[[model_name]] <- as.numeric(pred_train[, idx])
    probs_test[[model_name]] <- as.numeric(pred_test[, idx])
    probs_external[[model_name]] <- as.numeric(pred_external[, idx])
  } else {
    probs_train[[model_name]] <- predict_model_prob(final_models[[model_name]], model_name, x_train)
    probs_test[[model_name]] <- predict_model_prob(final_models[[model_name]], model_name, x_test)
    probs_external[[model_name]] <- predict_model_prob(final_models[[model_name]], model_name, x_external)
  }
}

results_train <- lapply(model_names, function(nm) calc_metrics(y_train, probs_train[[nm]]))
names(results_train) <- model_names
results_test <- lapply(model_names, function(nm) calc_metrics(y_test, probs_test[[nm]]))
names(results_test) <- model_names
results_external <- lapply(model_names, function(nm) calc_metrics(y_external, probs_external[[nm]]))
names(results_external) <- model_names

perf_train <- make_perf_df(results_train, "Train")
perf_test <- make_perf_df(results_test, "Test")
perf_external <- make_perf_df(results_external, "External")
perf_all <- bind_rows(nested_oof_metrics, perf_train, perf_test, perf_external)

auc_summary <- data.frame(
  Model = model_names,
  NestedCV_OutOfFold_AUC = nested_oof_metrics$AUC[match(model_names, nested_oof_metrics$Model)],
  Final_InnerCV_AUC = round(unlist(final_cv_auc[model_names]), 3),
  Train_AUC = round(sapply(results_train, function(x) x$AUC), 3),
  Test_AUC = round(sapply(results_test, function(x) x$AUC), 3),
  External_AUC = round(sapply(results_external, function(x) x$AUC), 3),
  stringsAsFactors = FALSE
) %>% arrange(desc(External_AUC))

write.xlsx(
  list(NestedCV_OutOfFold = nested_oof_metrics, Train = perf_train, Test = perf_test, External = perf_external, All = perf_all),
  file.path(table_dir, "Model_Performance_NestedCV_Train_Test_External.xlsx"),
  overwrite = TRUE
)
write.xlsx(auc_summary, file.path(table_dir, "AUC_Summary_NestedCV_Train_Test_External.xlsx"), overwrite = TRUE)

# 14. ROC plots -------------------------------------------------------------
build_roc_df <- function(y_true, probs_list, dataset_name) {
  bind_rows(lapply(names(probs_list), function(nm) {
    roc_obj <- roc(y_true, probs_list[[nm]], levels = c("No_MODS", "MODS"), direction = "<", quiet = TRUE)
    data.frame(FPR = 1 - roc_obj$specificities, TPR = roc_obj$sensitivities, Model = nm, AUC = as.numeric(auc(roc_obj)), Dataset = dataset_name)
  }))
}

build_nested_roc_df <- function(pred_df, dataset_name = "NestedCV Out-of-fold") {
  bind_rows(lapply(unique(pred_df$Model), function(nm) {
    tmp <- pred_df[pred_df$Model == nm, , drop = FALSE]
    roc_obj <- roc(tmp$Outcome, tmp$Predicted_Probability, levels = c("No_MODS", "MODS"), direction = "<", quiet = TRUE)
    data.frame(FPR = 1 - roc_obj$specificities, TPR = roc_obj$sensitivities, Model = nm, AUC = as.numeric(auc(roc_obj)), Dataset = dataset_name)
  }))
}

plot_roc_single <- function(roc_df, dataset_title) {
  label_map <- roc_df %>%
    group_by(Model) %>%
    summarise(AUC = unique(AUC)[1], .groups = "drop") %>%
    arrange(desc(AUC)) %>%
    mutate(Label = paste0(Model, " (AUC=", sprintf("%.3f", AUC), ")"))
  roc_df2 <- roc_df %>% left_join(label_map[, c("Model", "Label")], by = "Model")
  ggplot(roc_df2, aes(x = FPR, y = TPR, color = Model)) +
    geom_line(linewidth = 1) +
    geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "grey50") +
    scale_x_continuous(limits = c(0, 1), expand = c(0, 0)) +
    scale_y_continuous(limits = c(0, 1), expand = c(0, 0)) +
    labs(title = paste0("ROC Curves - ", dataset_title), x = "1 - Specificity", y = "Sensitivity", color = NULL) +
    theme_classic(base_size = 12) +
    theme(plot.title = element_text(hjust = 0.5, face = "bold"), legend.position = "right", legend.text = element_text(size = 9))
}

roc_nested_df <- build_nested_roc_df(outer_predictions)
roc_train_df <- build_roc_df(y_train, probs_train, "Train")
roc_test_df <- build_roc_df(y_test, probs_test, "Test")
roc_external_df <- build_roc_df(y_external, probs_external, "External")

save_plot_eps_jpeg <- function(plot_obj, base_filename, width = 11, height = 8, dpi = 600) {
  ggsave(file.path(plot_dir, paste0(base_filename, "_600dpi.jpeg")), plot_obj, width = width, height = height, dpi = dpi)
  ggsave(file.path(plot_dir, paste0(base_filename, ".eps")), plot_obj, width = width, height = height, device = cairo_ps, fallback_resolution = dpi)
}

save_plot_eps_jpeg(plot_roc_single(roc_nested_df, "10-fold Nested CV Out-of-fold"), "ROC_NestedCV_OutOfFold")
save_plot_eps_jpeg(plot_roc_single(roc_train_df, "Train Set"), "ROC_Train")
save_plot_eps_jpeg(plot_roc_single(roc_test_df, "Test Set"), "ROC_Test")
save_plot_eps_jpeg(plot_roc_single(roc_external_df, "External Validation Set"), "ROC_External")

# 15. AUC bar plot ----------------------------------------------------------
auc_plot_data <- data.frame(
  Model = rep(model_names, 5),
  Dataset = rep(c("Nested CV", "Final Inner CV", "Train", "Test", "External"), each = length(model_names)),
  AUC = c(
    auc_summary$NestedCV_OutOfFold_AUC[match(model_names, auc_summary$Model)],
    auc_summary$Final_InnerCV_AUC[match(model_names, auc_summary$Model)],
    auc_summary$Train_AUC[match(model_names, auc_summary$Model)],
    auc_summary$Test_AUC[match(model_names, auc_summary$Model)],
    auc_summary$External_AUC[match(model_names, auc_summary$Model)]
  )
)

auc_plot_data$Model <- factor(auc_plot_data$Model, levels = auc_summary$Model)

p_auc <- ggplot(auc_plot_data, aes(x = Model, y = AUC, fill = Dataset)) +
  geom_col(position = position_dodge(0.8), width = 0.7) +
  geom_text(aes(label = sprintf("%.3f", AUC)), position = position_dodge(0.8), vjust = -0.3, size = 2.7) +
  coord_cartesian(ylim = c(0, 1)) +
  scale_fill_manual(values = c(
    "Nested CV" = "#D55E00",
    "Final Inner CV" = "#999999",
    "Train" = "#666666",
    "Test" = "#0072B2",
    "External" = "#009E73"
  )) +
  labs(title = "Model AUC Performance", x = NULL, y = "AUC") +
  theme_classic(base_size = 12) +
  theme(axis.text.x = element_text(angle = 30, hjust = 1), plot.title = element_text(hjust = 0.5, face = "bold"))

save_plot_eps_jpeg(p_auc, "AUC_Comparison_NestedCV", width = 15, height = 8)

# 16. Save models and results ----------------------------------------------
saveRDS(pp_final$preprocess, file.path(model_dir, "preprocess_params.rds"))
saveRDS(pp_final$medians, file.path(model_dir, "train_medians.rds"))
saveRDS(final_models, file.path(model_dir, "final_models_all.rds"))

for (name in model_names) {
  saveRDS(final_models[[name]], file.path(model_dir, paste0(gsub(" ", "_", tolower(name)), ".rds")))
}

saveRDS(
  list(
    predictors_master = predictors_master,
    candidate_predictors = candidate_predictors,
    selected_predictors_lasso = selected_predictors,
    nested_cv_predictions = outer_predictions,
    nested_cv_metrics = nested_oof_metrics,
    final_probs_train = probs_train,
    final_probs_test = probs_test,
    final_probs_external = probs_external,
    y_train = y_train,
    y_test = y_test,
    y_external = y_external
  ),
  file.path(output_dir, "NestedCV_Final_Predictions_and_Results.rds")
)

cat("\n", strrep("=", 90), "\n", sep = "")
cat("10-fold nested CV machine learning analysis completed.\n")
cat("Predictors selected by LASSO:\n")
cat(paste(selected_predictors, collapse = ", "), "\n")
best_external_model <- auc_summary$Model[1]
cat("\nBest model by external-validation AUC:", best_external_model, "\n")
print(auc_summary)
cat("\nOutput directory:\n", output_dir, "\n", sep = "")
cat(strrep("=", 90), "\n", sep = "")

parallel::stopCluster(cl)
registerDoSEQ()