# Nest survival model using JAGS 
# Fit a nest level categorical variable (habitat)

# Analysis of Shelly redstart data using a Bayesian model

# Sherry TW, Wilson S, Hunter S, Holmes RT (2015) 
# Impacts of nest predators and weather on reproductive success and 
# population limitation in a long-distance migratory songbird. 
# Journal of Avian Biology 46(6): 559-569. https://doi.org/10.1111/jav.00536

# Data from
# Sherry TW, Wilson S, Hunter S, Holmes RT (2015) 
# Data from: Impacts of nest predators and weather on reproductive success and 
# population limitation in a long-distance migratory songbird. 
# Dryad Digital Repository. https://doi.org/10.5061/dryad.73870

#
# 2019-06-28 CHJS First Edition
#

library("R2jags")  # used for call to JAGS
## Loading required package: rjags
## Loading required package: coda
## Linked to JAGS 4.3.0
## Loaded modules: basemod,bugs
## 
## Attaching package: 'R2jags'
## The following object is masked from 'package:coda':
## 
##     traceplot
library(coda)
library(ggplot2)
## Registered S3 methods overwritten by 'ggplot2':
##   method         from 
##   [.quosures     rlang
##   c.quosures     rlang
##   print.quosures rlang
library(reshape2)

options(width=200)

source(file.path("..","..","jags-nest-survival-fixed-effects.r"))


# The dataframe must contain the following fields with the following names
#
#    FirstFound: day the nest was first found
#    LastPresent: last day that a chick was present in the nest
#    LastChecked: last day the nest was checked
#    Fate: fate of the nest; 0=hatch an
#    Freq: number of nests with this data
#
# In this example, the multiple visits to a nest have been collapsed
# to a single record for each nest.
# In more complex examples, you may have multple records per nest
# as shown in the mallard example.
#

nestdata <- readxl::read_excel(file.path("..","Sherry.xlsx"), 
                              sheet="NestData")
head(nestdata)
## # A tibble: 6 x 11
##   NestId       FirstFound LastPresent LastChecked  Fate  Freq Height BaffleStatus   DBH AgeDay1 Year 
##   <chr>             <dbl>       <dbl>       <dbl> <dbl> <dbl>  <dbl> <chr>        <dbl>   <dbl> <chr>
## 1 /*1983-175*/         11          11          18     1     1   16.9 N             39        -8 1983 
## 2 /*1983-176*/         13          22          23     1     1   10.5 N             10.7     -10 1983 
## 3 /*1983-178*/          7          26          26     0     1    4.8 N              8.8      -4 1983 
## 4 /*1983-180*/         12          14          20     1     1    3.7 N              5.1      -9 1983 
## 5 /*1983-181*/          6          20          25     1     1   10.7 N             15.6      -3 1983 
## 6 /*1983-182*/         11          19          20     1     1   17.8 N             46.2      -8 1983
# as noted in the paper, we remove nests from years with no baffeling
dim(nestdata)
## [1] 537  11
nestdata <- nestdata[ !(nestdata$BaffleStatus=="N" & nestdata$Year %in% c(1983, 1984, 1990)),]
dim(nestdata)
## [1] 414  11
# create factor variable for year because it numeric
nestdata$YearF <- factor(nestdata$Year)


# Unfortunately, JAGS cannot deal with alpha numeric code and 
# so we need to convert the alphanumberic NestID to numeric codes
# by declaring NestId as a factor and extracting the level values
nestdata$NestId.num <- as.numeric(factor(nestdata$NestId))

# We must create a file with every combination of next x day nest was "active"
# being every day from FirstCound to LastChecked-1

nesttime <- plyr::adply(nestdata, 1, function(x){
     nesttime <- expand.grid(NestId.num=x$NestId.num, 
                             Day=x$FirstFound:(x$LastChecked-1),
                             stringsAsFactors=FALSE)
     nesttime
})


# Extract the nest level covariates (including AgeNest1)
# The next level covariates should be indexed using NestId
# If AgeNest1 variable is present then the age of the nest is computed
#
nest.covariates <- NULL

if( !is.null(nest.covariates)){
   nesttime <- merge(nesttime, nest.covariates, by="NestId")
}


# Extract any survey time covariates such as time, time^2, early/late
# weather covariates ect.
# All of these covariates will affect all nests simultaneouls

# none added here

# if there is a AgeDay1 variable, we compute the nest age for each time for each nest
if( !is.null(nesttime$AgeDay1)){
   nesttime$NestAge <- nesttime$AgeDay1 + nesttime$Day -1
}
head(nesttime)
##         NestId FirstFound LastPresent LastChecked Fate Freq Height BaffleStatus DBH AgeDay1 Year YearF NestId.num Day NestAge
## 1 /*1985-434*/          7          27          27    0    1    2.6            N 3.3      -4 1985  1985          1   7       2
## 2 /*1985-434*/          7          27          27    0    1    2.6            N 3.3      -4 1985  1985          1   8       3
## 3 /*1985-434*/          7          27          27    0    1    2.6            N 3.3      -4 1985  1985          1   9       4
## 4 /*1985-434*/          7          27          27    0    1    2.6            N 3.3      -4 1985  1985          1  10       5
## 5 /*1985-434*/          7          27          27    0    1    2.6            N 3.3      -4 1985  1985          1  11       6
## 6 /*1985-434*/          7          27          27    0    1    2.6            N 3.3      -4 1985  1985          1  12       7
# Add any next x day survey covariates to the nesttime data
#

# there is nothing here for this example


