## Code to run behavior-specific and site selection timber rattlesnake models
## to use when fitting 'bernoulli' brms models to snake behavior data


# Libraries ---------------------------------------------------------------
library(brms)
library(dplyr)
library(bayestestR)
library(cowplot)
library(ggplot2)
library(lubridate)
library(lme4)
library(rstanarm)
library(bayesplot)
library(loo)
library(ggThemeAssist)
ggplot2::theme_set(theme_cowplot())

# Import data ------------------------------------------------------------------

setwd("~/DRYAD")

#Read in site use dataset, set sex and year as factors
dat_use <- read.csv("dat_use_5m.csv")
dat_use$Male <- as.factor(dat_use$Male)
dat_use$Adult <- as.factor(dat_use$Adult)
dat_use$Year <- as.Date(dat_use$Date, format = "%Y-%m-%d") %>% lubridate::year(.)

#Scale site use dataset
dat_use.s <- scale(dat_use[,c(7:21)])
colnames(dat_use.s) <- paste0(colnames(dat_use.s), ".s")
dat.use5 <- cbind(dat_use[,c(1:6)], dat_use.s, dat_use[,c(7:21)])

#Read in four different scale datasets for behaviors, set sex and year as factors
dat5 <- read.csv("dat_5m.csv")
dat5$Male <- as.factor(dat5$Male)
dat5$Adult <- as.factor(dat5$Adult)
dat5$Year <- as.Date(dat5$Date, format = "%m/%d/%Y") %>% lubridate::year(.)
dat25 <- read.csv("dat_25m.csv")
dat25$Male <- as.factor(dat25$Male)
dat25$Adult <- as.factor(dat25$Adult)
dat25$Year <- as.Date(dat25$Date, format = "%m/%d/%Y") %>% lubridate::year(.)
dat55 <- read.csv("dat_55m.csv")
dat55$Male <- as.factor(dat55$Male)
dat55$Adult <- as.factor(dat55$Adult)
dat55$Year <- as.Date(dat55$Date, format = "%m/%d/%Y") %>% lubridate::year(.)
dat105 <- read.csv("dat_105m.csv")
dat105$Male <- as.factor(dat105$Male)
dat105$Adult <- as.factor(dat105$Adult)
dat105$Year <- as.Date(dat105$Date, format = "%m/%d/%Y") %>% lubridate::year(.)

#Create datasets at each scale
Behav5 <- data.frame(dat5)
Behav25 <- data.frame(dat25)
Behav55 <- data.frame(dat55)
Behav105 <- data.frame(dat105)

#Create multiscale ecdysis dataset from 105m dataset
datEM <- data.frame(dat105)
datEM$meanT <- dat5$meanT
datEM$stnd_age <- dat25$stnd_age
datEM$mid <- dat5$mid
datEM$und <- dat5$und
datEM$ndms1 <- dat55$ndms1
datEM$ndms3 <- dat5$ndms3

#Create multiscale gestation dataset from 105m dataset
datGM <- data.frame(dat105)
datGM$strm_dist <- dat25$strm_dist
datGM$rf_tba <- dat55$rf_tba
datGM$ove <- dat25$ove
datGM$mid <- dat25$mid
datGM$ndms2 <- dat5$ndms2

#Create multiscale digestion dataset from 105m dataset
datDM <- data.frame(dat105)
datDM$meanT <- dat5$meanT
datDM$slope <- dat55$slope

# > Site Use Model --------------------------------------------------------

mod.use <- brm(Used ~
                   meanT.s + sol_rad.s + stnd_age.s + slope.s + strm_dist.s +
                   rf_tba.s + rf_tde.s + chm.s + ove.s + mid.s + und.s + ndms1.s + ndms2.s + ndms3.s +
                   (1|Snake),
                   data = dat.use5,
                   family = 'bernoulli',
                   chains = 4,
                   cores = 4,
                   iter = 15000,
                   warmup = 3000,
                   control = list(adapt_delta = 0.975,
                                  max_treedepth = 20),
                   save_all_pars = TRUE,
                   file = 'mod.use'
)


#Describe Posteror
describe_posterior(mod.use)

# Foraging ----------------------------------------------------------------
#Must change dataset used depending on which scale is desired
dat.behavF <- Behav25

