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


# SETUP ####

library(embarcadero)
library(raster)


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


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"))
preds <- read.csv(paste0("../outputs/", my_species, "/predictions/BART_preds.csv"))
dat <- merge(dat, preds, by = "cells")

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!

mod_full <- readRDS(paste0("../outputs/", my_species, "/models/mod_full.rds"))
mod_varsel_emb <- readRDS(paste0("../outputs/", my_species, "/models/mod_varsel_emb.rds"))
mod_varsel_fwd <- readRDS(paste0("../outputs/", my_species, "/models/mod_varsel_fwd.rds"))
mod_varsel_bwd <- readRDS(paste0("../outputs/", my_species, "/models/mod_varsel_bwd.rds"))

layers <- terra::rast(paste0("../outputs/", my_species, "/layers_cut/layers_cut.tif"))


# PLOT THE PREDICTED RESPONSE VALUES ####

preds_sorted <- preds[order(preds$BART_P), ]

par(mfrow = c(1, 1))

plot(preds_sorted$BART_P, pch = ".")
points(preds_sorted$BART_P_upper, pch = ".", col = "darkgrey")
points(preds_sorted$BART_P_lower, pch = ".", col = "darkgrey")

plot(preds_sorted$BART_F, pch = ".")
points(preds_sorted$BART_F_upper, pch = ".", col = "darkgrey")
points(preds_sorted$BART_F_lower, pch = ".", col = "darkgrey")


# plot predictions against uncertainty:

plot(dat$BART_P, dat$BART_P_uncert, pch = 20)
plot(dat$BART_F, dat$BART_F_uncert, pch = 20)


# PARTIAL DEPENDENCE PLOTS ####

# choose a model to play with in this and the next sections:
mod <- mod_full  # may be better for the marine example, as the models made with variable selection may have too little variation (too few variables)
mod <- mod_varsel_emb  # or any other BART model you have in your workspace

# get the variables in that model:
vars <- colnames(mod$varcount)
vars


# partial dependence plots with 'dbarts':

partial_1_pdb <- dbarts::pdbart(x.train = dat[ , vars], y.train = dat[ , spc_col], xind = vars[1])
plot(partial_1_pdb)

# choose another variable:
partial_2_pdb <- dbarts::pdbart(x.train = dat[ , vars], y.train = dat[ , spc_col], xind = vars[2])
plot(partial_2_pdb)


# partial dependence plots with 'embarcadero':

partial_1_emb <- embarcadero::partial(mod, x.vars = vars[1], trace = FALSE)  # takes time! and would take even longer with the default trace=TRUE
partial_1_emb

# choose another variable:
partial_2_emb <- embarcadero::partial(mod, x.vars = vars[2], trace = FALSE)
partial_2_emb

# you can get the partial dependence plots for all variables in the model at once, but note this may take a very long time! you can try it with the marine example data, as it has fewer variables
# partials <- partial(mod, trace = FALSE)
# partials


# BIDIMENSIONAL PARTIAL DEPENDENCE PLOTS ####

partial_1and2 <- dbarts::pd2bart(mod, xind = c(vars[1], vars[2]))  # takes time!
par(mfrow = c(1, 1), mar = c(4.5, 4, 2, 1))
plot(partial_1and2)

# (later) try this also with other combinations of variables, especially those with "strange" partial response plots


# SPATIAL PARTIAL DEPENDENCE PLOTS ####

# these are a cartographic version of partial dependence plots, showing where a particular variable most and least favours species presence

layers_stack <- raster::stack(layers)  # 'embarcadero' still requires old 'raster' pkg formats - this may change in future versions

spartial_1 <- embarcadero::spartial(mod, envs = layers_stack, x.vars = vars[1])
plot(spartial_1, main = paste("partial prediction:\n", vars[1]))

spartial_2 <- embarcadero::spartial(mod, envs = layers_stack, x.vars = vars[2])
plot(spartial_2, main = paste("partial prediction:\n", vars[2]))


# generate plots with different variables and combinations, to see how each variable or interaction affects your species
# try interpreting the results in light of your species' ecology
# read each function's help file and try other argument options


# PLOT INDIVIDUAL TREES ####

mod_full$fit$plotTree(sampleNum = 2, treeNum = 2)

mod_full$fit$plotTree(sampleNum = 5, treeNum = 3)


# make a simpler 10-tree, 10-sample model to play with, else the following will take too long to run:

mod_simple <- dbarts::bart(x.train = dat[ , vars], y.train = dat[ , spc_col], keeptrees = TRUE, seed = 10, ntree = 10, ndpost = 10)


# extract all model trees:

trees <- dbarts::extract(mod_simple, "trees")

head(trees)
tail(trees)


# plot all model trees:

sample_seq <- sort(unique(trees$sample))
tree_seq <- unique(trees$tree)
n_trees <- length(sample_seq) * length(tree_seq)
n_trees

for (s in sample_seq) for (t in tree_seq) {
  mod_simple$fit$plotTree(sampleNum = s, treeNum = t)
  legend("topleft", paste("sample", s, "\ntree", t), bty = "n", text.font = 2)
}


# MAP THE BART LEARNING PROCESS ####

source("https://raw.githubusercontent.com/AMBarbosa/embarcadero/master/R/plot_mcmc.R")  # (I introduced a minor fix so maps with NA pixels are not disorganized; this should eventually be fixed in 'embarcadero' too - check for package updates on GitHub)
plot.mcmc(mod, layers_stack, wait = 0.5, iter = 50)  # see the 'plots' pane