## Make Predictions for Full Time-period 

# Load packages 
library(readr)
library(dplyr)
library(caret)
library(xgboost)
library(ggplot2)

## LOAD AND PROCESS DATA 

# Load dataset and preprocess
fitting_dataset_flies <- read_rds("fitting_dataset_flies.rds")
fitting_dataset_disease <- read_rds("fitting_dataset_disease.rds")

# Average to NUTS2 
fitting_dataset_flies <- fitting_dataset_flies %>%
  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()

fitting_dataset_disease <- fitting_dataset_disease %>%
  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()

## FIT BEST FIT FLY MODELS 

# Fit the fly models first 
# Split for training and testing 
split <- rsample::initial_split(fitting_dataset_flies, prop = 0.75)  
train <- rsample::training(split)
test  <- rsample::testing(split)

# Define a function to train XGBoost models for each outcome
train_xgb_model <- function(features, target, params, data) {
  # Convert the target column to a factor for binary classification
  data[[target]] <- as.factor(data[[target]])
  
  # Prepare training data
  train_data <- data[, c(features, target)]
  
  # Define the tuning grid (including nrounds)
  tune_grid <- expand.grid(
    nrounds = 100,  # You can change this value
    eta = params$eta,
    max_depth = params$max_depth,
    min_child_weight = params$min_child_weight,
    subsample = params$subsample,
    colsample_bytree = params$colsample_bytree,
    gamma = 0  # You can also tune this if necessary
  )
  
  # Train the model using caret
  fit_control <- trainControl(method = "cv", number = 5, search = "grid")
  
  model <- train(
    as.formula(paste(target, "~ .")), data = train_data,
    method = "xgbTree",
    trControl = fit_control,
    tuneGrid = tune_grid,
    objective = params$objective,
    eval_metric = params$eval_metric
  )
  
  return(model)
}

# Define the feature sets and parameters for each species
models_list <- list()

# Define models for each species and their respective parameters (these are saved in "best_mods.txt")
# For the flies 
species_params <- list(
  "neglectus" = list(features = c("arope", "Shrub", "hurs_coldest", "BIO04", "Forest", "population"), target = "neglectus", params = list(eta = 0.3, max_depth = 6, min_child_weight = 1, subsample = 0.6, colsample_bytree = 0.6, objective = "binary:logistic", eval_metric = "auc")),
  "perfiliewi" = list(features = c("arope", "Shrub", "BIO02", "BIO11", "elevation", "Forest"), target = "perfiliewi", params = list(eta = 0.1, max_depth = 3, min_child_weight = 1, subsample = 0.6, colsample_bytree = 0.6, objective = "binary:logistic", eval_metric = "auc")),
  "ariasi" = list(features = c("arope", "BIO06", "elevation", "PermanentWater", "BIO02", "Forest"), target = "ariasi", params = list(eta = 0.3, max_depth = 3, min_child_weight = 1, subsample = 0.6, colsample_bytree = 0.6, objective = "binary:logistic", eval_metric = "auc")),
  "perniciosus" = list(features = c("BIO06", "arope", "Forest", "elevation", "Grass", "population"), target = "perniciosus", params = list(eta = 0.3, max_depth = 9, min_child_weight = 5, subsample = 1, colsample_bytree = 0.8, objective = "binary:logistic", eval_metric = "auc")),
  "tobbi" = list(features = c("BIO02", "Shrub", "Grass", "BuiltUp", "BIO04", "Crops", "Forest"), target = "tobbi", params = list(eta = 0.3, max_depth = 3, min_child_weight = 1, subsample = 0.6, colsample_bytree = 0.6, objective = "binary:logistic", eval_metric = "auc"))
)

# Train each model
for (species in names(species_params)) {
  models_list[[species]] <- train_xgb_model(species_params[[species]]$features, species_params[[species]]$target, species_params[[species]]$params, train)
}

# Double check the perforance 
# Create a performance testing list to save the results 
performnance_list <- list()

# Initialize the final_predictions dataframe with the same number of rows as the test set
performance <- data.frame(matrix(ncol = 0, nrow = nrow(test)))

