PROJECT_DIR <- paste0(
  "C:/Users/baitu/OneDrive - Texas State University/",
  "Das, Subasish's files - Hackathon_Anika/",
  "Papers/TRBAM 2027/",
  "6. bike_Travel lane vs crossing_FARS/",
  "02_Code"
)

DATA_FILE <- file.path(
  PROJECT_DIR,
  "Bike01_copy_no duplicate.xlsx"
)

SHEET_NAME <- "Lump"

TARGET_VARIABLE <- "Bicyclist Path Type"

RF_SEED <- 2027
RF_TREES <- 2000
RF_NODE_SIZE <- 1

# Used only to create a separate selected-variable table.
# All non-target variables are still included in the Random Forest.
N_SELECTED_VARIABLES <- 13

FONT_FAMILY <- "serif"

RF_PALETTE <- c(
  "#CBB4A7",
  "#A56A6C",
  "#7C2A35",
  "#4B1D24"
)

OUTPUT_DIR <- file.path(
  PROJECT_DIR,
  "Analysis_Output",
  "01_RF_Variable_Selection"
)

dir.create(
  OUTPUT_DIR,
  recursive = TRUE,
  showWarnings = FALSE
)

# ============================================================
# LOAD MAIN RANDOM FOREST DATA
#
# This chunk restores bike_data before verify-variables.
# It uses the existing DATA_FILE and SHEET_NAME from Project settings.
# It does NOT use the annual distribution sheet and does NOT change
# the Random Forest analysis logic.
# ============================================================

required_packages_core <- c(
  "readxl",
  "dplyr",
  "stringr",
  "tibble",
  "randomForest",
  "ggplot2",
  "scales",
  "writexl"
)

packages_to_install_core <- required_packages_core[
  !required_packages_core %in% rownames(installed.packages())
]

if (length(packages_to_install_core) > 0) {
  install.packages(
    packages_to_install_core,
    repos = "https://cloud.r-project.org",
    dependencies = TRUE
  )
}

invisible(
  lapply(
    required_packages_core,
    library,
    character.only = TRUE
  )
)

if (!file.exists(DATA_FILE)) {
  stop(
    paste0(
      "Main Random Forest data file not found:\n",
      DATA_FILE
    )
  )
}

available_sheets_main <- readxl::excel_sheets(DATA_FILE)

if (!SHEET_NAME %in% available_sheets_main) {
  stop(
    paste0(
      "The main Random Forest sheet was not found: ",
      SHEET_NAME,
      "\nAvailable sheets are:\n",
      paste(available_sheets_main, collapse = "\n")
    )
  )
}

bike_data <- readxl::read_excel(
  path = DATA_FILE,
  sheet = SHEET_NAME,
  guess_max = 100000
)

cat("Main Random Forest data loaded successfully.\n")

cat("File:", DATA_FILE, "\n")

cat("Sheet:", SHEET_NAME, "\n")

cat("Rows:", nrow(bike_data), "\n")

cat("Columns:", ncol(bike_data), "\n")

EXPECTED_VARIABLES <- c(
  "Hit and Run",
  "Road Relation",
  "Weather",
  "Lighting Condition",
  "Rural/Urban",
  "Bicyclist Path Type",
  "Intersection Related",
  "Trafficway Configuration",
  "Road Alignment",
  "Surface Condition",
  "Vehicle Movement",
  "Bicyclist age",
  "Driver Age",
  "Daytime",
  "Vehicle Type",
  "Speed Limit",
  "Contributing Factors"
)

missing_variables <- setdiff(
  EXPECTED_VARIABLES,
  names(bike_data)
)

if (length(missing_variables) > 0) {
  stop(
    paste0(
      "The following required Excel variable(s) are missing:\n",
      paste(
        missing_variables,
        collapse = ", "
      )
    )
  )
}

if (!identical(names(bike_data), EXPECTED_VARIABLES)) {
  stop(
    paste0(
      "The variables or their order do not exactly match the Lump worksheet ",
      "structure used to prepare this code."
    )
  )
}

knitr::kable(
  tibble::tibble(
    `Excel Variable Name` = names(bike_data)
  ),
  caption = "Variables imported exactly as written in the Lump worksheet"
)

clean_text <- function(x) {
  x <- as.character(x)
  x <- stringr::str_squish(x)

  x[
    is.na(x) |
      x == ""
  ] <- "Not Reported"

  x
}

# Clean text spacing only. No categories are merged or renamed.
bike_data <- bike_data |>
  dplyr::mutate(
    dplyr::across(
      where(is.character),
      clean_text
    )
  )

approved_target_categories <- c(
  "Travel Lane",
  "Crossing Lane"
)

unexpected_target_categories <- setdiff(
  unique(
    bike_data[[TARGET_VARIABLE]]
  ),
  approved_target_categories
)

if (length(unexpected_target_categories) > 0) {
  stop(
    paste0(
      "Unexpected target category or categories found: ",
      paste(
        unexpected_target_categories,
        collapse = ", "
      )
    )
  )
}

# Explicit classification outcome.
rf_target <- factor(
  bike_data[[TARGET_VARIABLE]],
  levels = c(
    "Travel Lane",
    "Crossing Lane"
  )
)

if (!is.factor(rf_target)) {
  stop("The Random Forest outcome is not a factor.")
}

if (nlevels(rf_target) != 2) {
  stop("The Random Forest outcome must contain exactly two classes.")
}

# Automatically use EVERY Excel variable except the target.
RF_PREDICTORS <- setdiff(
  names(bike_data),
  TARGET_VARIABLE
)

rf_predictor_data <- bike_data[
  ,
  RF_PREDICTORS,
  drop = FALSE
]

