# Nest survival model using JAGS 
# Fit a nest level continuous variable (robel height)

#
# 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(  ~ Robel, data=nesttime)

head(fe.design)
##   (Intercept) Robel
## 1           1     6
## 2           1     6
## 3           1     6
## 4           1     6
## 5           1     6
## 6           1     6
# 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: 2
##    Total graph size: 41574
## 
## 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~Robel model, so all rows of nesttime have the same estimated DSR
plotdata <- plyr::ddply(fitted.model$nesttime, "Robel", 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        337  14         14          37          37    0    1 0.625   0.9248        9      -4       N Nest.401       9 0.9485401 0.006815479 0.9342339 0.9442886 0.9489545 0.9531515 0.9611876 1.004390
## 2        158  27         27          30          34    1    1 0.750   0.9671        6     -20       N Nest.240       6 0.9487515 0.006545154 0.9350973 0.9446663 0.9491234 0.9531565 0.9608803 1.004272
## 3        267  12         12          21          25    1    1 0.875   0.8765        1     -10       R Nest.339       1 0.9489609 0.006279118 0.9358897 0.9449859 0.9493109 0.9532019 0.9606358 1.004144
## 4        163  22         22          45          53    1    1 1.000   0.6587        1     -20       R Nest.245       1 0.9491681 0.006017548 0.9367014 0.9453605 0.9494649 0.9532140 0.9604129 1.004007
## 5        136  30         30          56          56    0    1 1.125   0.6587        6     -23       P Nest.220       6 0.9493733 0.005760658 0.9374876 0.9456992 0.9496557 0.9532294 0.9601783 1.003858
## 6        352  11         11          35          35    0    1 1.250   0.2290        6      -4       N Nest.415       6 0.9495764 0.005508699 0.9382221 0.9460314 0.9498343 0.9532832 0.9599637 1.003697
##   n.eff
## 1   890
## 2   910
## 3   930
## 4   960
## 5   990
## 6  1000
ggplot(data=plotdata, aes(x=Robel, y=mean))+
  ggtitle("Estimated DSR as a function of Robel height")+
  geom_line(group=1)+
  geom_ribbon(aes(ymin=X2.5., ymax=X97.5.), alpha=0.1)+
  ylab("DSR (95% ci)")