#################################################################################################################
#Code for Weegman et al. "Assessing bias in demographic estimates from joint live and dead encounter models"
#Described in Kery and Schaub (2012), adapted by M. D. Weegman and S. Wilson 
#Here we present the dead recovery model, then the joint encounter model

library(jagsUI)

#Joint encounter model with simulated data
n.occasions<-10 # number of years
n.states<-4 # columns in the state transition matrix
n.obs<-3 # columns in the observation matrix
marked<-matrix(0,ncol=n.states,nrow=n.occasions) # create empty matrix
marked[,1]<-rep(100,n.occasions) # 100 individuals marked in each occasion

s <- c(0.9331172, 0.8798684, 0.9088144, 0.7409958, 0.8971829, 0.8559600, 0.9461083, 0.8986286, 0.8345931)

F <- c(0.8410903, 0.9284602, 0.6463537, 0.8295567, 0.7218508, 0.9922987, 0.8973964, 0.7416404, 0.8889492)

#r = 0.10
r <- c(0.096, 0.062, 0.13, 0.085, 0.118, 0.128, 0.111, 0.083, 0.095)

#r = 0.19
#r <- c(0.2185036, 0.2428580, 0.1739697, 0.1701664, 0.2029492, 0.2179421, 0.1957932, 0.1887818, 0.1350191)

#p = 0.01
p <- c(0.028888428, 0.011673430, 0.009540074, 0.007690508, 0.005852380, 0.008081784, 0.014782862, 0.017102660, 0.008751518)

#p = 0.20
#p0.20 <- c(0.30849219, 0.08051274, 0.12637864, 0.16525122, 0.36454542, 0.16144649, 0.21491193, 0.31294984, 0.20229925)

#p = 0.50
#p <- c(0.6415830, 0.3656194, 0.2716996, 0.5731303, 0.6288523, 0.5428100, 0.6428730, 0.3267440, 0.5878842)

#p = 0.90
#p <- c(0.9076922, 0.9401909, 0.8471082, 0.8408722, 0.9286633, 0.9171718, 0.8767357, 0.9238857, 0.9342344)

#Simulate multistate capture histories (from Kery and Schaub 2012)
simul.ms<-function(PSI.STATE,PSI.OBS,marked,unobservable=NA){
n.occasions<-dim(PSI.STATE)[4]+1
CH<-CH.TRUE<-matrix(NA,ncol=n.occasions,nrow=sum(marked)) # n.occasions = number of intervals and a row for each individual
# define vector with occasion of marking
mark.occ<-matrix(0,ncol=dim(PSI.STATE)[1],nrow=sum(marked))
g<-colSums(marked)
   for (s in 1:dim(PSI.STATE)[1]){
      if (g[s]==0) next  # To avoid error message if nothing to replace
      mark.occ[(cumsum(g[1:s])-g[s]+1)[s]:cumsum(g[1:s])[s],s] <-
      rep(1:n.occasions, marked[1:n.occasions,s])
      } #s
for (i in 1:sum(marked)){
	for (s in 1:dim(PSI.STATE)[1]){
	if (mark.occ[i,s]==0) next
	first<-mark.occ[i,s]
	CH[i,first]<-s
	CH.TRUE[i,first]<-s
	}
# above creates a matrix with the occasion of marking for each individual
	for (t in (first+1):n.occasions){ # for each occasion after marking
	# multinomial trials for state transitions
	if (first==n.occasions) next # except for birds banded on last occasion
	state<-which(rmultinom(1,1,PSI.STATE[CH.TRUE[i,t-1],,i,t-1])==1)
	CH.TRUE[i,t]<-state
      event <- which(rmultinom(1, 1, PSI.OBS[CH.TRUE[i,t],,i,t-1])==1)
      CH[i,t] <- event
         } #t
      } #i
# replace the NA and highest state number (4=dead) with 0
CH[is.na(CH)]<-0
CH[CH==dim(PSI.STATE)[1]] <- 0
   CH[CH==unobservable] <- 0
   id <- numeric(0)
   for (i in 1:dim(CH)[1]){
      z <- min(which(CH[i,]!=0))
      ifelse(z==dim(CH)[2], id <- c(id,i), id <- c(id))
      }
   return(list(CH=CH[-id,], CH.TRUE=CH.TRUE[-id,]))
   # CH: capture histories to be used
   # CH.TRUE: capture histories with perfect observation
   } 