# The Lump worksheet contains recoded categorical predictors.
# Convert every predictor to a factor while preserving its original name.
rf_predictor_data[] <- lapply(
  rf_predictor_data,
  function(x) {
    factor(
      clean_text(x)
    )
  }
)

# Confirm that all 16 non-target variables are present.
if (ncol(rf_predictor_data) != ncol(bike_data) - 1) {
  stop("Not all non-target variables entered the Random Forest.")
}

if (!setequal(
  names(rf_predictor_data),
  setdiff(
    names(bike_data),
    TARGET_VARIABLE
  )
)) {
  stop("The Random Forest predictor names do not match the Excel variables.")
}

cat(
  "Target variable:",
  TARGET_VARIABLE,
  "\n"
)

cat(
  "Number of predictors included:",
  ncol(rf_predictor_data),
  "\n"
)

cat(
  "Predictors included:\n",
  paste(
    names(rf_predictor_data),
    collapse = "\n"
  ),
  "\n"
)

outcome_distribution <- tibble::tibble(
  `Bicyclist Path Type` = rf_target
) |>
  dplyr::count(
    `Bicyclist Path Type`,
    name = "Count"
  ) |>
  dplyr::mutate(
    Percentage = round(
      100 * Count / sum(Count),
      2
    )
  )

knitr::kable(
  outcome_distribution,
  caption = "Travel Lane and Crossing Lane outcome distribution"
)

set.seed(RF_SEED)

rf_mtry <- max(
  1,
  floor(
    sqrt(
      ncol(
        rf_predictor_data
      )
    )
  )
)

rf_model <- randomForest::randomForest(
  x = rf_predictor_data,
  y = rf_target,
  ntree = RF_TREES,
  mtry = rf_mtry,
  nodesize = RF_NODE_SIZE,
  importance = TRUE,
  keep.forest = TRUE
)

cat("Model type: Classification\n")

cat("Number of trees:", RF_TREES, "\n")

cat("mtry:", rf_mtry, "\n")

cat("Predictors used:", ncol(rf_predictor_data), "\n")

importance_matrix <- randomForest::importance(
  rf_model,
  scale = FALSE
)

required_importance_columns <- c(
  "MeanDecreaseAccuracy",
  "MeanDecreaseGini"
)

missing_importance_columns <- setdiff(
  required_importance_columns,
  colnames(
    importance_matrix
  )
)

if (length(missing_importance_columns) > 0) {
  stop(
    paste0(
      "The following Random Forest importance column(s) are missing: ",
      paste(
        missing_importance_columns,
        collapse = ", "
      )
    )
  )
}

rf_importance <- tibble::tibble(
  `Variable Name` = rownames(
    importance_matrix
  ),
  `Mean Decrease Accuracy` = as.numeric(
    importance_matrix[
      ,
      "MeanDecreaseAccuracy"
    ]
  ),
  `Mean Decrease Gini` = as.numeric(
    importance_matrix[
      ,
      "MeanDecreaseGini"
    ]
  )
) |>
  dplyr::arrange(
    dplyr::desc(
      `Mean Decrease Gini`
    )
  ) |>
  dplyr::mutate(
    `Normalized Gini Importance` = round(
      100 *
        `Mean Decrease Gini` /
        max(
          `Mean Decrease Gini`,
          na.rm = TRUE
        ),
      2
    ),
    Rank = dplyr::row_number()
  ) |>
  dplyr::select(
    Rank,
    `Variable Name`,
    `Mean Decrease Accuracy`,
    `Mean Decrease Gini`,
    `Normalized Gini Importance`
  )

# Confirm that every non-target Excel variable appears in the importance output.
missing_from_importance <- setdiff(
  RF_PREDICTORS,
  rf_importance$`Variable Name`
)

extra_in_importance <- setdiff(
  rf_importance$`Variable Name`,
  RF_PREDICTORS
)

if (length(missing_from_importance) > 0) {
  stop(
    paste0(
      "The following predictor(s) are missing from the RF importance output: ",
      paste(
        missing_from_importance,
        collapse = ", "
      )
    )
  )
}

if (length(extra_in_importance) > 0) {
  stop(
    paste0(
      "Unexpected variable(s) appear in the RF importance output: ",
      paste(
        extra_in_importance,
        collapse = ", "
      )
    )
  )
}

if (nrow(rf_importance) != length(RF_PREDICTORS)) {
  stop(
    "The importance table does not contain all non-target variables."
  )
}

knitr::kable(
  rf_importance,
  digits = 3,
  caption = "Original Random Forest variable-importance results"
)

rf_plot_data <- rf_importance |>
  dplyr::mutate(
    `Variable Name` = factor(
      `Variable Name`,
      levels = rev(
        `Variable Name`
      )
    )
  )

