## Identify Best Fit Model 

# Load packages 
library(readr)
library(dplyr)
library(rsample)
library(xgboost)
library(caret)
library(pROC) 

## PREPARING THE DATA 

# Load fitting dataset 
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()

# Creating my fitting list 
fit_list <- list(fit_fly, fit_lei)

# Subset further to each species and disease 
fit_list2 <- 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
  )
}
fit_list2 <- map2(fit_list2, outcomes_list, split_by_outcome)
fit_list2 <- flatten(fit_list2)

# Load my feature subset 
feats <- read_delim("feature_subset.csv", 
                    delim = ";", escape_double = FALSE, trim_ws = TRUE)

# Filter fitting datasets to the selected features 
subsetted_list <- lapply(names(fit_list2), function(model_name) {
  relevant_features <- feats %>%
    filter(model == model_name) %>%
    pull(Feature)
  df <- fit_list2[[model_name]]
  relevant_features <- c(relevant_features, "outcome")
  subsetted_df <- df[, relevant_features, drop = FALSE]
  return(subsetted_df)
})
names(subsetted_list) <- names(fit_list2)

# Generate all the unique formulas for each model 
generate_formulas <- function(feats, model_name) {
  # Subset feats for the relevant model
  relevant_feats <- feats[feats$model == model_name, ]
  
  formulas_list <- list()
  
  for (i in 2:max(relevant_feats$clusters)) { 
    cluster_combinations <- combn(unique(relevant_feats$clusters), i, simplify = FALSE)
    for (combo in cluster_combinations) {
      selected_features <- lapply(combo, function(c) {
        sample(relevant_feats$Feature[relevant_feats$clusters == c], 1)
      })
      formula <- paste("outcome ~", paste(unlist(selected_features), collapse = "+"))
      formulas_list[[length(formulas_list) + 1]] <- formula
    }
  }
  
  formulas_df <- data.frame(formula = unlist(formulas_list))
  formulas_df$formula_id <- 1:nrow(formulas_df)
  return(formulas_df)
}

# For each model, split data to testing and training and combine with the corresponding formulas 
process_subset <- function(data, feats, model_name) {
  split <- initial_split(data, prop = 0.75)
  train <- training(split)
  test  <- testing(split)
  
  formulas_df <- generate_formulas(feats, model_name)
  
  return(list(train = train, test = test, formulas = formulas_df))
}

fit_list3 <- lapply(names(subsetted_list), function(model_name) {
  process_subset(subsetted_list[[model_name]], feats, model_name)
})
names(fit_list3) <- names(subsetted_list)
  
## FINDING BEST FIT MODEL 

# Create an empty data frame to store model performance metrics
results <- data.frame()

# Set the number of k-folds 
k <- 5