# Set up the design matrix for the fixed effects
fe.design <- model.matrix(  ~ DBH+BaffleStatus+YearF, data=nesttime)

head(fe.design)
##   (Intercept) DBH BaffleStatusY YearF1986 YearF1987 YearF1988 YearF1989 YearF1991 YearF1992 YearF1993 YearF1994 YearF1995
## 1           1 3.3             0         0         0         0         0         0         0         0         0         0
## 2           1 3.3             0         0         0         0         0         0         0         0         0         0
## 3           1 3.3             0         0         0         0         0         0         0         0         0         0
## 4           1 3.3             0         0         0         0         0         0         0         0         0         0
## 5           1 3.3             0         0         0         0         0         0         0         0         0         0
## 6           1 3.3             0         0         0         0         0         0         0         0         0         0
# Finally, the actual call to JAGS
fitted.model <- jags.nest.survival.fixed.effects(
         nestdata=nestdata,    # nest data
         nesttime=nesttime,    # daily nest values with nest, time, nest x time covariates
         fe.design=fe.design,  # fixed effects design matrix
         init.seed=12321312)   # initial seed)  
## module glm loaded
## Compiling data graph
##    Declaring variables
##    Resolving undeclared variables
##    Allocating nodes
##    Initializing
##    Reading data back into data table
## Compiling model graph
##    Declaring variables
##    Resolving undeclared variables
##    Allocating nodes
## Graph information:
##    Observed stochastic nodes: 561
##    Unobserved stochastic nodes: 12
##    Total graph size: 82394
## 
## Initializing model
# the results list has lots of other stuff
results <- fitted.model$results



# the nesttime dataframe has the estimated DSR for every combination of NestId.num and Day
# evaluated at the covariates for that nest and time

# This is problematic as DBH will vary among the nests.
# We want the DSR at the mean DBH for each baffle status x year combination
# Consequently, we need to create a design matrix for every baffle status x year and
# the mean DBH and then multiply by the estimated betas 

pred.df <- unique(fitted.model$nesttime[, c("BaffleStatus","YearF")])
pred.df$DBH <- mean(nestdata$DBH)

# Notice that the model must MATCH exactly in the order of terms etc as used to generate
# the fe.design matrix seen earlier
pred.design <- model.matrix( ~ DBH+BaffleStatus+YearF, data=pred.df)

head(fe.design)
##   (Intercept) DBH BaffleStatusY YearF1986 YearF1987 YearF1988 YearF1989 YearF1991 YearF1992 YearF1993 YearF1994 YearF1995
## 1           1 3.3             0         0         0         0         0         0         0         0         0         0
## 2           1 3.3             0         0         0         0         0         0         0         0         0         0
## 3           1 3.3             0         0         0         0         0         0         0         0         0         0
## 4           1 3.3             0         0         0         0         0         0         0         0         0         0
## 5           1 3.3             0         0         0         0         0         0         0         0         0         0
## 6           1 3.3             0         0         0         0         0         0         0         0         0         0
head(pred.design)
##     (Intercept)      DBH BaffleStatusY YearF1986 YearF1987 YearF1988 YearF1989 YearF1991 YearF1992 YearF1993 YearF1994 YearF1995
## 1             1 16.12215             0         0         0         0         0         0         0         0         0         0
## 21            1 16.12215             1         0         0         0         0         0         0         0         0         0
## 43            1 16.12215             0         0         1         0         0         0         0         0         0         0
## 108           1 16.12215             1         0         1         0         0         0         0         0         0         0
## 225           1 16.12215             0         0         0         1         0         0         0         0         0         0
## 297           1 16.12215             1         0         0         1         0         0         0         0         0         0
# extract the posterior sample of the beta as a matrix
# the results$BUGSoutput$sims.matrix is a 2-d object [iterations, variables] with chains stacked
# on top of each other
dim(  results$BUGSoutput$sims.matrix)
## [1] 4500 5236
select <- grepl("^beta", colnames(results$BUGSoutput$sims.matrix))
colnames(results$BUGSoutput$sims.matrix)[select]
##  [1] "beta[1]"  "beta[2]"  "beta[3]"  "beta[4]"  "beta[5]"  "beta[6]"  "beta[7]"  "beta[8]"  "beta[9]"  "beta[10]" "beta[11]" "beta[12]"
beta.matrix <- results$BUGSoutput$sims.matrix[,select]

# create posterior sample of DSR (first on the logit scale) and then back transformed
DSR <- beta.matrix %*% t(pred.design)
DSR <- 1/(1+exp(-DSR))
dim(DSR)
## [1] 4500   20
pred.df$mean <- apply(DSR,2,mean)
pred.df$sd   <- apply(DSR,2,sd)
pred.df$lcl  <- apply(DSR,2, quantile, prob=.025)
pred.df$ucl  <- apply(DSR,2,  quantile, prob=.975)

# now we are ready to plot
pred.df$Year <- as.numeric(as.character(pred.df$YearF))
# in this case, we fit a S~Habitat model, so all rows of nesttime have the same estimated DSR

ggplot(data=pred.df, aes(x=Year, y=mean, color=BaffleStatus))+
  ggtitle("Estimated DSR by Baffle Status and Year",
          subtitle="Evaluated at the mean DBH")+
  geom_line( position=position_dodge(w=0.2))+
  geom_errorbar(aes(ymin=lcl, ymax=ucl), width=0.1, position=position_dodge(w=0.2))+
  ylab("DSR (95% ci)")+
  theme(legend.position=c(0,0),legend.justification=c(0,0))