rf_importance_plot <- ggplot2::ggplot(
  rf_plot_data,
  ggplot2::aes(
    x = `Mean Decrease Gini`,
    y = `Variable Name`,
    fill = `Mean Decrease Gini`
  )
) +
  ggplot2::geom_col(
    width = 0.68,
    color = "black",
    linewidth = 0.25
  ) +
  ggplot2::geom_text(
    ggplot2::aes(
      label = sprintf(
        "%.1f",
        `Mean Decrease Gini`
      )
    ),
    hjust = -0.12,
    family = FONT_FAMILY,
    fontface = "bold",
    size = 3.7,
    color = "black"
  ) +
  ggplot2::scale_fill_gradientn(
    colours = RF_PALETTE,
    guide = "none"
  ) +
  ggplot2::scale_x_continuous(
    expand = ggplot2::expansion(
      mult = c(
        0,
        0.12
      )
    )
  ) +
  ggplot2::labs(
    title = "Random Forest Variable Importance",
    subtitle = "Travel Lane versus Crossing Lane fatal bicycle crashes",
    x = "Mean Decrease Gini",
    y = NULL
  ) +
  ggplot2::theme_minimal(
    base_family = FONT_FAMILY,
    base_size = 12
  ) +
  ggplot2::theme(
    plot.title = ggplot2::element_text(
      face = "bold",
      size = 16,
      hjust = 0
    ),
    plot.subtitle = ggplot2::element_text(
      size = 11,
      hjust = 0,
      margin = ggplot2::margin(
        b = 10
      )
    ),
    axis.text.y = ggplot2::element_text(
      face = "bold",
      size = 10.5,
      color = "black"
    ),
    axis.text.x = ggplot2::element_text(
      size = 10,
      color = "black"
    ),
    axis.title.x = ggplot2::element_text(
      face = "bold",
      size = 11,
      margin = ggplot2::margin(
        t = 10
      )
    ),
    panel.grid.major.y = ggplot2::element_blank(),
    panel.grid.minor = ggplot2::element_blank(),
    panel.grid.major.x = ggplot2::element_line(
      color = "grey85",
      linewidth = 0.35
    ),
    plot.background = ggplot2::element_rect(
      fill = "white",
      color = NA
    ),
    panel.background = ggplot2::element_rect(
      fill = "white",
      color = NA
    ),
    plot.margin = ggplot2::margin(
      t = 12,
      r = 35,
      b = 12,
      l = 12
    )
  )

rf_importance_plot

selected_variable_table <- rf_importance |>
  dplyr::slice_head(
    n = min(
      N_SELECTED_VARIABLES,
      nrow(
        rf_importance
      )
    )
  )

selected_variables <- selected_variable_table$`Variable Name`

knitr::kable(
  selected_variable_table,
  digits = 3,
  caption = paste0(
    "Top ",
    length(
      selected_variables
    ),
    " variables based on original Mean Decrease Gini importance"
  )
)

writexl::write_xlsx(
  list(
    Outcome_Distribution = outcome_distribution,
    All_RF_Importance = rf_importance,
    Selected_Variables = selected_variable_table
  ),
  path = file.path(
    OUTPUT_DIR,
    "RF_Original_Variable_Importance.xlsx"
  )
)

utils::write.csv(
  rf_importance,
  file = file.path(
    OUTPUT_DIR,
    "RF_Original_Variable_Importance.csv"
  ),
  row.names = FALSE
)

utils::write.csv(
  selected_variable_table,
  file = file.path(
    OUTPUT_DIR,
    "RF_Selected_Variables.csv"
  ),
  row.names = FALSE
)

ggplot2::ggsave(
  filename = file.path(
    OUTPUT_DIR,
    "RF_Original_Variable_Importance.png"
  ),
  plot = rf_importance_plot,
  width = 10,
  height = 9,
  units = "in",
  dpi = 600,
  bg = "white"
)

ggplot2::ggsave(
  filename = file.path(
    OUTPUT_DIR,
    "RF_Original_Variable_Importance.pdf"
  ),
  plot = rf_importance_plot,
  width = 10,
  height = 9,
  units = "in",
  device = grDevices::cairo_pdf,
  bg = "white"
)

saveRDS(
  rf_model,
  file = file.path(
    OUTPUT_DIR,
    "RF_Model_All_Predictors.rds"
  )
)

saveRDS(
  rf_predictor_data,
  file = file.path(
    OUTPUT_DIR,
    "RF_All_Predictor_Data.rds"
  )
)

saveRDS(
  selected_variables,
  file = file.path(
    OUTPUT_DIR,
    "RF_Selected_Variable_Names.rds"
  )
)

cat(
  "Random Forest outputs were saved to:\n",
  OUTPUT_DIR,
  "\n"
)

# ============================================================
# FAST TREESHAP ANALYSIS
# Sheet: Selected Variables
# Target: Bicyclist Path Type
# Explained outcome: Crossing Lane probability
# Output:
# 1. SHAP matrix for SOM
# 2. Reference-paper-style 4-panel SHAP figure
# 3. SHAP variable-importance plot
# ============================================================

options(pkgType = "binary")

required_packages <- c(
  "readxl",
  "dplyr",
  "stringr",
  "ranger",
  "treeshap",
  "ggplot2",
  "gridExtra",
  "grid",
  "writexl"
)

packages_to_install <- required_packages[
  !required_packages %in% rownames(installed.packages())
]

if (length(packages_to_install) > 0) {
  install.packages(
    packages_to_install,
    repos = "https://cloud.r-project.org",
    type = "binary",
    dependencies = TRUE
  )
}

invisible(
  lapply(
    required_packages,
    library,
    character.only = TRUE
  )
)

# ----------------------------
# Settings
# ----------------------------

PROJECT_DIR <- paste0(
  "C:/Users/baitu/OneDrive - Texas State University/",
  "Das, Subasish's files - Hackathon_Anika/",
  "Papers/TRBAM 2027/",
  "6. bike_Travel lane vs crossing_FARS/",
  "02_Code"
)

DATA_FILE <- file.path(
  PROJECT_DIR,
  "Bike01_copy_no duplicate.xlsx"
)

SHEET_NAME <- "Selected Variables"
TARGET_VARIABLE <- "Bicyclist Path Type"

NEGATIVE_CLASS <- "Travel Lane"
POSITIVE_CLASS <- "Crossing Lane"

EXCLUDED_VARIABLES <- c(
  "Hit and Run",
  "Road Relation"
)

RF_SEED <- 2027
TRAIN_PROPORTION <- 0.80
RF_TREES <- 500

MAROON_POS <- "#7C2A35"
BROWN_NEG <- "#CBB4A7"
MAROON_DARK <- "#4B1D24"
GRID_COLOR <- "grey85"

OUT_DIR <- file.path(
  PROJECT_DIR,
  "Analysis_Output",
  "02_TreeSHAP_Reference_Style"
)