# Make predictions and store them in performance_list
for (species in names(models_list)) {
  model <- models_list[[species]]
  
  # Make predictions on the test set (probabilities for the positive class)
  y_pred <- predict(model, test, type = "prob")[, 2]  # 2nd column for positive class probabilities
  
  # Get the actual target values from the test set
  y_true <- test[[species_params[[species]]$target]]
  
  # Store predictions and actual values in the performance_list
  performance_list[[species]] <- list(predictions = y_pred, actual = y_true)
  
  # Optionally, print out the AUC score
  auc_score <- pROC::auc(y_true, y_pred)
  print(paste(species, "AUC Score:", auc_score))
  
  # Calculate RMSE
  rmse <- sqrt(mean((y_true - y_pred)^2))
  print(paste(species, "RMSE:", rmse))
  
  # Add actual values for each disease as a separate column in the performance dataframe
  performance[[paste("actual_", species, sep = "")]] <- y_true
  
  # Add model predictions for each disease as a separate column in the performance dataframe
  performance[[paste(species, "_predictions", sep = "")]] <- y_pred
}

## MAKE FLY PREDICTIONS 

# Load my predictions data (features data for 1970-2024)
predicting_dataset <- read_rds("predicting_dataset.rds")

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

# Remove NAs from predicting_dataset
predicting_dataset <- na.omit(predicting_dataset)

# Save and remove identifying data
itfy <- predicting_dataset %>% select(year, NUTS_ID, CNTR_CODE)
predicting_dataset <- predicting_dataset %>% ungroup %>% select(-year, -NUTS_ID, -CNTR_CODE)

# Initialize the predictions dataframe with the correct number of rows based on predicting_dataset
predictions_df <- data.frame(matrix(ncol = length(models_list), nrow = nrow(predictions_df)))

# Set the column names to match the model names (if the models in `models_list` have names)
colnames(predictions_df) <- names(models_list)

# Loop through each model to generate predictions
for (i in 1:length(models_list)) {
  # Extract the model
  model <- models_list[[i]]
  
  # Generate predictions for the current model (adjust the feature set if necessary)
  predictions <- predict(model, predicting_dataset, type = "prob")[, 2]
  
  # Assign the predictions to the column corresponding to the model name in the list
  predictions_df[[names(models_list)[i]]] <- predictions
}

# Merge back to my identify data
predictions_df <- cbind(predictions_df, itfy)

# Save results 
write_csv(predictions_df, "prediction_results_flies.csv")

## FIT BEST FIT DISEASE MODELS 

# Fit the disease models second, with and without the fly predictions as features  
# Join the fly predictions to the fitting dataset 
fitting_dataset_disease <- left_join(fitting_dataset_disease, predictions_df)

# Split for training and testing 
split <- rsample::initial_split(fitting_dataset_disease, prop = 0.75)  
train <- rsample::training(split)
test  <- rsample::testing(split)

# Define the feature sets and parameters for each model
models_list <- list()
# You can use your function you made above to fit the models "train_xgb_model"
# You just need to define your parameters 

# Define models for each disease and their respective parameters (these are saved in "best_mods.txt")
# For the flies 
disease_params <- list(
  "cl" = list(features = c("arope", "BIO01", "elevation", "Shrub", "hurs_driest"), target = "cl", params = list(eta = 0.1, max_depth = 9, min_child_weight = 1, subsample = 0.6, colsample_bytree = 0.6, objective = "binary:logistic", eval_metric = "auc")),
  "vl" = list(features = c("arope", "Shrub", "elevation", "BIO10", "hurs_driest", "BIO04"), target = "vl", params = list(eta = 0.1, max_depth = 9, min_child_weight = 5, subsample = 1, colsample_bytree = 1, objective = "binary:logistic", eval_metric = "auc")),
  "cl_fly" = list(features = c("arope", "BIO01", "elevation", "Shrub", "hurs_driest", "neglectus", "perfiliewi", "ariasi", "perniciosus", "tobbi"), target = "cl", params = list(eta = 0.1, max_depth = 9, min_child_weight = 1, subsample = 0.6, colsample_bytree = 0.6, objective = "binary:logistic", eval_metric = "auc")),
  "vl_fly" = list(features = c("arope", "Shrub", "elevation", "BIO10", "hurs_driest", "BIO04", "neglectus", "perfiliewi", "ariasi", "perniciosus", "tobbi"), target = "vl", params = list(eta = 0.1, max_depth = 9, min_child_weight = 5, subsample = 1, colsample_bytree = 1, objective = "binary:logistic", eval_metric = "auc"))
)

# Train each model
for (diseases in names(disease_params)) {
  models_list[[diseases]] <- train_xgb_model(disease_params[[diseases]]$features, disease_params[[diseases]]$target, disease_params[[diseases]]$params, train)
}

# Double check the performance 
# Create a performance testing list to save the results 
performance_list <- list()