# Loop over each model in fit_list3
for (model_name in names(fit_list3)) {
  # Access the data and formulas for the current model
  model_data <- fit_list3[[model_name]]$train
  model_test <- fit_list3[[model_name]]$test
  formulas_df <- fit_list3[[model_name]]$formulas
  
  # Loop through each formula in the formulas dataframe
  for (i in 1:nrow(formulas_df)) {
    formula_str <- formulas_df$formula[i]
    formula <- as.formula(formula_str)
    
    # Create model matrix for train and test
    model_data_matrix <- model.matrix(formula, data = model_data)
    test_data_matrix <- model.matrix(formula, data = model_test)
    
    target <- model_data$outcome
    target_test <- model_test$outcome
    predictors <- model_data_matrix[, -1]  # Exclude the intercept column
    predictors_test <- test_data_matrix[, -1]
    
    # Cross-validation
    folds <- createFolds(target, k = k, list = TRUE)
    metrics <- data.frame(rmse = numeric(k), sensitivity = numeric(k), specificity = numeric(k), auc = numeric(k))
    
    for (j in 1:k) {
      val_idx <- folds[[j]]
      train_idx <- setdiff(seq_len(nrow(predictors)), val_idx)
      
      train_X <- predictors[train_idx, , drop = FALSE]
      train_y <- target[train_idx]
      val_X <- predictors[val_idx, , drop = FALSE]
      val_y <- target[val_idx]
      
      # Fit the XGBoost model
      xgb_model <- xgboost(data = as.matrix(train_X), label = train_y, nrounds = 100, objective = "binary:logistic", verbose = 0)
      
      # Predict and calculate metrics
      preds <- predict(xgb_model, as.matrix(val_X))
      preds_class <- ifelse(preds > 0.5, 1, 0)
      
      metrics$rmse[j] <- sqrt(mean((preds - val_y)^2))
      confusion <- confusionMatrix(factor(preds_class), factor(val_y))
      metrics$sensitivity[j] <- confusion$byClass["Sensitivity"]
      metrics$specificity[j] <- confusion$byClass["Specificity"]
      roc_curve <- roc(val_y, preds)
      metrics$auc[j] <- auc(roc_curve)
    }
    
    # Train the final model on the full training set
    final_xgb_model <- xgboost(data = as.matrix(predictors), label = target, nrounds = 100, objective = "binary:logistic", verbose = 0)
    
    # Predict on the test set
    test_preds <- predict(final_xgb_model, as.matrix(predictors_test))
    test_preds_class <- ifelse(test_preds > 0.5, 1, 0)
    
    # Test set metrics
    test_rmse <- sqrt(mean((test_preds - target_test)^2))
    test_confusion <- confusionMatrix(factor(test_preds_class), factor(target_test))
    test_sensitivity <- test_confusion$byClass["Sensitivity"]
    test_specificity <- test_confusion$byClass["Specificity"]
    test_roc_curve <- roc(target_test, test_preds)
    test_auc <- auc(test_roc_curve)
    
    # Store the results for the current formula
    results <- rbind(results, data.frame(
      model = model_name,
      formula = formulas_df$formula_id[i], 
      cv_rmse = mean(metrics$rmse, na.rm = TRUE), 
      cv_sensitivity = mean(metrics$sensitivity, na.rm = TRUE), 
      cv_specificity = mean(metrics$specificity, na.rm = TRUE), 
      cv_auc = mean(metrics$auc, na.rm = TRUE),
      test_rmse = test_rmse,
      test_sensitivity = test_sensitivity,
      test_specificity = test_specificity,
      test_auc = test_auc
    ))
  }
}

# Save the full results 
row.names(results)<- NULL
formulas_df <- map_dfr(names(fit_list3), function(name) {
  fit_list3[[name]]$formulas %>%
    mutate(source = name) 
})
formulas_df <- formulas_df %>% rename("model" = source)
results <- results %>% rename("formula_id" = formula)
results <- left_join(results, formulas_df)
results <- results %>% 
  mutate(w_socio = if_else(stringr::str_detect(formula, "arope|population"), 1, 0)) %>% 
  mutate(w_bioclim = if_else(stringr::str_count(formula, "BIO(0[1-9]|1[0-9])|hurs_(warmest|driest|coldest|wettest)|sm_(warmest|driest|coldest|wettest)") >= 2, 1, 0))
write_csv(results, "model_performance_metrics.csv")

# Identify best models 
best_mods <- results %>% group_by(model, w_socio, w_bioclim) %>% 
  arrange(desc(cv_auc), desc(test_auc)) %>% 
  slice_head(n = 2) %>% filter(w_bioclim == 1)

## HYPERPARAMETER TUNING OF BEST ENVIRO-SOCIO MODEL 

# Filter best_mods to keep relevant columns
best_mods <- best_mods[, c("model", "formula", "formula_id")]

# Initialize an empty list to store tuning results for all datasets
all_tuning_results <- list()

