# R scripts for short course "Species distribution models with Bayesian statistics in R"
# A. Marcia Barbosa (https://modtools.wordpress.com)


# NOTE: most commands on this and the next script can take a LONG TIME to run, so start soon and prepare to be patient!!!


# SETUP ####

library(terra)  # mapping
library(collinear)  # multicollinearity reduction
library(embarcadero)  # BART variable selection
library(nonlinvarsel)  # alternative BART variable selection
library(modEvA)  # model evaluation

setwd(dirname(rstudioapi::getActiveDocumentContext()$path))


# IMPORT DATA AND MODEL ####

my_species <- "Loxodonta africana"  # or another species for which you ran the previous scripts

# example marine species:
# my_species <- "Phocoena sinus"

dat <- read.csv(paste0("../outputs/", my_species, "/occ_data/presences_gridded.csv"))
nrow(dat)

preds <- read.csv(paste0("../outputs/", my_species, "/predictions/BART_preds.csv"))
nrow(preds)

head(preds)
head(dat)

dat <- merge(dat, preds, by = "cells")
nrow(dat)
head(dat)

dat_sv <- terra::vect(dat, geom = c("x", "y"), keepgeom = TRUE, crs = "EPSG:4326")
plot(dat_sv, "presence", col = c("pink", "darkblue"), cex = 0.5)
plot(dat_sv, "BART_P", cex = 0.5)

mod_full <- readRDS(paste0("../outputs/", my_species, "/models/mod_full.rds"))
summary(mod_full)

names(dat)
spc_col <- "presence"  # name of the species presence/absence column IN THE EXAMPLE DATASET (change if appropriate)
var_cols <- 5:23  # variables are in these column positions IN THE EXAMPLE DATASET - change as appropriate, ask me for help if needed!
# var_cols <- c(5:9)  # in the marine example
names(dat)[var_cols]  # confirm it's OK!

constants <- which(sapply(dat[ , var_cols], function(x) length(unique(x)) <= 1))
if (length(constants) > 0) var_cols <- var_cols[-constants]


# VARIABLE IMPORTANCE ####

par(mar = c(5, 4, 2, 1))

# get the importance of each variable in the full model:
embarcadero::varimp(mod_full, plots = TRUE)


# get a diagnostic plot of variable importance for these data and variables:
#(run the following only if you have time at the end, as it is time-consuming, results are just visual and they're not used further in the script)

# varimp_diag <- embarcadero::varimp.diag(x.data = dat[ , var_cols], y.data = dat[ , spc_col])
# plot(varimp_diag)


# EXCLUDE HIGHLY CORRELATED VARIABLES ####

# this exclusion isn't mandatory, but your variable selection may be less stable without it, as the effects of correlated variables are difficult to distinguish

# use all variables:
var_sel <- names(dat)[var_cols]  # for the marine example species, this is advised INSTEAD of the next command, as it has already very few variables

# or select non highly correlated variables (advisable for the terrestrial species example, with more variables):
var_sel <- collinear::collinear(df = dat, response = spc_col, predictors = names(dat)[var_cols], max_cor = 0.7, max_vif = 10)

# note the 'response' argument is only relevant when there are categorical variables (see ?collinear)

var_sel


# VARIABLE SELECTION ####

# several options, no unique solution


# with 'embarcadero' ####

# not meaningful if 'var_sel' has 3 variables or less!

set.seed(456)
varsel_emb <- embarcadero::variable.step(x.data = dat[ , var_sel], y.data = dat[ , spc_col])  # takes time!

varsel_emb

# save selected variables to disk:
write.csv(varsel_emb, paste0("../outputs/", my_species, "/models/varsel_emb.csv"), row.names = FALSE)


# with 'nonlinvarsel' ####

names(dat)
complete_rows <- complete.cases(dat[ , c(var_sel, "BART_P")])  # required below

# forward selection:
varsel_fwd <- nonlinvarsel::vsf(dat[complete_rows, var_sel], dat[complete_rows, "BART_P"])

# backward selection:
varsel_bwd <- nonlinvarsel::vsb(dat[complete_rows, var_sel], dat[complete_rows, "BART_P"])

# all subsets selection:
# very memory-intensive if there are several variables, may crash!!! you can try it with the MARINE EXAMPLE DATA
# varsel_all <- nonlinvarsel::vsa(dat[complete_rows, var_sel], dat[complete_rows, "BART_P"])

print(varsel_fwd)
print(varsel_bwd)
# print(varsel_all)

par(mar = c(5, 4, 2, 1))
plot(varsel_fwd, cex.axis = 0.6)
abline(h = 0.95, lty = 2, col = "darkgrey")
plot(varsel_bwd, cex.axis = 0.6)
abline(h = 0.95, lty = 2, col = "darkgrey")
# plot(varsel_all, cex.axis = 0.6)
# abline(h = 0.95, lty = 2, col = "darkgrey")


R2_thresh <- 0.95  # minimum cumulative R-squared for variable to be selected

# forward minimal subset:
varsel_fwd_ind <- min(which(sort(varsel_fwd$R2) > R2_thresh))
varsel_fwd_ind
varsel_fwd_names <- names(print(varsel_fwd))[1:varsel_fwd_ind]
varsel_fwd_names