loop.dat <- dat.behavF %>%
  group_by(Snake, Year, Forage) %>%  ## Need to specify behavior (here, Forage)
  summarise(count = n()) %>% as.data.frame()

select.list <- vector('list', sum(loop.dat[,3]))
obs <- 0
for(i in which(loop.dat[,3] == 1)) {
  set.seed(12300)
  obs <- obs + 1
  snake <- loop.dat[i,1]
  year <- loop.dat[i,2]

  samp.dat <- dat.behavF[dat.behavF$Snake == snake &
                          dat.behavF$Year == year &
                           dat.behavF$Forage == 0,] ## Need to specify behavior == 0 (here, Forage)

  ## Take a sample that has 4x zeros as observations
  samp.size <- ifelse(4 * loop.dat[i,4] > nrow(samp.dat),
                      nrow(samp.dat),
                      4 * loop.dat[i,4])
  select.df <- samp.dat[sort(sample(nrow(samp.dat),
                                    samp.size,
                                    replace = F)),]

  obs.df <- dat.behavF[dat.behavF$Snake == snake &
                          dat.behavF$Year == year &
                          dat.behavF$Forage == 1,] ## Need to specify behavior == 1 (here, Forage)

  select.list[[obs]] <- rbind(obs.df, select.df)
}

select.forage <- plyr::ldply(select.list)  ## Change the name of the final data set
dat.s <- scale(select.forage[,c(10:24)])
colnames(dat.s) <- paste0(colnames(dat.s), ".s")
select.forage <- cbind(select.forage[,c(1:9)], dat.s, select.forage[,c(10:24)])
select.forage[is.na(select.forage$meanT.s),]


# > Forage Model ---------------------------------------------------------------

#Full model, must change scale for model name and file name
mod.for_25 <- brm(Forage ~
                    meanT.s + sol_rad.s + stnd_age.s + slope.s + strm_dist.s + beers.s +
                    rf_tba.s + rf_tde.s + chm.s + ove.s + mid.s + und.s + ndms1.s + ndms2.s + ndms3.s +
                    (1 | Snake),
                  data = select.forage,
                  family = 'bernoulli',
                  chains = 4,
                  cores = 4,
                  iter = 15000,
                  warmup = 3000,
                  control = list(adapt_delta = 0.975,
                                 max_treedepth = 20),
                  save_all_pars = TRUE,
                  file = 'mod_for25.rds'
)

#Reduced model, must change scale for model name and file name
rmod.for_25 <- brm(Forage ~
                     meanT.s + sol_rad.s + slope.s + strm_dist.s + beers.s + chm.s + ndms1.s +
                     (1|Snake),
                   data = select.forage,
                   family = 'bernoulli',
                   chains = 4,
                   cores = 4,
                   iter = 15000,
                   warmup = 3000,
                   control = list(adapt_delta = 0.975,
                                  max_treedepth = 20),
                   save_all_pars = TRUE,
                   file = 'rmod_for25.rds'
)

#Describe Posteror, change model name for different scales
describe_posterior(mod.for_25)
describe_posterior(rmod.for_25)


# Ecdysis ----------------------------------------------------------------

#Must change dataset used depending on which scale is desired
dat.behavE <- datEM

loop.dat <- dat.behavE %>%
  group_by(Snake, Year, Ecdysis) %>% ## Need to specify behavior (here, Ecdysis)
  summarise(count = n()) %>% as.data.frame()

select.list <- vector('list', sum(loop.dat[,3]))
obs <- 0
for(i in which(loop.dat[,3] == 1)) {
  set.seed(12300)
  obs <- obs + 1
  snake <- loop.dat[i,1]
  year <- loop.dat[i,2]

  samp.dat <- dat.behavE[dat.behavE$Snake == snake &
                           dat.behavE$Year == year &
                           dat.behavE$Ecdysis == 0,] ## Need to specify behavior == 0 (here, Ecdysis)

  ## Take a sample that has 4x zeros as observations
  samp.size <- ifelse(4 * loop.dat[i,4] > nrow(samp.dat),
                      nrow(samp.dat),
                      4 * loop.dat[i,4])
  select.df <- samp.dat[sort(sample(nrow(samp.dat),
                                    samp.size,
                                    replace = F)),]

  obs.df <- dat.behavE[dat.behavE$Snake == snake &
                        dat.behavE$Year == year &
                        dat.behavE$Ecdysis == 1,] ## Need to specify behavior == 1 (here, Ecdysis)

  select.list[[obs]] <- rbind(obs.df, select.df)
}