dir.create(
  OUT_DIR,
  recursive = TRUE,
  showWarnings = FALSE
)

clean_text <- function(x) {
  x <- as.character(x)
  x <- stringr::str_squish(x)
  x[is.na(x) | x == ""] <- "Not Reported"
  x
}

# ----------------------------
# Read data
# ----------------------------

bike_data <- readxl::read_excel(
  path = DATA_FILE,
  sheet = SHEET_NAME,
  guess_max = 100000
)

bike_data <- bike_data |>
  dplyr::mutate(
    dplyr::across(
      where(is.character),
      clean_text
    )
  )

if (!TARGET_VARIABLE %in% names(bike_data)) {
  stop(
    paste0(
      "Target variable not found: ",
      TARGET_VARIABLE
    )
  )
}

bike_data[[TARGET_VARIABLE]] <- factor(
  clean_text(
    bike_data[[TARGET_VARIABLE]]
  ),
  levels = c(
    NEGATIVE_CLASS,
    POSITIVE_CLASS
  )
)

bike_data <- bike_data[
  !is.na(
    bike_data[[TARGET_VARIABLE]]
  ),
  ,
  drop = FALSE
]

FINAL_PREDICTORS <- setdiff(
  names(bike_data),
  c(
    TARGET_VARIABLE,
    EXCLUDED_VARIABLES
  )
)

model_data <- bike_data[
  ,
  c(
    FINAL_PREDICTORS,
    TARGET_VARIABLE
  ),
  drop = FALSE
]

for (v in FINAL_PREDICTORS) {
  if (is.character(model_data[[v]])) {
    model_data[[v]] <- factor(
      clean_text(
        model_data[[v]]
      )
    )
  }
}

model_data <- model_data[
  complete.cases(model_data),
  ,
  drop = FALSE
]

cat("Rows used:", nrow(model_data), "\n")

cat("Predictors used:", length(FINAL_PREDICTORS), "\n")

cat(
  paste(
    FINAL_PREDICTORS,
    collapse = "\n"
  )
)

cat("\n")

# ----------------------------
# Stratified 80/20 split
# ----------------------------

set.seed(RF_SEED)

row_id <- seq_len(
  nrow(model_data)
)

split_ids <- split(
  row_id,
  model_data[[TARGET_VARIABLE]]
)

train_index <- unlist(
  lapply(
    split_ids,
    function(ids) {
      sample(
        ids,
        size = floor(
          TRAIN_PROPORTION * length(ids)
        )
      )
    }
  )
)

test_index <- setdiff(
  row_id,
  train_index
)

train_data <- model_data[
  train_index,
  ,
  drop = FALSE
]

test_data <- model_data[
  test_index,
  ,
  drop = FALSE
]

cat("Training rows:", nrow(train_data), "\n")

cat("Testing rows:", nrow(test_data), "\n")

cat("Training distribution:\n")

print(table(train_data[[TARGET_VARIABLE]]))

cat("Testing distribution:\n")

print(table(test_data[[TARGET_VARIABLE]]))

# ----------------------------
# Safe internal names and dummy encoding
# ----------------------------

name_key <- data.frame(
  Original = FINAL_PREDICTORS,
  Safe = make.names(
    FINAL_PREDICTORS,
    unique = TRUE
  ),
  stringsAsFactors = FALSE
)

x_all_original <- model_data[
  ,
  FINAL_PREDICTORS,
  drop = FALSE
]

x_all_safe <- x_all_original

names(x_all_safe) <- name_key$Safe

x_all_dummy <- model.matrix(
  ~ . - 1,
  data = x_all_safe
)

x_all_dummy <- as.data.frame(
  x_all_dummy
)

x_train_dummy <- x_all_dummy[
  train_index,
  ,
  drop = FALSE
]

x_test_dummy <- x_all_dummy[
  test_index,
  ,
  drop = FALSE
]

y_train_class <- droplevels(
  train_data[[TARGET_VARIABLE]]
)

y_test_class <- droplevels(
  test_data[[TARGET_VARIABLE]]
)

y_train_numeric <- ifelse(
  y_train_class == POSITIVE_CLASS,
  1,
  0
)

# ----------------------------
# Final Random Forest
# Regression forest on 0/1 outcome gives predicted Crossing Lane probability.
# This avoids treeshap probability-classification error.
# ----------------------------

rf_train_data <- x_train_dummy
rf_train_data$Crossing_Lane_Probability_Target <- y_train_numeric

set.seed(RF_SEED)

rf_final <- ranger::ranger(
  dependent.variable.name = "Crossing_Lane_Probability_Target",
  data = rf_train_data,
  num.trees = RF_TREES,
  mtry = max(
    1,
    floor(
      sqrt(
        ncol(x_train_dummy)
      )
    )
  ),
  seed = RF_SEED,
  importance = "impurity",
  num.threads = max(
    1,
    parallel::detectCores() - 1
  )
)

test_pred_prob <- predict(
  rf_final,
  data = x_test_dummy
)$predictions

test_pred_prob <- pmin(
  pmax(
    test_pred_prob,
    0
  ),
  1
)

test_pred_class <- ifelse(
  test_pred_prob >= 0.5,
  POSITIVE_CLASS,
  NEGATIVE_CLASS
)

prediction_output <- data.frame(
  Original_Row_ID = test_index,
  Actual = as.character(
    y_test_class
  ),
  Predicted = as.character(
    test_pred_class
  ),
  Predicted_Probability_Crossing_Lane = as.numeric(
    test_pred_prob
  ),
  stringsAsFactors = FALSE
)

cat("Random Forest completed.\n")

# ----------------------------
# Fast TreeSHAP for all test observations
# ----------------------------

cat("Starting TreeSHAP for all test observations...\n")

unified_rf <- treeshap::ranger.unify(
  rf_final,
  x_train_dummy
)