# backward minimal subset:
varsel_bwd_ind <- min(which(sort(varsel_bwd$R2) > R2_thresh))
varsel_bwd_ind
varsel_bwd_names <- names(print(varsel_bwd))[1:varsel_bwd_ind]
varsel_bwd_names

# all-subsets minimal subset:
# varsel_all_ind <- min(which(sort(varsel_all$R2) > R2_thresh))
# varsel_all_ind
# varsel_all_names <- print(varsel_all)[[varsel_fwd_ind]]
# varsel_all_names


# save selected variables to disk:

write.csv(varsel_fwd_names, paste0("../outputs/", my_species, "/models/varsel_fwd.csv"), row.names = FALSE)
write.csv(varsel_bwd_names, paste0("../outputs/", my_species, "/models/varsel_bwd.csv"), row.names = FALSE)
# write.csv(varsel_all_names, paste0("../outputs/", my_species, "/models/varsel_all.csv"), row.names = FALSE)


# build models with the selected variables:

# [note that functions like varimp() and predict() may fail if there aren't enough variables, in which case ignore variable selection and use all variables instead]

mod_varsel_emb <- dbarts::bart(x.train = dat[ , varsel_emb], y.train = dat[ , spc_col], keeptrees = TRUE, seed = 654)

summary(mod_varsel_emb)

embarcadero::varimp(mod_varsel_emb, plots = TRUE)


mod_varsel_fwd <- dbarts::bart(x.train = dat[ , varsel_fwd_names], y.train = dat[ , spc_col], keeptrees = TRUE, seed = 654)

summary(mod_varsel_fwd)

embarcadero::varimp(mod_varsel_fwd, plots = TRUE)


mod_varsel_bwd <- dbarts::bart(x.train = dat[ , varsel_bwd_names], y.train = dat[ , spc_col], keeptrees = TRUE, seed = 654)

summary(mod_varsel_bwd)

embarcadero::varimp(mod_varsel_bwd, plots = TRUE)


# mod_varsel_all <- dbarts::bart(x.train = dat[ , varsel_all_names], y.train = dat[ , spc_col], keeptrees = TRUE, seed = 654)

# summary(mod_varsel_all)

# embarcadero::varimp(mod_varsel_all, plots = TRUE)


# save models to disk:

invisible(mod_varsel_emb$fit$state)  # see "Saving" section under ?bart Details
invisible(mod_varsel_fwd$fit$state)
invisible(mod_varsel_bwd$fit$state)
# invisible(mod_varsel_all$fit$state)

saveRDS(mod_varsel_emb, paste0("../outputs/", my_species, "/models/mod_varsel_emb.rds"))
saveRDS(mod_varsel_fwd, paste0("../outputs/", my_species, "/models/mod_varsel_fwd.rds"))
saveRDS(mod_varsel_bwd, paste0("../outputs/", my_species, "/models/mod_varsel_bwd.rds"))
# saveRDS(mod_varsel_all, paste0("../outputs/", my_species, "/models/mod_varsel_all.rds"))


# COMPARE PREDICTIONS FROM DIFFERENT SELECTED VARIABLES ####

# model performance:

par(mfrow = c(2, 2))
modEvA::AUC(model = mod_full, simplif = TRUE, main = "mod_full")
modEvA::AUC(model = mod_varsel_emb, simplif = TRUE, main = "mod_varsel_emb")
modEvA::AUC(model = mod_varsel_fwd, simplif = TRUE, main = "mod_varsel_fwd")
modEvA::AUC(model = mod_varsel_bwd, simplif = TRUE, main = "mod_varsel_bwd")
# modEvA::AUC(model = mod_varsel_all, simplif = TRUE)

modEvA::MillerCalib(model = mod_full, main = "mod_full")
modEvA::MillerCalib(model = mod_varsel_emb, main = "mod_varsel_emb")
modEvA::MillerCalib(model = mod_varsel_fwd, main = "mod_varsel_fwd")
modEvA::MillerCalib(model = mod_varsel_bwd, main = "mod_varsel_bwd")
# modEvA::MillerCalib(model = mod_varsel_all)


# check correlations between predictions:

source("https://raw.githubusercontent.com/AMBarbosa/unpackaged/master/predict_bart_df")

dat$BART_varsel_emb_P <- predict_bart_df(mod_varsel_emb, dat)
dat$BART_varsel_fwd_P <- predict_bart_df(mod_varsel_fwd, dat)
dat$BART_varsel_bwd_P <- predict_bart_df(mod_varsel_bwd, dat)
# dat$BART_varsel_all_P <- predict_bart_df(mod_varsel_all, dat)

pred_corrs <- cor(dat[ , c("BART_P", "BART_varsel_emb_P", "BART_varsel_fwd_P", "BART_varsel_bwd_P")])  # , "BART_varsel_all_P"

pred_corrs
pred_corrs[ , 1, drop = FALSE]


# map the different predictions:

dat_sv <- terra::vect(dat, geom = c("x", "y"), keepgeom = TRUE, crs = "EPSG:4326")

columns_to_map <- grep("BART_P$|BART_varsel", names(dat_sv))
par(oma = c(0, 0, 1.5, 0))
plot(dat_sv, columns_to_map, type = "continuous", range = c(0, 1), cex = 0.4, mar = c(1, 1, 1.5, 3))
title(my_species, outer = TRUE)


#### ASSIGNMENT 7: if you're using different data than the example, post the figure you obtained with the above commands (if you have time, convert to favourability first)