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

#
# 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 input dataframe must contain the following fields with the following names
#
#    NestID: id code of the nest (alpha numeric)
#    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 = success; 1=fail
#    AgheDay1 = age of the nest on day 1 (if you are fitting age of nest models) 
#
# You could also have a nest level covariates, survey level covariates, and
# next x survey time covariates as well

nestdata <- readxl::read_excel(file.path("..","mallard.xlsx"), 
                               sheet="mallard")
# we need to add a NestId variables
nestdata$NestId <- paste("Nest", 1:nrow(nestdata),sep=".")
head(nestdata)
## # A tibble: 6 x 11
##   FirstFound LastPresent LastChecked  Fate  Freq Robel PpnGrass AgeFound AgeDay1 Habitat NestId
##        <dbl>       <dbl>       <dbl> <dbl> <dbl> <dbl>    <dbl>    <dbl>   <dbl> <chr>   <chr> 
## 1         73          89          89     0     1  6       0.800       13     -59 R       Nest.1
## 2         63          90          90     0     1  3.75    0.668        5     -57 P       Nest.2
## 3         70          70          76     1     1  3.75    0.800       13     -56 P       Nest.3
## 4         63          81          85     1     1  3.12    0.668        6     -56 N       Nest.4
## 5         61          61          66     1     1  4.5     0.668        4     -56 P       Nest.5
## 6         57          57          61     1     1  4.25    0.668        1     -55 P       Nest.6
# 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)
##   FirstFound LastPresent LastChecked Fate Freq Robel PpnGrass AgeFound AgeDay1 Habitat NestId NestId.num Day NestAge
## 1         73          89          89    0    1     6   0.8002       13     -59       R Nest.1          1  73      13
## 2         73          89          89    0    1     6   0.8002       13     -59       R Nest.1          1  74      14
## 3         73          89          89    0    1     6   0.8002       13     -59       R Nest.1          1  75      15
## 4         73          89          89    0    1     6   0.8002       13     -59       R Nest.1          1  76      16
## 5         73          89          89    0    1     6   0.8002       13     -59       R Nest.1          1  77      17
## 6         73          89          89    0    1     6   0.8002       13     -59       R Nest.1          1  78      18
# 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(  ~ Habitat, data=nesttime)

head(fe.design)
##   (Intercept) HabitatP HabitatR HabitatW
## 1           1        0        1        0
## 2           1        0        1        0
## 3           1        0        1        0
## 4           1        0        1        0
## 5           1        0        1        0
## 6           1        0        1        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: 762
##    Unobserved stochastic nodes: 4
##    Total graph size: 56228
## 
## 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

# in this case, we fit a S~Habitat model, so all rows of nesttime have the same estimated DSR
plotdata <- plyr::ddply(fitted.model$nesttime, "Habitat", function(x){ x[1,]})
head(plotdata)
##   NestId.num Day FirstFound LastPresent LastChecked Fate Freq Robel PpnGrass AgeFound AgeDay1 Habitat   NestId NestAge      mean          sd     X2.5.      X25.      X50.      X75.    X97.5.     Rhat
## 1        101  56         56          56          63    1    1 3.875   0.6684        4     -51       N  Nest.19       4 0.9459935 0.005195808 0.9354093 0.9425664 0.9462027 0.9495567 0.9556759 1.002369
## 2         10  43         43          55          60    1    1 3.375   0.5079        8     -34       P Nest.107       8 0.9563080 0.003375550 0.9496128 0.9540733 0.9563889 0.9585766 0.9627762 1.001515
## 3          1  73         73          89          89    0    1 6.000   0.8002       13     -59       R   Nest.1      13 0.9558656 0.008919650 0.9366432 0.9504301 0.9564252 0.9620530 0.9716889 1.001141
## 4         16  42         42          66          66    0    1 2.875   0.3344        7     -34       W Nest.112       7 0.9503930 0.010613186 0.9278923 0.9436996 0.9512304 0.9579799 0.9686667 1.000679
##   n.eff
## 1  3200
## 2  2400
## 3  4200
## 4  4500
ggplot(data=plotdata, aes(x=Habitat, y=mean))+
  ggtitle("Estimated DSR by Habitat")+
  geom_line(group=1)+
  geom_errorbar(aes(ymin=X2.5., ymax=X97.5.), width=0.1)+
  ylab("DSR (95% ci)")