treeshap_output <- treeshap::treeshap(
  unified_rf,
  x_test_dummy,
  verbose = TRUE
)

shap_encoded <- as.data.frame(
  treeshap_output$shaps
)

names(shap_encoded) <- names(
  x_test_dummy
)

cat("TreeSHAP completed.\n")

cat("SHAP rows:", nrow(shap_encoded), "\n")

cat("Encoded SHAP columns:", ncol(shap_encoded), "\n")

# ----------------------------
# Aggregate encoded SHAP back to original variables
# ----------------------------

shap_original <- data.frame(
  matrix(
    0,
    nrow = nrow(shap_encoded),
    ncol = length(FINAL_PREDICTORS)
  )
)

names(shap_original) <- FINAL_PREDICTORS

for (i in seq_len(nrow(name_key))) {

  original_name <- name_key$Original[i]
  safe_name <- name_key$Safe[i]

  matched_cols <- names(shap_encoded)[
    startsWith(
      names(shap_encoded),
      safe_name
    )
  ]

  if (length(matched_cols) > 0) {
    shap_original[[original_name]] <- rowSums(
      shap_encoded[
        ,
        matched_cols,
        drop = FALSE
      ],
      na.rm = TRUE
    )
  }
}

# ----------------------------
# SHAP matrix for SOM
# ----------------------------

shap_matrix_for_som <- cbind(
  prediction_output,
  shap_original
)

# ----------------------------
# SHAP importance table
# ----------------------------

shap_importance <- data.frame(
  Variable_Name = names(shap_original),
  Mean_Absolute_SHAP = colMeans(
    abs(
      shap_original
    ),
    na.rm = TRUE
  ),
  Mean_SHAP = colMeans(
    shap_original,
    na.rm = TRUE
  ),
  stringsAsFactors = FALSE
) |>
  dplyr::arrange(
    dplyr::desc(
      Mean_Absolute_SHAP
    )
  ) |>
  dplyr::mutate(
    Rank = dplyr::row_number(),
    Normalized_Mean_Absolute_SHAP = round(
      100 *
        Mean_Absolute_SHAP /
        max(
          Mean_Absolute_SHAP,
          na.rm = TRUE
        ),
      2
    )
  ) |>
  dplyr::select(
    Rank,
    Variable_Name,
    Mean_Absolute_SHAP,
    Mean_SHAP,
    Normalized_Mean_Absolute_SHAP
  )

print(
  shap_importance
)

# ----------------------------
# Reference-style SHAP plot data
# ----------------------------

# ============================================================
# CORRECTED SHAP PLOTTING SECTION
# Fixes:
# 1. No NA labels
# 2. Four separate large figures
# 3. Caption/heading below each figure
# 4. Original Excel variable names preserved
# ============================================================

MAROON_POS <- "#7C2A35"
BROWN_NEG <- "#CBB4A7"
GRID_COLOR <- "grey85"

# Keep original Excel variable names exactly
test_predictor_values <- test_data[, FINAL_PREDICTORS, drop = FALSE]

for (v in FINAL_PREDICTORS) {
  test_predictor_values[[v]] <- as.character(test_predictor_values[[v]])
  test_predictor_values[[v]][is.na(test_predictor_values[[v]]) | test_predictor_values[[v]] == ""] <- "Not Reported"
}

# Prefix SHAP columns, but keep original variable names after prefix
shap_prefixed <- shap_original
names(shap_prefixed) <- paste0("SHAP__", names(shap_original))

plot_data_base <- data.frame(
  prediction_output,
  test_predictor_values,
  shap_prefixed,
  check.names = FALSE,
  stringsAsFactors = FALSE
)

plot_data_base$CM_Group <- NA_character_

plot_data_base$CM_Group[
  plot_data_base$Predicted == NEGATIVE_CLASS &
    plot_data_base$Actual == NEGATIVE_CLASS
] <- "TN"

plot_data_base$CM_Group[
  plot_data_base$Predicted == NEGATIVE_CLASS &
    plot_data_base$Actual == POSITIVE_CLASS
] <- "FN"

plot_data_base$CM_Group[
  plot_data_base$Predicted == POSITIVE_CLASS &
    plot_data_base$Actual == NEGATIVE_CLASS
] <- "FP"

plot_data_base$CM_Group[
  plot_data_base$Predicted == POSITIVE_CLASS &
    plot_data_base$Actual == POSITIVE_CLASS
] <- "TP"

group_info <- data.frame(
  CM_Group = c("TN", "FN", "FP", "TP"),
  Bottom_Heading = c(
    "TN: True negative | Predicted Travel Lane, actual Travel Lane",
    "FN: False negative | Predicted Travel Lane, actual Crossing Lane",
    "FP: False positive | Predicted Crossing Lane, actual Travel Lane",
    "TP: True positive | Predicted Crossing Lane, actual Crossing Lane"
  ),
  File_Name = c(
    "SHAP_TN_True_Negative.png",
    "SHAP_FN_False_Negative.png",
    "SHAP_FP_False_Positive.png",
    "SHAP_TP_True_Positive.png"
  ),
  stringsAsFactors = FALSE
)

make_feature_summary_label <- function(data_group, variable_name) {

  x <- data_group[[variable_name]]

  if (is.null(x)) {
    return(paste0(variable_name, " = Not Reported: 0%"))
  }

  x <- as.character(x)
  x[is.na(x) | x == "" | x == "NA"] <- "Not Reported"

  tab <- sort(
    table(x),
    decreasing = TRUE
  )

  if (length(tab) == 0) {
    return(paste0(variable_name, " = Not Reported: 0%"))
  }

  top_level <- names(tab)[1]

  if (is.na(top_level) || top_level == "" || top_level == "NA") {
    top_level <- "Not Reported"
  }

  pct <- round(
    100 * as.numeric(tab[1]) / length(x),
    0
  )

  if (is.na(pct)) {
    pct <- 0
  }

  paste0(
    variable_name,
    " = ",
    top_level,
    ": ",
    pct,
    "%"
  )
}