# Define matrices with survival, transition and recapture probabilities
# These are 4-dimensional matrices, with 
   # Dimension 1: state of departure
   # Dimension 2: state of arrival
   # Dimension 3: individual
   # Dimension 4: time
# 1. State process matrix
totrel <- sum(marked)*(n.occasions-1)
PSI.STATE <- array(NA, dim=c(n.states, n.states, totrel, n.occasions-1))
for (i in 1:totrel){
   for (t in 1:(n.occasions-1)){
      PSI.STATE[,,i,t] <- matrix(c(
      s[t]*F[t], s[t]*(1-F[t]), 1-s[t], 0,
      0,   s[t],       1-s[t], 0,
      0,   0,       0,   1,
      0,   0,       0,   1), nrow = n.states, byrow = TRUE)
      } #t
   } #i

# 2.Observation process matrix
PSI.OBS <- array(NA, dim=c(n.states, n.obs, totrel, n.occasions-1))
for (i in 1:totrel){
   for (t in 1:(n.occasions-1)){
      PSI.OBS[,,i,t] <- matrix(c(
      p[t], 0, 1-p[t],
      0, 0, 1,
      0, r[t], 1-r[t],
      0, 0, 1), nrow = n.states, byrow = TRUE)
      } #t
   } #i

# Execute simulation function
sim <- simul.ms(PSI.STATE, PSI.OBS, marked)
CH <- sim$CH

# Compute date of first capture
get.first <- function(x) min(which(x!=0))
f <- apply(CH, 1, get.first)

# Recode CH matrix: note, a 0 is not allowed!
# 1 = alive and in study are, 2 = recovered dead, 3 = not seen or recovered
rCH <- CH  # Recoded CH
rCH[rCH==0] <- 3

