## Fit Preliminary Models 

# Load packages 
library(readr)
library(dplyr)
library(purrr)
library(corrplot)
library(xgboost)
library(shapviz)
library(ggplot2)
library(corrplot)
library(patchwork)

# Load fit data 
fit_fly <- read_rds("fitting_dataset_flies.rds")
fit_lei <- read_rds("fitting_dataset_disease.rds")

# Average to NUTS2 
fit_fly <- fit_fly %>%
  group_by(NUTS_ID, year) %>%
  mutate(across(BIO01:elevation, ~ mean(.x, na.rm = TRUE))) %>%
  mutate(across(neglectus:tobbi, ~ max(.x, na.rm = TRUE))) %>%
  select(-x, -y) %>% ungroup() %>% 
  distinct()

fit_lei <- fit_lei %>%
  group_by(NUTS_ID, year) %>%
  mutate(across(BIO01:elevation, ~ mean(.x, na.rm = TRUE))) %>%
  mutate(across(cl:vl, ~ max(.x, na.rm = TRUE))) %>%
  select(-x, -y) %>% ungroup() %>% 
  distinct()

# Transform to a list
fit_list <- list(fit_fly, fit_lei)

## CORRELATIONS 

# Select columns 3 to 40 from each dataframe in the list
cor_list <- map(fit_list, ~ select(.x, 2:39))

# Calculate Pearson correlations
cor_list <- map(cor_list, ~ cor(.x, use = "complete.obs", method = "pearson"))

# Plot the correlation matrices
map(cor_list, ~ corrplot(.x, method = "circle", type = "upper", order = "hclust", 
                         tl.col = "black", tl.srt = 45))
# Save as table 
cor_list <- map(cor_list, reshape2::melt)

# Assign a cut off to prevent over-fitting 
cor_list <- map(cor_list, ~.x %>% 
  mutate(r_cutoff = if_else(value > 0.7 | value < -0.7, 1, 0)))

# Merge to one dataframe 
names(cor_list) <- c("fly_model", "leish_model")
cor_data <- imap_dfr(cor_list, ~ mutate(.x, source = .y))

# Save 
write_csv(cor_data, "correlation_results.csv")

## SHAP

# Create lists for each outcome I am interested in 
shap_list <- map(fit_list, ~.x %>% select(-year, -NUTS_ID, -CNTR_CODE, -ISO3))
outcomes_list <- list(
  df1 = c("neglectus", "perfiliewi", "ariasi", "perniciosus", "tobbi"),
  df2 = c("cl", "vl")
)
split_by_outcome <- function(df, outcomes) {
  common_cols <- setdiff(names(df), outcomes)
  set_names(
    map(outcomes, ~ df %>% select(all_of(common_cols), all_of(.x)) %>% rename(outcome = .x)),
    outcomes
  )
}
shap_list <- map2(shap_list, outcomes_list, split_by_outcome)
shap_list <- flatten(shap_list)

# Transform outcome variable to numeric
shap_list <- map(shap_list, ~ .x %>% mutate(across(where(is.character), as.numeric)))

# Define the function to train the xgboost model
calculate_shap <- function(shap_dat) {
  # Ensure the outcome column is numeric
  shap_dat <- shap_dat %>% mutate(outcome = as.numeric(outcome))
  
  # Convert to matrix for xgboost
  X <- as.matrix(shap_dat[, -ncol(shap_dat)])  
  y <- shap_dat$outcome  
  
  # Train xgboost model
  dtrain <- xgb.DMatrix(X, label = y)
  params <- list(objective = "binary:logistic")
  model <- xgb.train(params, dtrain, nrounds = 50)
  
  # Extract the SHAP values
  shap_values <- shapviz(model, X = X, X_pred = dtrain)
  mean_abs_shap <- apply(abs(shap_values$S), 2, mean)
  importance_df <- data.frame(
    Feature = colnames(shap_values$S),
    Importance = mean_abs_shap
  )
  row.names(importance_df) <- NULL
  
  return(importance_df)
}

shap_results <- map(shap_list, calculate_shap)

# Merge to one dataframe for saving 
importance_results <- imap_dfr(shap_results, ~ mutate(.x, model = .y))
write_csv(importance_results, "shap_values.csv")

# Make a plot for each model
shap_plots <- map(shap_results, ~
                    ggplot(.x, aes(x = reorder(Feature, Importance), y = Importance)) + 
                    geom_bar(stat = "identity", fill = "#FFA500") + 
                    coord_flip() + 
                    labs(x = "", y = "mean(|SHAP value|)") + 
                    theme_minimal() + 
                    geom_text(aes(label = round(Importance, 2)), hjust = -0.1, color = "black") +
                    theme(text = element_text(face = "bold"))
)

# Facet to one plot for saving 
shap_plots <- map2(shap_plots, names(shap_plots), ~ .x + ggtitle(.y))
reduce(shap_plots, `+`) + plot_layout(ncol = 2)

## FEATURE SUBSET 

# Subset to the top 10 features and identify those which are clustered according to correlation 

shap <- read_csv("shap_values.csv")
corr <- read_csv("correlation_results.csv")

# Filter to the top 10 features (based on importance)
# Make sure that it doesn't filter out all the bioclimatic variables 
feat_sub <- shap %>%
  group_by(model) %>%
  arrange(desc(Importance)) %>%
  # Get the top 5 bioclimatic/seasonal variables
  filter(grepl("^BIO[0-9]{2}$|^(hurs|sm)_(warmest|coldest|driest|wettest)$", Feature)) %>%
  slice_head(n = 5) %>%
  ungroup() -> top_bioclim

# Now, select the remaining top variables (ensuring no more than 5 additional variables)
remaining_feat <- shap %>%
  group_by(model) %>%
  arrange(desc(Importance)) %>%
  filter(!grepl("^BIO[0-9]{2}$|^(hurs|sm)_(warmest|coldest|driest|wettest)$", Feature)) %>%
  slice_head(n = 5) %>%
  ungroup()

# Combine the two dataframes and ensure 10 total features
feat_sub <- bind_rows(top_bioclim, remaining_feat) %>%
  group_by(model) %>%
  arrange(desc(Importance)) %>%
  slice_head(n = 10) %>%
  ungroup()
rm(remaining_feat)

# Subset correlations to the top 10 features and those which meet the correlations threshold 
feat <- unique(feat_sub$Feature)
corr_sub <- corr %>% filter(Var1 %in% feat) %>% filter(r_cutoff == 1)

# Identify clusters for the features 

feat_sub$clusters <- c(1, 2, 3, 4, 5, 1, 6, 5, 1, 7, 1, 2, 3, 4, 5, 6, 4, 7, 7, 8, 
                       1, 2, 3, 4, 3, 5, 6, 5, 7, 7, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 
                       1, 2, 3, 4, 5, 6, 7, 8, 8, 9, 1, 2, 3, 4, 4, 5, 6, 5, 7, 8, 
                       1, 2, 3, 4, 5, 6, 7, 8, 6, 9)
feat_sub$Importance <- NULL

# Save selected features 
write_csv(feat_sub, "feature_subset.csv")


