###########################################################################
#Script used to run Cormack-Jolly-Seber models using JAGS 4.0 to estimate stopover duration of migrant songbirds
#This script runs one model and outputs three files (summary, trace plots, list of draws)
#Text files of code ("cjshierarchical_1.txt") for hierarchical models available in supplemental materials
#Code written by N.N.Dorian
#Last revised: December 17, 2019
############################################################################


###############################
#Establish working environment
###############################
#set working directory
setwd()

#load functions (available in supplemental materials) to set initial values, from Kéry and Schaub 2012 chpt 7 
#NB:known.state.cjs takes one arguments: the data matrix 
source("cjs_ms_functions.R")

#you may need to use install.packages("package name") first
library(rjags) #run JAGS 4.0 in R
library(jagsUI) #streamlined JAGS interface
library(MCMCvis) #to visualize model results using trace plots

#load functions (available in supplemental materials, S2D) to set initial values, from Kéry and Schaub 2012 chpt 7 and 9
#NB:known.state.cjs takes one arguments: the data matrix 
source("msfunctions.R")

###############################
#Format data for model run
###############################
#load capture histories for early and late period and concatenate
#cleaned captured histories for each species available on Dryad
CH<-read.csv("veerch.fall.7015.csv")

#ensure that no rows in the data contain all zeros
torem<-which(apply(CH[,5:dim(CH)[2]],1, sum) == 0)
if(length(torem > 0)){CH<-CH[-torem,]}

#grouping vector with the year of capture for each individual
totyr<-length(unique(CH$year))
yeartomatch<-as.data.frame(cbind(1:totyr,sort(unique(as.character(CH$year)))))
CH<-merge(CH, yeartomatch, by.x = "year", by.y = "V2")
year<-as.numeric(as.character(CH$V1))

#extract only capture histories
rCH<-CH[,4:(dim(CH)[2]-1)]
rCH<-as.matrix(apply(rCH, 2, as.numeric))
dim(rCH)

#get vector of initial values--when was each bird captured first?
f<-apply(rCH, 1, get.first)

###############################
#Gather data for model run
###############################
#y = capture history matrix
#f = occasions of first capture
#n.occasions = number of banding days per season (columns)
#nind = number of individuals captured in dataset (rows)
#z = capture history matrix with known values included (e.g. a capture history of 10010 inputted into this function becomes 11110 because the bird had to be alive between the two capture occasions)
#g = number of periods, always 2
#group = period grouping vector
#nyear = number of years (hierarchical model only)
#year = year grouping vector (hierarchical model only)
jags.data<-list(y = rCH, f = f, n.occasions = dim(rCH)[2], nind = dim(rCH)[1], z = known.state.cjs(rCH), nyear = length(unique(year)), year = year)

#set initial values
#for phi.logit and p.logit, values are on the logit scale
inits<-function(){list(z = cjs.init.z(rCH, f), phi.logit = 1, p.logit = 0, sigma.phi = runif(1,0,5),sigma.p = runif(1,0,5))}

#choose parameters to monitor
parameters<-c("phi.mu","stay.mu","p.mu", "year.phi", "year.p", "sigma2.phi", "sigma2.p")


cjsrun<-jags(data = jags.data,#data list 
             inits = inits, #vectors of initial values
             parameters.to.save = parameters, #vector of parameters
             model.file = "cjshierarchical_1.txt", 
             n.thin = 3, #retain a third of all samples
             n.chains = 3, #run across 3 chains
             n.adapt = 100, #run the model for 5000 samples to initialize
             n.burnin = 10000, #throw-away first 10,000 samples
             n.iter = 30000, #run for total of 30,000 samples 
             parallel = TRUE) #run across multiple core processes

###############################
#Export model results
###############################
#save traces of chains, uses package MCMCvis
MCMCtrace(cjsrun, open_pdf = F, filename = "cjsrun.trace")

#export model summary to CSV
cjsrun.summary<-MCMCsummary(cjsrun, digits = 3, n.eff = T)
write.csv(cjsrun.summary, "cjsrun.summary.csv")

#export simslist to CSV, needed for posterior predictive checks
cjsrun.simslist<-write.csv(cjsrun$sims.list, "cjsrun.simslist.csv")