sink("burn.sim.jags")
cat("
model {

# Parameters:
# s: true survival probability
# F: fidelity probability
# r: recovery probability
# p: recapture/resighting probability
# -------------------------------------------------
# States (S):
# 1 alive in study area
# 2 alive outside study area
# 3 recently dead and recovered
# 4 recently dead, but not recovered, or dead (absorbing)
# Observations (O):
# 1 seen alive
# 2 recovered dead
# 3 neither seen nor recovered
# -------------------------------------------------
for (t in 1:(n.occasions-1)){ 
   logit(s[t]) <- mu.s + eps.s[t]
   logit(F[t]) <- mu.F + eps.F[t]
   logit(r[t]) <- mu.r + eps.r[t]
   logit(p[t]) <- mu.p + eps.p[t]
	}

mu.s<-log(mean.s/(1-mean.s))
mean.s~dunif(0,1)
mean.F~dunif(0,1)
mu.F<-log(mean.F/(1-mean.F))
mean.r~dunif(0,1)
mu.r<-log(mean.r/(1-mean.r))
mean.p~dunif(0,1)
mu.p<-log(mean.p/(1-mean.p))

sd.s~dunif(0,10)
tau.s<-pow(sd.s,-2)
sd.F~dunif(0,10)
tau.F<-pow(sd.F,-2)
sd.r~dunif(0,10)
tau.r<-pow(sd.r,-2)
sd.p~dunif(0,10)
tau.p<-pow(sd.p,-2)

# error terms - time effect
for (t in 1:(n.occasions-1)){
	eps.s[t] ~ dnorm(0,tau.s)T(-10,10)
	eps.F[t] ~ dnorm(0,tau.F)T(-10,10)
	eps.r[t] ~ dnorm(0,tau.r)T(-10,10)
	eps.p[t] ~ dnorm(0,tau.p)T(-10,10)
	}

# define state and transition probabilities
for (t in 1:(n.occasions-1)){
      ps[1,t,1] <- s[t]*F[t] # 1 to 1 is a marked individual recaptured in study area
      ps[1,t,2] <- s[t]*(1-F[t])
      ps[1,t,3] <- (1-s[t])*r[t]
      ps[1,t,4] <- (1-s[t])*(1-r[t])
      ps[2,t,1] <- 0
      ps[2,t,2] <- s[t]
      ps[2,t,3] <- (1-s[t])*r[t]
      ps[2,t,4] <- (1-s[t])*(1-r[t])
      ps[3,t,1] <- 0
      ps[3,t,2] <- 0
      ps[3,t,3] <- 0
      ps[3,t,4] <- 1
      ps[4,t,1] <- 0 
      ps[4,t,2] <- 0
      ps[4,t,3] <- 0
      ps[4,t,4] <- 1

      # Define probabilities of O(t) given S(t)
      po[1,t,1] <- p[t]
      po[1,t,2] <- 0
      po[1,t,3] <- 1-p[t]
      po[2,t,1] <- 0
      po[2,t,2] <- 0
      po[2,t,3] <- 1
      po[3,t,1] <- 0
      po[3,t,2] <- 1
      po[3,t,3] <- 0
      po[4,t,1] <- 0
      po[4,t,2] <- 0
      po[4,t,3] <- 1
      } #t

# Likelihood 
for (i in 1:nind){
   # Define latent state at first capture
   z[i,f[i]] <- y[i,f[i]]
   for (t in (f[i]+1):n.occasions){
      # State process: draw S(t) given S(t-1)
      z[i,t] ~ dcat(ps[z[i,t-1], t-1,])
      # Observation process: draw O(t) given S(t)
      y[i,t] ~ dcat(po[z[i,t], t-1,])
      } #t
   } #i
}

",fill = TRUE)
sink()

ld.init <- function(ch, f){
   ch[ch==3] <- NA
   v2 <- which(ch==2, arr.ind = T)
   ch[v2] <- 3
   for (i in 1:nrow(v2)){
      ifelse(v2[i,2]!=ncol(ch), ch[v2[i,1], (v2[i,2]+1):ncol(ch)] <- 4, next)}
   for (i in 1:nrow(ch)){
      m <- max(which(ch[i,]==1))
      ch[i,f[i]:m] <- 1
      }
   for (i in 1:nrow(v2)){
      u1 <- min(which(ch[v2[i,1],]==1))
      ch[v2[i],u1:(v2[i,2]-1)] <- 1
      }
   for (i in 1:nrow(ch)){
      for (j in f[i]:ncol(ch)){
         if(is.na(ch[i,j])==1) ch[i,j] <- 1
         }
      ch[i,f[i]] <- NA
      }
   return(ch)
   }

ch <- rCH
ch[ch==3] <- NA
z.known <- ld.init(rCH, f)
z.known[is.na(ch)] <- NA
for (i in 1:nrow(ch)){
   z.known[i,f[i]] <- NA
   }
z <- ld.init(rCH, f)
z[!is.na(ch)] <- NA

# Bundle data
jags.data <- list(y = rCH, f = f, n.occasions = dim(rCH)[2], nind = dim(rCH)[1], z = z.known)

# Initial values 
inits <- function(){list(mean.s = 0.90, mean.F = 0.91, mean.p = 0.01, mean.r = 0.10,sd.s=1,sd.r=1,sd.F=1,sd.p=1,z = z)} 

# Parameters monitored
parameters <- c("mean.s", "s","mean.F","F", "mean.r", "r","mean.p","p","sd.s","sd.p","sd.r","sd.F")

# MCMC settings
#r-hat 1.1 
ni <- 500000
nt <- 10
nb <- 100000
nc <- 3

# Call JAGS from R 
burnham <- jags(jags.data, inits=inits, parallel=TRUE, parameters.to.save=parameters,  "burn.sim.jags", n.thin=nt, n.chains=nc, n.burnin=nb, n.iter=ni, DIC=FALSE)
print(burnham, digit = 3)

#################################################################################################################

#Dead recovery model with snow goose data
ad<-read.csv("Dead recovery model snow goose data.csv",header=FALSE)

DR<-as.matrix(ad)

## define function create m-arrays from the CH matrices
marray.dead<-function(DR){
	nind<-dim(DR)[1] # create matrix to hold information
	n.occasions<-dim(DR)[2]
	m.array<-matrix(data=0,ncol=n.occasions+1,nrow=n.occasions)
	#create vector with occasion of first marking
	get.first<-function(x)min(which(x!=0)) 
	f<-apply(DR,1,get.first) 
	# calculate the number of released individuals at each time period
	first<-as.numeric(table(f))
	for (t in 1:n.occasions){
		m.array[t,1]<-first[t]
		}
	# fill m-array with recovered individuals
   	# Fill m-array with recovered individuals
   	rec.ind <- which(apply(DR, 1, sum)==2)
   	rec <- numeric()
  	for (i in 1:length(rec.ind)){
      d <- which(DR[rec.ind[i],(f[rec.ind[i]]+1):n.occasions]==1)
      rec[i] <- d + f[rec.ind[i]]
      m.array[f[rec.ind[i]],rec[i]] <- m.array[f[rec.ind[i]],rec[i]] + 1
      }
   	# Calculate the number of individuals that are never recovered
   	for (t in 1:n.occasions){
      	m.array[t,n.occasions+1] <- m.array[t,1]-sum(m.array[t,2:n.occasions])
      }
   	out <- m.array[1:(n.occasions-1),2:(n.occasions+1)]
   return(out)
   }

# create m-arrays 
marr<-marray.dead(DR)

sink("snow.dr.txt")
cat("
model {

# Priors and constraints
for (t in 1:n.occasions){
   s[t]~dunif(0,1)
   r[t]~dunif(0.1,0.4)
  }

# Define the multinomial likelihood
for (t in 1:n.occasions){
   marr[t,1:(n.occasions+1)] ~ dmulti(pr[t,], rel[t])
   }
# Define the cell probabilities of the m-array
# Main diagonal
for (t in 1:n.occasions){
   pr[t,t] <- (1-s[t])*r[t]
   # Above main diagonal
   for (j in (t+1):n.occasions){
      pr[t,j] <- prod(s[t:(j-1)])*(1-s[j])*r[j]
      } #j
   # Below main diagonal
   for (j in 1:(t-1)){
      pr[t,j] <- 0
      } #j
   } #t
# Last column: probability of non-recovery
for (t in 1:n.occasions){
   pr[t,n.occasions+1] <- 1-sum(pr[t,1:n.occasions])
   } #t
}
",fill = TRUE)
sink()

# Bundle data
jags.data <- list(marr = marr, n.occasions = dim(marr)[2]-1, rel = rowSums(marr))

# Initial values
inits <- function(){list(s = runif(18, 0, 1), r = rep(0.2,18))}  

# Parameters monitored
parameters <- c("s", "r")

# MCMC settings
ni <- 50000
nt <- 10
nb <- 25000
nc <- 3

# Call JAGS from R (BRT <1 min)
dead.recovery <- jags(jags.data, inits, parameters, "snow.dr.txt", n.chains = nc, n.thin = nt, n.iter = ni, n.burnin = nb)

# Summarize posteriors
print(dead.recovery, digits = 3)

#################################################################################################################

#Joint encounter model with snow goose data
sngo<-read.csv("Joint encounter model snow goose data.csv",header=FALSE)
CH<-data.matrix(sngo)
get.first<-function(x)min(which(x==1)) 
f<-apply(CH,1,get.first) 

# recode matrix; 1=alive in study area, 2=recover dead, 3=not seen or recovered
# (recode so that 2=1, 3=2, 4=3 and 0=3)
rCH <- CH  # Recoded CH
rCH[rCH==0] <- 3

sink("snow.je.txt")
cat("
model {
#--------------Multistate Burnham---------------------------------------------
# Parameters:
# s: true survival probability
# F: fidelity probability
# r: recovery probability
# p: recapture/resighting probability
# -------------------------------------------------
# States (S):
# 1 alive in study area
# 2 alive outside study area
# 3 recently dead and recovered
# 4 recently dead, but not recovered, or dead (absorbing)
# Observations (O):
# 1 seen alive
# 2 recovered dead
# 3 neither seen nor recovered
# -------------------------------------------------
for (t in 1:(n.ms.occasions-1)){ 
   logit(s[t]) <- mu.s + eps.s[t]
   logit(F[t]) <- mu.F + eps.F[t]
   logit(r[t]) <- mu.r + eps.r[t]
   logit(p[t]) <- mu.p + eps.p[t]
	}
mean.s~dunif(0.5,1)
mu.s<-log(mean.s/(1-mean.s))
mean.F~dunif(0.5,1)
mu.F<-log(mean.F/(1-mean.F))
mean.r~dunif(0,0.5)
mu.r<-log(mean.r/(1-mean.r))
mean.p~dunif(0,0.3)
mu.p<-log(mean.p/(1-mean.p))

sd.s~dunif(0,10)
tau.s<-pow(sd.s,-2)
sd.F~dunif(0,10)
tau.F<-pow(sd.F,-2)
sd.r~dunif(0,10)
tau.r<-pow(sd.r,-2)
sd.p~dunif(0,10)
tau.p<-pow(sd.p,-2)

# error terms - time effect
for (t in 1:(n.ms.occasions-1)){
	eps.s[t] ~ dnorm(0,tau.s)T(-10,10)
	eps.F[t] ~ dnorm(0,tau.F)T(-10,10)
	eps.r[t] ~ dnorm(0,tau.r)T(-10,10)
	eps.p[t] ~ dnorm(0,tau.p)T(-10,10)
}
# define state and transition probabilities
for (t in 1:(n.ms.occasions-1)){
      ps[1,t,1] <- s[t]*F[t] # 1 to 1 is a marked individual recaptured in study area
      ps[1,t,2] <- s[t]*(1-F[t])
      ps[1,t,3] <- (1-s[t])*r[t]
      ps[1,t,4] <- (1-s[t])*(1-r[t])
      ps[2,t,1] <- 0
      ps[2,t,2] <- s[t]
      ps[2,t,3] <- (1-s[t])*r[t]
      ps[2,t,4] <- (1-s[t])*(1-r[t])
      ps[3,t,1] <- 0
      ps[3,t,2] <- 0
      ps[3,t,3] <- 0
      ps[3,t,4] <- 1
      ps[4,t,1] <- 0
      ps[4,t,2] <- 0
      ps[4,t,3] <- 0
      ps[4,t,4] <- 1

      # Define probabilities of O(t) given S(t)
      po[1,t,1] <- p[t]
      po[1,t,2] <- 0
      po[1,t,3] <- 1-p[t]
      po[2,t,1] <- 0
      po[2,t,2] <- 0
      po[2,t,3] <- 1
      po[3,t,1] <- 0
      po[3,t,2] <- 1
      po[3,t,3] <- 0
      po[4,t,1] <- 0
      po[4,t,2] <- 0
      po[4,t,3] <- 1
      } #t

# Likelihood 
for (i in 1:nind){
   # Define latent state at first capture
   z[i,f[i]] <- y[i,f[i]]
   for (t in (f[i]+1):n.ms.occasions){
      # State process: draw S(t) given S(t-1)
      z[i,t] ~ dcat(ps[z[i,t-1], t-1,])
      # Observation process: draw O(t) given S(t)
      y[i,t] ~ dcat(po[z[i,t], t-1,])
      } #t
   } #i
}
",fill = TRUE)
sink()

# Function from Kery and Schaub (2012):
ld.init <- function(ch, f){
   ch[ch==3] <- NA
   v2 <- which(ch==2, arr.ind = T)
   ch[v2] <- 3
   for (i in 1:nrow(v2)){
      ifelse(v2[i,2]!=ncol(ch), ch[v2[i,1], (v2[i,2]+1):ncol(ch)] <- 4, next)}
   for (i in 1:nrow(ch)){
      m <- max(which(ch[i,]==1))
      ch[i,f[i]:m] <- 1
      }
   for (i in 1:nrow(v2)){
      u1 <- min(which(ch[v2[i,1],]==1))
      ch[v2[i],u1:(v2[i,2]-1)] <- 1
      }
   for (i in 1:nrow(ch)){
      for (j in f[i]:ncol(ch)){
         if(is.na(ch[i,j])==1) ch[i,j] <- 1
         }
      ch[i,f[i]] <- NA
      }
   return(ch)
   }

ch <- rCH
ch[ch==3] <- NA
z.known <- ld.init(rCH, f)
z.known[is.na(ch)] <- NA
for (i in 1:nrow(ch)){
   z.known[i,f[i]] <- NA
   }
z <- ld.init(rCH, f)
z[!is.na(ch)] <- NA

# Bundle data
jags.data <- list(y = rCH,f = f,n.ms.occasions = dim(rCH)[2],nind = dim(rCH)[1],z = z.known)

# Initial values 
inits <- function(){list(mean.s = 0.9,mean.F = 0.9,mean.p = 0.015,mean.r = 0.18,sd.s=1,sd.r=1,sd.F=1,sd.p=1,z = z)}  

# Parameters monitored
parameters <- c("s","F","r","p","mean.s","mean.F","mean.r","mean.p","sd.s","sd.F","sd.r","sd.p")

# MCMC settings
ni <- 80000
nt <- 10
nb <- 40000
nc <- 3

joint.encounter <- jags(jags.data, inits=inits, parallel = TRUE, parameters.to.save=parameters,  "snow.je.txt", n.thin=nt, n.chains=nc, n.burnin=nb, n.iter=ni)
print(joint.encounter, digit = 3)
#################################################################################################################