make_single_group_plot <- function(group_name, bottom_heading) {

  group_data <- plot_data_base[
    plot_data_base$CM_Group == group_name,
    ,
    drop = FALSE
  ]

  n_group <- nrow(group_data)

  shap_cols <- paste0("SHAP__", names(shap_original))

  shap_means <- data.frame(
    Variable_Name = names(shap_original),
    Mean_SHAP = colMeans(
      group_data[, shap_cols, drop = FALSE],
      na.rm = TRUE
    ),
    stringsAsFactors = FALSE
  )

  shap_means$Label <- sapply(
    shap_means$Variable_Name,
    function(v) {
      make_feature_summary_label(group_data, v)
    }
  )

  shap_means <- shap_means |>
    dplyr::arrange(Mean_SHAP)

  shap_means$Label <- factor(
    shap_means$Label,
    levels = shap_means$Label
  )

  full_bottom_heading <- paste0(
    bottom_heading,
    " | n = ",
    format(n_group, big.mark = ",")
  )

  p_bar <- ggplot2::ggplot(
    shap_means,
    ggplot2::aes(
      x = Mean_SHAP,
      y = Label,
      fill = Mean_SHAP > 0
    )
  ) +
    ggplot2::geom_col(
      color = "black",
      linewidth = 0.35,
      width = 0.68
    ) +
    ggplot2::scale_fill_manual(
      values = c(
        "TRUE" = MAROON_POS,
        "FALSE" = BROWN_NEG
      ),
      guide = "none"
    ) +
    ggplot2::geom_vline(
      xintercept = 0,
      color = "black",
      linewidth = 0.45
    ) +
    ggplot2::labs(
      x = "Mean SHAP value",
      y = NULL
    ) +
    ggplot2::theme_minimal(
      base_family = "serif",
      base_size = 15
    ) +
    ggplot2::theme(
      axis.text.y = ggplot2::element_text(
        face = "bold",
        color = "black",
        size = 13
      ),
      axis.text.x = ggplot2::element_text(
        color = "black",
        size = 13
      ),
      axis.title.x = ggplot2::element_text(
        face = "bold",
        size = 15
      ),
      panel.grid.major.y = ggplot2::element_blank(),
      panel.grid.minor = ggplot2::element_blank(),
      panel.grid.major.x = ggplot2::element_line(
        color = GRID_COLOR,
        linewidth = 0.35
      ),
      plot.margin = ggplot2::margin(
        t = 8,
        r = 10,
        b = 4,
        l = 8
      )
    )

  p_box <- ggplot2::ggplot(
    group_data,
    ggplot2::aes(
      x = "",
      y = Predicted_Probability_Crossing_Lane
    )
  ) +
    ggplot2::geom_boxplot(
      fill = "white",
      color = "black",
      width = 0.32,
      outlier.size = 1.0,
      linewidth = 0.45
    ) +
    ggplot2::coord_flip() +
    ggplot2::scale_y_continuous(
      limits = c(0, 1),
      breaks = seq(0, 1, 0.25)
    ) +
    ggplot2::labs(
      x = NULL,
      y = "Predicted probability of Crossing Lane"
    ) +
    ggplot2::theme_minimal(
      base_family = "serif",
      base_size = 13
    ) +
    ggplot2::theme(
      axis.text.y = ggplot2::element_blank(),
      axis.ticks.y = ggplot2::element_blank(),
      axis.text.x = ggplot2::element_text(
        color = "black",
        size = 12
      ),
      axis.title.x = ggplot2::element_text(
        face = "bold",
        size = 13
      ),
      panel.grid.major.y = ggplot2::element_blank(),
      panel.grid.minor = ggplot2::element_blank(),
      panel.grid.major.x = ggplot2::element_line(
        color = GRID_COLOR,
        linewidth = 0.30
      ),
      plot.margin = ggplot2::margin(
        t = 2,
        r = 10,
        b = 2,
        l = 8
      )
    )

  combined_plot <- gridExtra::arrangeGrob(
    p_bar,
    p_box,
    ncol = 1,
    heights = c(4.8, 1.0),
    bottom = grid::textGrob(
      full_bottom_heading,
      gp = grid::gpar(
        fontfamily = "serif",
        fontsize = 18,
        fontface = "bold"
      )
    )
  )

  combined_plot
}

separate_group_plots <- list()

for (i in seq_len(nrow(group_info))) {

  group_name <- group_info$CM_Group[i]
  bottom_heading <- group_info$Bottom_Heading[i]
  file_name <- group_info$File_Name[i]

  this_plot <- make_single_group_plot(
    group_name = group_name,
    bottom_heading = bottom_heading
  )

  separate_group_plots[[group_name]] <- this_plot

  ggplot2::ggsave(
    filename = file.path(
      OUT_DIR,
      file_name
    ),
    plot = this_plot,
    width = 13,
    height = 9,
    dpi = 600,
    bg = "white"
  )
}

# Show the four figures one by one inside R Markdown
grid::grid.newpage()
grid::grid.draw(separate_group_plots[["TN"]])

grid::grid.newpage()
grid::grid.draw(separate_group_plots[["FN"]])

grid::grid.newpage()
grid::grid.draw(separate_group_plots[["FP"]])

grid::grid.newpage()
grid::grid.draw(separate_group_plots[["TP"]])

cat("Four separate corrected SHAP figures saved to:\n")

cat(OUT_DIR, "\n")

# ----------------------------
# Save outputs
# ----------------------------