select.ecdysis <- plyr::ldply(select.list) ## Change the name of the final data set
dat.s <- scale(select.ecdysis[,c(10:24)])
colnames(dat.s) <- paste0(colnames(dat.s), ".s")
select.ecdysis <- cbind(select.ecdysis[,c(1:9)], dat.s, select.ecdysis[,c(10:24)])
select.ecdysis[is.na(select.ecdysis$meanT.s),]

# > Ecdysis Model ---------------------------------------------------------------

#Full model, must change scale for model name and file name
mod.ecd_MS <- brm(Ecdysis ~
                   meanT.s + sol_rad.s + stnd_age.s + slope.s + strm_dist.s + beers.s +
                   rf_tba.s + rf_tde.s + chm.s + ove.s + mid.s + und.s + ndms1.s + ndms2.s + ndms3.s +
                   (1 | Snake),
                 data = select.ecdysis,
                 family = 'bernoulli',
                 chains = 4,
                 cores = 4,
                 iter = 15000,
                 warmup = 3000,
                 control = list(adapt_delta = 0.99,
                                max_treedepth = 20),
                 save_all_pars = TRUE,
                 file = 'fitted/mod.ecdMS'
)

#Reduced model, must change scale for model name and file name
rmod.ecd_MS <- brm(Ecdysis ~
                     meanT.s + sol_rad.s + stnd_age.s + slope.s + strm_dist.s + beers.s +
                     rf_tde.s + chm.s + ove.s + mid.s + und.s + ndms1.s + ndms3.s +
                     (1|Snake),
                   data = select.ecdysis,
                   family = 'bernoulli',
                   chains = 4,
                   cores = 4,
                   iter = 15000,
                   warmup = 3000,
                   control = list(adapt_delta = 0.99,
                                  max_treedepth = 20),
                   save_all_pars = TRUE,
                   file = 'rmod.ecdMS'
)

#Describe Posteror, change model name for different scales
describe_posterior(mod.ecd_MS)
describe_posterior(rmod.ecd_MS)


# Gestation ----------------------------------------------------------------

dat.behavG <- datGM

loop.dat <- dat.behavG %>%
  group_by(Snake,
           # Year, ## No year
           Gestate) %>% ## Need to specify behavior (here, Gestation)
  summarise(count = n()) %>% as.data.frame()

select.list <- vector('list', sum(loop.dat[,3]))
obs <- 0
for(i in which(loop.dat[,2] == 1)) {
  set.seed(12300)
  obs <- obs + 1
  snake <- loop.dat[i,1]
  # year <- loop.dat[i,2]

  samp.dat <- dat.behavG[dat.behavG$Snake == snake &
                           dat.behavG$Year == year &
                           dat.behavG$Gestate == 0,] ## Need to specify behavior == 0 (here, Gestation)

  ## Take a sample that has 4x zeros as observations
  samp.size <- ifelse(4 * loop.dat[i,3] > nrow(samp.dat),
                      nrow(samp.dat),
                      4 * loop.dat[i,3])
  select.df <- samp.dat[sort(sample(nrow(samp.dat),
                                    samp.size,
                                    replace = F)),]

  obs.df <- dat.behavG[dat.behavG$Snake == snake &
                         dat.behavG$Year == year &
                         dat.behavG$Gestate == 1,] ## Need to specify behavior == 1 (here, Gestation)

  select.list[[obs]] <- rbind(obs.df, select.df)
}

select.gestate <- plyr::ldply(select.list) ## Change the name of the final data set
dat.s <- scale(select.gestate[,c(10:24)])
colnames(dat.s) <- paste0(colnames(dat.s), ".s")
select.gestate <- cbind(select.gestate[,c(1:9)], dat.s, select.gestate[,c(10:24)])
select.gestate[is.na(select.gestate$meanT.s),]


# > Gestation Model ---------------------------------------------------------------