# Initialize the final_predictions dataframe with the same number of rows as the test set
performance <- data.frame(matrix(ncol = 0, nrow = nrow(test)))

# Make predictions and store them in performance_list
for (diseases in names(models_list)) {
  model <- models_list[[diseases]]
  
  # Make predictions on the test set (probabilities for the positive class)
  y_pred <- predict(model, test, type = "prob")[, 2]  # 2nd column for positive class probabilities
  
  # Get the actual target values from the test set
  y_true <- test[[disease_params[[diseases]]$target]]
  
  # Store predictions and actual values in the performance_list
  performance_list[[diseases]] <- list(predictions = y_pred, actual = y_true)
  
  # Optionally, print out the AUC score
  auc_score <- pROC::auc(y_true, y_pred)
  print(paste(diseases, "AUC Score:", auc_score))
  
  # Calculate RMSE
  rmse <- sqrt(mean((y_true - y_pred)^2))
  print(paste(diseases, "RMSE:", rmse))
  
  # Add actual values for each disease as a separate column in the performance dataframe
  performance[[paste("actual_", diseases, sep = "")]] <- y_true
  
  # Add model predictions for each disease as a separate column in the performance dataframe
  performance[[paste(diseases, "_predictions", sep = "")]] <- y_pred
}

## MAKE DISEASE PREDICTIONS 

# Merge the fly predictions to the predicting dataset 
predicting_dataset <- left_join(predicting_dataset, prediction_results_flies)

# Save and remove identifying data
itfy <- predicting_dataset %>% select(year, NUTS_ID, CNTR_CODE)
predictions_df <- predicting_dataset %>% ungroup %>% select(-year, -NUTS_ID, -CNTR_CODE)

# Initialize the predictions dataframe with the correct number of rows based on predicting_dataset
predictions_df <- data.frame(matrix(ncol = length(models_list), nrow = nrow(predictions_df)))

# Set the column names to match the model names (if the models in `models_list` have names)
colnames(predictions_df) <- names(models_list)

# Loop through each model to generate predictions
for (i in 1:length(models_list)) {
  # Extract the model
  model <- models_list[[i]]
  
  # Generate predictions for the current model (adjust the feature set if necessary)
  predictions <- predict(model, predicting_dataset, type = "prob")[, 2]
  
  # Assign the predictions to the column corresponding to the model name in the list
  predictions_df[[names(models_list)[i]]] <- predictions
}

# Merge back to my identify data
predictions_df <- cbind(predictions_df, itfy)

# Save results 
write_csv(predictions_df, "prediction_results_disease.csv")







# ````````````````````````````````````````````````````````````````````````````````````````

## OPTIONAL: PLOT PARTIAL DEPENDENCY PLOTS  

# Define the function to generate PDF plots for each feature in each model
plot_feature_predictions <- function(models_list, data, species_params, output_pdf_path) { # change to disease_params as needed
  # Create an empty list to store the plots
  plot_list <- list()
  
  # Loop through each species/model in models_list
  for (species in names(models_list)) {
    model <- models_list[[species]]
    target_column <- species_params[[species]]$target # change to disease_params as needed
    feature_columns <- species_params[[species]]$features # change to disease_params as needed
    
    # Loop through each feature in the current model
    for (feature in feature_columns) {
      # Generate the predictions for the test set using the model
      y_pred <- predict(model, data, type = "prob")[, 2]  # Get probabilities for the positive class (2nd column)
      
      # Create a data frame to hold the feature and predictions
      plot_data <- data.frame(
        Feature = data[[feature]], 
        Prediction = y_pred,
        Species = species
      )
      
      # Create a plot for the current feature vs the predictions
      p <- ggplot(plot_data, aes(x = Feature, y = Prediction)) +
        geom_point(alpha = 0.5) +
        geom_smooth(method = "loess", col = "blue") +
        labs(title = paste(species, "-", feature),
             x = feature,
             y = "Predicted Probability") +
        theme_minimal()
      
      # Add the plot to the plot list
      plot_list[[paste(species, feature, sep = "_")]] <- p
    }
  }
  
  # Create a PDF to save the plots
  pdf(output_pdf_path, height = 8, width = 6)
  
  # Print each plot to the PDF (using facet_wrap for faceting)
  for (p in plot_list) {
    print(p)
  }
  
  # Close the PDF device
  dev.off()
}

output_pdf_path <- "feature_predictions_plots_flies.pdf"
plot_feature_predictions(models_list, train, species_params, output_pdf_path) # change to disease_params as needed
 