writexl::write_xlsx(
  list(
    Predictions = prediction_output,
    TreeSHAP_Importance = shap_importance,
    SHAP_Matrix_for_SOM = shap_matrix_for_som
  ),
  file.path(
    OUT_DIR,
    "TreeSHAP_Output_All_Test_Observations.xlsx"
  )
)

write.csv(
  shap_matrix_for_som,
  file.path(
    OUT_DIR,
    "SHAP_Matrix_for_SOM.csv"
  ),
  row.names = FALSE
)

write.csv(
  shap_importance,
  file.path(
    OUT_DIR,
    "TreeSHAP_Importance.csv"
  ),
  row.names = FALSE
)

# ----------------------------
# TreeSHAP variable importance plot
# ----------------------------

importance_plot_data <- shap_importance |>
  dplyr::mutate(
    Variable_Name = factor(
      Variable_Name,
      levels = rev(Variable_Name)
    )
  )

importance_plot <- ggplot2::ggplot(
  importance_plot_data,
  ggplot2::aes(
    x = Mean_Absolute_SHAP,
    y = Variable_Name
  )
) +
  ggplot2::geom_col(
    fill = "#7C2A35",
    color = "black",
    linewidth = 0.35,
    width = 0.68
  ) +
  ggplot2::labs(
    title = "TreeSHAP Variable Importance",
    x = "Mean absolute SHAP value",
    y = NULL
  ) +
  ggplot2::theme_minimal(
    base_family = "serif",
    base_size = 14
  ) +
  ggplot2::theme(
    plot.title = ggplot2::element_text(
      face = "bold",
      size = 18,
      color = "black"
    ),
    axis.text.y = ggplot2::element_text(
      face = "bold",
      color = "black",
      size = 12
    ),
    axis.text.x = ggplot2::element_text(
      color = "black",
      size = 11
    ),
    axis.title.x = ggplot2::element_text(
      face = "bold",
      size = 13
    ),
    panel.grid.major.y = ggplot2::element_blank(),
    panel.grid.minor = ggplot2::element_blank()
  )

print(importance_plot)

ggplot2::ggsave(
  file.path(
    OUT_DIR,
    "TreeSHAP_Variable_Importance.png"
  ),
  importance_plot,
  width = 10,
  height = 8,
  dpi = 600,
  bg = "white"
)


saveRDS(
  shap_original,
  file.path(
    OUT_DIR,
    "SHAP_Matrix_for_SOM.rds"
  )
)

saveRDS(
  rf_final,
  file.path(
    OUT_DIR,
    "Ranger_RF_TreeSHAP_Model.rds"
  )
)

cat(
  "TreeSHAP analysis complete. Outputs saved to:\n",
  OUT_DIR,
  "\n"
)

cat(
  "Use SHAP_Matrix_for_SOM.csv for SOM grid selection.\n"
)

# ============================================================
# SOM GRID-SIZE SELECTION
# Input: SHAP_Matrix_for_SOM.csv
# Purpose: select SOM grid size before final SOM training
# ============================================================

options(pkgType = "binary")

required_packages <- c(
  "kohonen",
  "dplyr",
  "ggplot2",
  "writexl"
)

packages_to_install <- required_packages[
  !required_packages %in% rownames(installed.packages())
]

if (length(packages_to_install) > 0) {
  install.packages(
    packages_to_install,
    repos = "https://cloud.r-project.org",
    type = "binary",
    dependencies = TRUE
  )
}

invisible(
  lapply(
    required_packages,
    library,
    character.only = TRUE
  )
)

# ------------------------------------------------------------
# Output folder
# ------------------------------------------------------------

SOM_OUT_DIR <- file.path(
  PROJECT_DIR,
  "Analysis_Output",
  "03_SOM_Grid_Selection"
)

dir.create(
  SOM_OUT_DIR,
  recursive = TRUE,
  showWarnings = FALSE
)

# ------------------------------------------------------------
# Read SHAP matrix
# ------------------------------------------------------------

SHAP_SOM_FILE <- file.path(
  OUT_DIR,
  "SHAP_Matrix_for_SOM.csv"
)

if (!file.exists(SHAP_SOM_FILE)) {
  stop(
    paste0(
      "SHAP_Matrix_for_SOM.csv not found here: ",
      SHAP_SOM_FILE
    )
  )
}

shap_som_data <- read.csv(
  SHAP_SOM_FILE,
  check.names = FALSE,
  stringsAsFactors = FALSE
)

cat("Rows in SHAP matrix:", nrow(shap_som_data), "\n")

cat("Columns in SHAP matrix:", ncol(shap_som_data), "\n")

# ------------------------------------------------------------
# Select only SHAP columns for SOM
# Do not use ID, actual class, predicted class, or probability
# ------------------------------------------------------------

metadata_columns <- c(
  "Original_Row_ID",
  "Actual",
  "Predicted",
  "Predicted_Probability_Crossing_Lane"
)

shap_columns <- setdiff(
  names(shap_som_data),
  metadata_columns
)

som_input_raw <- shap_som_data[, shap_columns, drop = FALSE]

# Force all SHAP columns to numeric
for (v in names(som_input_raw)) {
  som_input_raw[[v]] <- as.numeric(som_input_raw[[v]])
}

# Remove columns with all missing values
all_missing_columns <- names(som_input_raw)[
  sapply(
    som_input_raw,
    function(x) all(is.na(x))
  )
]

if (length(all_missing_columns) > 0) {
  som_input_raw <- som_input_raw[
    ,
    !names(som_input_raw) %in% all_missing_columns,
    drop = FALSE
  ]
}

# Replace remaining NA values with 0
som_input_raw[is.na(som_input_raw)] <- 0

# Remove zero-variance SHAP columns
zero_variance_columns <- names(som_input_raw)[
  sapply(
    som_input_raw,
    function(x) {
      stats::sd(x, na.rm = TRUE) == 0
    }
  )
]