mod.ges_MS <- brm(Gestate ~
                   meanT.s + sol_rad.s + stnd_age.s + slope.s + strm_dist.s + beers.s +
                   rf_tba.s + rf_tde.s + chm.s + ove.s + mid.s + und.s + ndms1.s + ndms2.s + ndms3.s +
                   (1 | Snake),
                 data = select.gestate,
                 family = 'bernoulli',
                 chains = 4,
                 cores = 4,
                 iter = 15000,
                 warmup = 3000,
                 control = list(adapt_delta = 0.975,
                                max_treedepth = 20),
                 save_all_pars = TRUE,
                 file = 'fitted/mod.gesMS'
)

rmod.ges_MS <- brm(Gestate ~
                    meanT.s + sol_rad.s + stnd_age.s + slope.s + strm_dist.s + beers.s +
                    rf_tba.s + rf_tde.s + chm.s + ove.s + mid.s + und.s + ndms1.s + ndms2.s +
                    (1 | Snake),
                  data = select.gestate,
                  family = 'bernoulli',
                  chains = 4,
                  cores = 4,
                  iter = 15000,
                  warmup = 3000,
                  control = list(adapt_delta = 0.975,
                                 max_treedepth = 20),
                  save_all_pars = TRUE,
                  file = 'rmod.gesMS'
)

#Describe Posteror, change model name for different scales
describe_posterior(mod.ges_MS)
describe_posterior(rmod.ges_MS)


# Digest ----------------------------------------------------------------

dat.behavD <- datDM

loop.dat <- dat.behavD %>%
  group_by(Snake, Year, Digest) %>% ## Need to specify behavior (here, Digest)
  summarise(count = n()) %>% as.data.frame()

select.list <- vector('list', sum(loop.dat[,3]))
obs <- 0
for(i in which(loop.dat[,3] == 1)) {
  set.seed(12300)
  obs <- obs + 1
  snake <- loop.dat[i,1]
  year <- loop.dat[i,2]

  samp.dat <- dat.behavD[dat.behavD$Snake == snake &
                           dat.behavD$Year == year &
                           dat.behavD$Digest == 0,] ## Need to specify behavior == 0 (here, Digest)

  ## Take a sample that has 4x zeros as observations
  samp.size <- ifelse(4 * loop.dat[i,4] > nrow(samp.dat),
                      nrow(samp.dat),
                      4 * loop.dat[i,4])
  select.df <- samp.dat[sort(sample(nrow(samp.dat),
                                    samp.size,
                                    replace = F)),]

  obs.df <- dat.behavD[dat.behavD$Snake == snake &
                         dat.behavD$Year == year &
                         dat.behavD$Digest == 1,] ## Need to specify behavior == 1 (here, Digest)

  select.list[[obs]] <- rbind(obs.df, select.df)
}

select.digestion <- plyr::ldply(select.list) ## Change the name of the final data set
dat.s <- scale(select.digestion[,c(10:24)])
colnames(dat.s) <- paste0(colnames(dat.s), ".s")
select.digestion <- cbind(select.digestion[,c(1:9)], dat.s, select.digestion[,c(10:24)])
select.digestion[is.na(select.digestion$meanT.s),]



# > Digestion Model ---------------------------------------------------------------

#Full model, must change scale for model name and file name
mod.dig_MS <- brm(Digest ~
                   meanT.s + sol_rad.s + stnd_age.s + slope.s + strm_dist.s + beers.s +
                   rf_tba.s + rf_tde.s + chm.s + ove.s + mid.s + und.s + ndms1.s + ndms2.s + ndms3.s +
                   (1 | Snake),
                 data = select.digestion,
                 family = 'bernoulli',
                 chains = 4,
                 cores = 4,
                 iter = 15000,
                 warmup = 3000,
                 control = list(adapt_delta = 0.975,
                                max_treedepth = 20),
                 save_all_pars = TRUE,
                 file = 'fitted/mod.digMS'
)

#Reduced model, must change scale for model name and file name
rmod.dig_MS <- brm(Digest ~
                   meanT.s + sol_rad.s + slope.s + rf_tba.s + chm.s + ndms3.s +
                    (1 | Snake),
                 data = select.digestion,
                 family = 'bernoulli',
                 chains = 4,
                 cores = 4,
                 iter = 15000,
                 warmup = 3000,
                 control = list(adapt_delta = 0.975,
                                max_treedepth = 20),
                 save_all_pars = TRUE,
                 file = 'rmod.digMS'
)

#Describe Posterior, change model name for different scales
describe_posterior(mod.dig_MS)
describe_posterior(rmod.dig_MS)