# Loop through each dataset in fit_list3
for (dataset_name in names(fit_list3)) {
  dataset <- fit_list3[[dataset_name]]
  train <- dataset$train
  test <- dataset$test
  formulas_df <- dataset$formulas
  
  # Select the best formulas for this dataset from best_mods
  top_formulas <- best_mods[best_mods$model == dataset_name, ]
  
  # Skip dataset if there are no formulas
  if (nrow(top_formulas) == 0) next 
  
  # Store results for this dataset
  dataset_tuning_results <- list()
  
  for (i in 1:nrow(top_formulas)) {
    best_form <- top_formulas[i, ]
    best_formula <- as.formula(paste(best_form$formula))
    formula_vars <- all.vars(best_formula)
    
    # Subset train and test data to only include relevant variables
    train_subset <- train[, formula_vars, drop = FALSE]
    test_subset <- test[, formula_vars, drop = FALSE]
    
    # Create model matrices
    model_data_train <- model.matrix(best_formula, data = train_subset)
    model_data_test <- model.matrix(best_formula, data = test_subset)
    
    # Separate target variable
    target_train <- train_subset$outcome
    target_test <- test_subset$outcome
    
    # Prepare predictor matrices (remove intercept column)
    predictors_train <- model_data_train[, -1, drop = FALSE]
    predictors_test <- model_data_test[, -1, drop = FALSE]
    
    # Define grid search for tuning parameters
    param_grid <- expand.grid(
      eta = c(0.01, 0.1, 0.3),  # Learning rate
      max_depth = c(3, 6, 9),   # Tree depth
      min_child_weight = c(1, 5, 10),  # Minimum child weight
      subsample = c(0.6, 0.8, 1.0),  # Subsampling fraction
      colsample_bytree = c(0.6, 0.8, 1.0)  # Feature subsampling
    )
    
    # Store tuning results for this formula
    tuning_results <- data.frame(
      eta = numeric(), max_depth = numeric(), min_child_weight = numeric(),
      subsample = numeric(), colsample_bytree = numeric(),
      rmse = numeric(), auc = numeric()
    )
    
    # Loop over each combination of hyperparameters
    for (j in 1:nrow(param_grid)) {
      params <- param_grid[j, ]
      
      # Train XGBoost model with current hyperparameters
      xgb_model <- xgboost(
        data = as.matrix(predictors_train), label = target_train,
        nrounds = 100, eta = params$eta, max_depth = params$max_depth,
        min_child_weight = params$min_child_weight, subsample = params$subsample,
        colsample_bytree = params$colsample_bytree, objective = "binary:logistic",
        verbose = 0
      )
      
      # Make predictions on the test set
      preds_prob <- predict(xgb_model, as.matrix(predictors_test))
      preds_class <- ifelse(preds_prob > 0.5, 1, 0)
      
      # Compute RMSE
      rmse <- sqrt(mean((preds_prob - target_test)^2))
      
      # Compute AUC using pROC
      roc_curve <- roc(target_test, preds_prob)
      auc_value <- auc(roc_curve)
      
      # Store results
      tuning_results <- rbind(tuning_results, data.frame(
        eta = params$eta, max_depth = params$max_depth,
        min_child_weight = params$min_child_weight, subsample = params$subsample,
        colsample_bytree = params$colsample_bytree,
        rmse = rmse, auc = auc_value
      ))
    }
    
    # Select the best model based on highest AUC
    best_params <- tuning_results[which.max(tuning_results$auc), ]
    
    # Store results for this formula only if best_params is not NULL
    if (!is.null(best_params) && nrow(best_params) > 0) {
      best_params$formula_id <- best_form$formula_id  # Add formula_id to best_params
      dataset_tuning_results[[best_form$formula_id]] <- best_params
    }
  } 
  
  # Store results for this dataset only if there are tuning results
  if (length(dataset_tuning_results) > 0) {
    all_tuning_results[[dataset_name]] <- dataset_tuning_results
  }
}

# Convert to a data frame 
tuning_results_df <- all_tuning_results %>%
  map(~ map_dfr(.x, as.data.frame, .id = "formula_id")) %>%  # Convert inner lists to dataframes
  bind_rows(.id = "model")

# Merge back to formulas 
tuning_results_df$formula_id <- as.numeric(tuning_results_df$formula_id)
tuning_results_df <- left_join(tuning_results_df, best_mods)

# Save my tuning results 
write_csv(tuning_results_df, "tuning_results.csv")