if (length(zero_variance_columns) > 0) {
  som_input_raw <- som_input_raw[
    ,
    !names(som_input_raw) %in% zero_variance_columns,
    drop = FALSE
  ]
}

cat("SHAP variables used for SOM:\n")

cat(paste(names(som_input_raw), collapse = "\n"))

cat("\n")

# ------------------------------------------------------------
# Scale SHAP matrix for SOM grid selection
# ------------------------------------------------------------

som_input_scaled <- scale(
  as.matrix(som_input_raw)
)

som_input_scaled[is.na(som_input_scaled)] <- 0

# ------------------------------------------------------------
# Candidate grid sizes
# ------------------------------------------------------------

grid_candidates <- data.frame(
  xdim = c(3, 4, 4, 5, 5, 6, 6, 7),
  ydim = c(3, 3, 4, 4, 5, 5, 6, 6)
)

grid_candidates$Grid_Size <- paste0(
  grid_candidates$xdim,
  " x ",
  grid_candidates$ydim
)

# ------------------------------------------------------------
# Function to calculate SOM quality metrics
# ------------------------------------------------------------

calculate_som_metrics <- function(xdim, ydim, som_data, seed_value = 2027) {

  set.seed(seed_value)

  som_grid <- kohonen::somgrid(
    xdim = xdim,
    ydim = ydim,
    topo = "hexagonal"
  )

 som_model <- kohonen::som(
  X = som_data,
  grid = som_grid,
  rlen = 500,
  alpha = c(0.05, 0.01),
  keep.data = TRUE
)

  codebook <- som_model$codes[[1]]

  distance_to_units <- sapply(
    seq_len(nrow(codebook)),
    function(k) {
      rowSums(
        sweep(
          som_data,
          2,
          codebook[k, ],
          "-"
        )^2
      )
    }
  )

  ordered_units <- t(
    apply(
      distance_to_units,
      1,
      order
    )
  )

  bmu_1 <- ordered_units[, 1]
  bmu_2 <- ordered_units[, 2]

  first_distance <- distance_to_units[
    cbind(
      seq_len(nrow(som_data)),
      bmu_1
    )
  ]

  quantization_error <- mean(
    sqrt(first_distance)
  )

  unit_distance_matrix <- kohonen::unit.distances(
    som_grid
  )

  bmu_neighbor_distance <- unit_distance_matrix[
    cbind(
      bmu_1,
      bmu_2
    )
  ]

  topographic_error <- mean(
    bmu_neighbor_distance > 1
  )

  kaski_lagus_error <- mean(
    sqrt(first_distance) + bmu_neighbor_distance
  )

  node_counts <- table(
    factor(
      bmu_1,
      levels = seq_len(nrow(codebook))
    )
  )

  empty_nodes <- sum(node_counts == 0)

  result <- data.frame(
    xdim = xdim,
    ydim = ydim,
    Grid_Size = paste0(xdim, " x ", ydim),
    Total_Nodes = xdim * ydim,
    Quantization_Error = quantization_error,
    Topographic_Error = topographic_error,
    Kaski_Lagus_Error = kaski_lagus_error,
    Empty_Nodes = empty_nodes,
    Empty_Node_Percent = round(
      100 * empty_nodes / (xdim * ydim),
      2
    ),
    Minimum_Records_Per_Node = min(node_counts),
    Maximum_Records_Per_Node = max(node_counts),
    Mean_Records_Per_Node = mean(node_counts),
    stringsAsFactors = FALSE
  )

  result
}

# ------------------------------------------------------------
# Run grid-size evaluation
# ------------------------------------------------------------

som_grid_results <- list()

for (i in seq_len(nrow(grid_candidates))) {

  this_xdim <- grid_candidates$xdim[i]
  this_ydim <- grid_candidates$ydim[i]

  cat(
    "Running SOM grid:",
    this_xdim,
    "x",
    this_ydim,
    "\n"
  )

  som_grid_results[[i]] <- calculate_som_metrics(
    xdim = this_xdim,
    ydim = this_ydim,
    som_data = som_input_scaled,
    seed_value = 2027
  )
}

som_grid_selection_table <- dplyr::bind_rows(
  som_grid_results
)

# ------------------------------------------------------------
# Create normalized ranking score
# Lower score is better
# ------------------------------------------------------------

normalize_metric <- function(x) {
  if (max(x, na.rm = TRUE) == min(x, na.rm = TRUE)) {
    return(rep(0, length(x)))
  }

  (x - min(x, na.rm = TRUE)) /
    (max(x, na.rm = TRUE) - min(x, na.rm = TRUE))
}

som_grid_selection_table <- som_grid_selection_table |>
  dplyr::mutate(
    QE_Normalized = normalize_metric(Quantization_Error),
    TE_Normalized = normalize_metric(Topographic_Error),
    KLE_Normalized = normalize_metric(Kaski_Lagus_Error),
    Empty_Normalized = normalize_metric(Empty_Node_Percent),
    Balanced_Selection_Score =
      0.40 * QE_Normalized +
      0.30 * TE_Normalized +
      0.20 * KLE_Normalized +
      0.10 * Empty_Normalized
  ) |>
  dplyr::arrange(
    Balanced_Selection_Score
  ) |>
  dplyr::mutate(
    Rank = dplyr::row_number()
  ) |>
  dplyr::select(
    Rank,
    Grid_Size,
    xdim,
    ydim,
    Total_Nodes,
    Quantization_Error,
    Topographic_Error,
    Kaski_Lagus_Error,
    Empty_Nodes,
    Empty_Node_Percent,
    Minimum_Records_Per_Node,
    Maximum_Records_Per_Node,
    Mean_Records_Per_Node,
    Balanced_Selection_Score
  )

print(som_grid_selection_table)