#Badia-Boher et al. - Code of manuscript "Strong impact of the recent Highly Pathogenic Avian Influenza panzootic on population dynamics of a long-lived bird"

#Code to analyse Dutch Peregrine Falcon data from 1993 to 2019, plus projections until 2034

#The model does not include the effect of HPAI

#Load necessary packages

require(jagsUI)
require(IPMbook)

#Load the data

load("data_nohpai.RData")


###Bundling the data####

jags.data <- list(
  
  #Mark-recapture data
  y = data_nohpai$ypool,
  nyears=ncol(data_nohpai$ypool), 
  nind = nrow(data_nohpai$ypool),
  f = data_nohpai$funic,
  age = data_nohpai$agesurvpool,
  agerec = data_nohpai$agerecpool,
  agedet = data_nohpai$agedetpool,
  nyscaled = as.numeric(scale(1:ncol(data_nohpai$ypool))),
  FR = data_nohpai$shortvec,
  fr = data_nohpai$shortvec,
  nstates = 6,
  
  #Productivity data
  
  J = data_nohpai$prodc$nfled,
  B = data_nohpai$prodc$nmonitored,
  
  #Count data and population dynamics model
  
  counts = data_nohpai$cens[18:28],
  pNinit1 = dUnif(1,70),
  pNinit2 = dUnif(1,30),
  pNinit3 = dUnif(1,10),
  pNinit4 = dUnif(1,3),
  pNinit5 = dUnif(1,100),
  yearspd = nrow(data_nohpai$prodc),
  yearsproj = 14
  
)


#JAGS code

cat(file = "pva2b.jags", "

    model{
    
    ###########################
    ###Productivity submodel###
    ###########################
    
    #Unstructured random effects over the years
    
    #Years with data (2010-2019)

  for(t in 1:yearspd){
    
    J[t] ~ dpois(B[t]*rho[t])
    log.rho[t] ~ dnorm(l.mean.rho, tau.rho)
    rho[t] <- exp(log.rho[t])
    
  }
  
    #Years with projections (2020-2034)
  
  for(t in (yearspd+1):(yearspd+yearsproj)){
  
    log.rho[t] ~ dnorm(l.mean.rho, tau.rho)
    rho[t] <- exp(log.rho[t])
  
  }
  
  #Vague priors for productivity
  
  mean.rho ~ dunif(0,5)
  l.mean.rho <- log(mean.rho)
  sigma.rho ~ dunif(0,5)
  tau.rho <- pow(sigma.rho, -2)
  
  ##################################
  ###Population dynamics submodel###
  ##################################  
  
  #T = 1 (2010)
  
  N[1,1] ~ dcat(pNinit1)
  N[2,1] ~ dcat(pNinit2)
  N[3,1] ~ dcat(pNinit3)
  N[4,1] ~ dcat(pNinit4)
  N[5,1] ~ dcat(pNinit3)
  N[6,1] ~ dcat(pNinit4)
  N[7,1] ~ dcat(pNinit4)
  N[8,1] ~ dcat(pNinit5)
  N[9,1] ~ dpois(omega[1])

  #Rest of years (2011-2034)
  
  for(t in 1:(yearspd-1+yearsproj)){
    
    N[1,t+1] ~ dpois(rho[t] / 2 * s[1,t+17] * NB[t]) #Juveniles
    N[2,t+1] ~ dbin(s[2,t+17]*gam[2], N[1,t]) #2yo first-time Breeders
    N[3,t+1] ~ dbin(s[2,t+17]*(1-gam[2]), N[1,t]) #2 yo NonBreeders
    N[4,t+1] ~ dbin(s[3,t+17]*gam[3], N[3,t]) #3yo first-time Breeders
    N[5,t+1] ~ dbin(s[3,t+17]*(1-gam[3]), N[3,t]) #3 yo NonBreeders
    N[6,t+1] ~ dbin(s[3,t+17]*gam[3], N[5,t]) #4yo first-time Breeders
    N[7,t+1] ~ dbin(s[3,t+17]*(1-gam[3]), N[5,t]) #4yo Nonbreeders
    N[8,t+1] ~ dbin(s[3,t+17], N[2,t]+N[4,t]+N[6,t]+N[7,t]+N[8,t]+N[9,t]) #Long-time breeders
    N[9,t+1] ~ dpois(omega[t+1]) #Immigrants 
    
    #Calculate the population growth rate
    
    lambda[t] <- Ntot[t+1]/(Ntot[t]+0.001) #Prevent division by zero
    lambdaNB[t] <- NB[t+1]/NB[t]
    
  }
  
      #Calculate the number of breeders, floaters, local recruits, total population size, proportions of every stage. Also, calculate the number of fledglings.
  
  for(t in 1:(yearspd+yearsproj)){
    
    NB[t] <- N[2,t]+N[4,t]+N[6,t]+N[8,t]+N[9,t] #Number of breeders 
    Nfloat[t] <- N[3,t]+N[5,t]+N[7,t] #Number of floaters
    Nlr[t] <- N[2,t]+N[4,t]+N[6,t] #Number of local recruits
    Ntot[t] <- sum(N[,t])
    propNB[t] <- NB[t]/Ntot[t]
    propfloat[t] <- Nfloat[t]/Ntot[t]
    propjuv[t] <- N[1,t]/Ntot[t]
    Nfledglings[t] ~ dpois(rho[t] / 2 * NB[t]) #Number of fledglings
    
  }
  
    #Link between expected counts and count data

  for(t in 1:yearspd){
  
    #Count model
    
    counts[t] ~ dpois(NB[t])
    
  }
  
  #Hidden parameter: Immigration - Unstructured random effects over the years (2010-2034)
  
  mean.omega ~ dunif(0.001, 50) 
  log.mean.omega <- log(mean.omega)
  sigma.omega ~ dunif(0.001,5) 
  tau.omega <- pow(sigma.omega, -2)
  
  for(t in 1:(yearspd+yearsproj)){
    
    log.omega[t] ~ dnorm(log.mean.omega, tau.omega)T(0, 4.6) #exp(4.6) = 100
    omega[t] <- exp(log.omega[t])
  }
    
    #############################
    ###Mark-recapture submodel###
    #############################

  for(a in 1:3){
  
   #Priors for age-structured survival
   
    mu.s[a] ~ dnorm(0, 0.001)
    mean.s[a] <- ilogit(mu.s[a])
    sigma.s[a] ~ dunif(0.001, 10)
    tau.s[a] <- pow(sigma.s[a],-2)
  
    #Priors for age-structured detection
    
    mu.p[a] ~ dnorm(0, 0.001)
    sigma.p[a] ~ dunif(0.001, 10)
    mean.p[a] <- ilogit(mu.p[a])
    tau.p[a] <- pow(sigma.p[a],-2)
    
        for(t in 1:(nyears-1)){

      #Unstructured random temporal effect on detection
      
      logit.p[a,t] ~ dnorm(mu.p[a], tau.p[a])
      p[a,t] <- ilogit(logit.p[a,t])
      
        }
  }
        
    for(a in 1:3){
      for(t in 1:(nyears-1+yearsproj)){

      #Unstructured random temporal effect on survival
      
      logit.s[a,t] ~ dnorm(mu.s[a], tau.s[a])
      s[a,t] <- ilogit(logit.s[a,t])
      
      }
    }
  
  #Recruitment probabilities
  
  gam[1] <- 0 #Assumed to be 0 for juveniles
  gam[2] ~ dbeta(1,1) #Estimate for subadults
  gam[3] ~ dbeta(1,1) #Estimate for 3 and 4 year olds
  gam[4] <- 1 #Assumed to be 1 for 5 year olds and older adults
  
  for(t in 1:(nyears-1)){
    
    #Linear trend over the years on recovery probability
    
    logit.r[t] <- alpha.r + beta.r * nyscaled[t]
    r[t] <- ilogit(logit.r[t])
    
  }
  
  #Priors for the linear effect on recovery over time
  alpha.r ~ dnorm(0, 0.001)
  beta.r ~ dnorm(0, 0.001)
  
  # Define state-transition and re-encounter probabilities
  
  for(i in 1:nind){
    for (t in f[i]:(nyears-1)){
      
      #STATES AND OBSERVATIONS:
      #1: Alive NB 2R
      #2: Alive NB 1R
      #3: Alive B 2R
      #4: Alive NB 2R
      #5: Dead
      #6: Absorbing state / Unseen
      
      #State transition probabilities
      
      psi[1,1,i,t] <- s[age[i,t],t]*(1-gam[agerec[i,t]])
      psi[1,2,i,t] <- 0
      psi[1,3,i,t] <- s[age[i,t],t]*gam[agerec[i,t]]
      psi[1,4,i,t] <- 0
      psi[1,5,i,t] <- 1-s[age[i,t],t]
      psi[1,6,i,t] <- 0
      
      psi[2,1,i,t] <- 0
      psi[2,2,i,t] <- s[age[i,t],t]*(1-gam[agerec[i,t]])
      psi[2,3,i,t] <- 0
      psi[2,4,i,t] <- s[age[i,t],t]*gam[agerec[i,t]]
      psi[2,5,i,t] <- 1-s[age[i,t],t]
      psi[2,6,i,t] <- 0
      
      psi[3,1,i,t] <- 0
      psi[3,2,i,t] <- 0
      psi[3,3,i,t] <- s[age[i,t],t]
      psi[3,4,i,t] <- 0
      psi[3,5,i,t] <- 1-s[age[i,t],t]
      psi[3,6,i,t] <- 0
      
      psi[4,1,i,t] <- 0
      psi[4,2,i,t] <- 0
      psi[4,3,i,t] <- 0
      psi[4,4,i,t] <- s[age[i,t],t]
      psi[4,5,i,t] <- 1-s[age[i,t],t]
      psi[4,6,i,t] <- 0
      
      psi[5,1,i,t] <- 0
      psi[5,2,i,t] <- 0
      psi[5,3,i,t] <- 0
      psi[5,4,i,t] <- 0
      psi[5,5,i,t] <- 0
      psi[5,6,i,t] <- 1
      
      psi[6,1,i,t] <- 0
      psi[6,2,i,t] <- 0
      psi[6,3,i,t] <- 0
      psi[6,4,i,t] <- 0
      psi[6,5,i,t] <- 0
      psi[6,6,i,t] <- 1
      
      #Observation probabilities
      
      #State NB2R
      
      po[1,1,i,t] <- p[agedet[i,t], t] #P(Seen as NB2R)
      po[1,2,i,t] <- 0 #P(Seen as NB1R)
      po[1,3,i,t] <- 0 #P(Seen as B2R)
      po[1,4,i,t] <- 0 #P(Seen as B1R)
      po[1,5,i,t] <- 0 #P(Seen as D)
      po[1,6,i,t] <- 1-p[agedet[i,t], t] #P(Not seen)

      #State NB1R
      
      po[2,1,i,t] <- 0 
      po[2,2,i,t] <- 0
      po[2,3,i,t] <- 0
      po[2,4,i,t] <- 0
      po[2,5,i,t] <- 0
      po[2,6,i,t] <- 1

      #State B2R
      
      po[3,1,i,t] <- 0 #P(Seen as NB2R)
      po[3,2,i,t] <- 0 #P(Seen as NB1R)
      po[3,3,i,t] <- p[agedet[i,t], t] #P(Seen as B2R)
      po[3,4,i,t] <- 0 #P(Seen as B1R)
      po[3,5,i,t] <- 0 #(Seen as D)
      po[3,6,i,t] <- 1-p[agedet[i,t], t] #P(Not Seen)

      #State B1R
      
      po[4,1,i,t] <- 0 
      po[4,2,i,t] <- 0
      po[4,3,i,t] <- 0
      po[4,4,i,t] <- 0
      po[4,5,i,t] <- 0
      po[4,6,i,t] <- 1
      
      #State D
      
      po[5,1,i,t] <- 0 
      po[5,2,i,t] <- 0
      po[5,3,i,t] <- 0
      po[5,4,i,t] <- 0
      po[5,5,i,t] <- r[t] #P(Seen as Dead)
      po[5,6,i,t] <- 1-r[t] #P(Not Seen)

      #Absorbing State
      
      po[6,1,i,t] <- 0 
      po[6,2,i,t] <- 0
      po[6,3,i,t] <- 0
      po[6,4,i,t] <- 0
      po[6,5,i,t] <- 0
      po[6,6,i,t] <- 1

    }
  }
  
  # Likelihood (Marginalized, Yackulic et al., 2020)
  
for (i in 1:nind){

  zeta[i,f[i],1] <- equals(1,y[i,f[i]])
  zeta[i,f[i],2] <- equals(2,y[i,f[i]])
  zeta[i,f[i],3] <- equals(3,y[i,f[i]])
  zeta[i,f[i],4] <- equals(4,y[i,f[i]])
  zeta[i,f[i],5] <- 0
  zeta[i,f[i],6] <- 0
  
  for (t in f[i]:(nyears-1)){
  
    zeta[i,t+1,1] <- (zeta[i,t,1:6] %*% psi[,1,i,t]) * po[1,y[i,t+1],i,t]
    zeta[i,t+1,2] <- (zeta[i,t,1:6] %*% psi[,2,i,t]) * po[2,y[i,t+1],i,t]
    zeta[i,t+1,3] <- (zeta[i,t,1:6] %*% psi[,3,i,t]) * po[3,y[i,t+1],i,t]
    zeta[i,t+1,4] <- (zeta[i,t,1:6] %*% psi[,4,i,t]) * po[4,y[i,t+1],i,t]
    zeta[i,t+1,5] <- (zeta[i,t,1:6] %*% psi[,5,i,t]) * po[5,y[i,t+1],i,t]
    zeta[i,t+1,6] <- (zeta[i,t,1:6] %*% psi[,6,i,t]) * po[6,y[i,t+1],i,t]
    
  } #t
  
  lik[i] <- sum(zeta[i,nyears,])
  fr[i] ~ dbin(lik[i], FR[i])
  }
    
    #Posterior predictive checks for the count and the productivity model

     for(t in 1:yearspd){
    
    #Posterior predictive checks of counts
    predcounts[t] ~ dpois(NB[t])
    #Distances from replicates to the model
    checkcounts[t,1] <- abs((predcounts[t] - NB[t])/(NB[t]+0.0001))
    #Distances from observations to the model
    checkcounts[t,2] <- abs((counts[t] - NB[t])/(NB[t]+0.0001))
    
    #Posterior predictive checks of productivity data.
    predfledg[t] ~ dpois(B[t]*rho[t])
    #Distances from replicates to the model
    checkfledg[t,1] <- pow(predfledg[t] - B[t]*rho[t], 2)/(B[t]*rho[t])
    #Distances from observations to the model
    checkfledg[t,2] <- pow(J[t] - B[t]*rho[t],2)/(B[t]*rho[t])
    
    }

    #Chi for count model is at row one
    chi2[1,1] <- sum(checkcounts[,1]) #Column 1 is simulated data
    chi2[1,2] <- sum(checkcounts[,2]) #Column 2 is estimated data
    
    #Chi for productivity model is at row 2
    chi2[2,1] <- sum(checkfledg[,1]) #Column 1 is simulated data
    chi2[2,2] <- sum(checkfledg[,2]) #Column 2 is estimated data

    
    }
    ")

nyears <- ncol(data_nohpai$ypool)
yearspd <- 11
yearsproj <- 14

#Initial values
inits <- function(){list(
  
  sigma.s = rep(.5,3),
  logit.s = matrix(nrow = 3, ncol = nyears-1+yearsproj, 0),
  mu.s = c(-1,0.2,1),
  mu.p = c(-1,0,1),
  sigma.p = rep(.5,3),
  logit.p = matrix(nrow = 3, ncol = nyears-1, -1),
  alpha.r = -1,
  beta.r = .5,
  gam = c(NA, rep(.5,2), NA),
  mean.rho = 2,
  sigma.rho = 1,
  log.rho = rep(1, yearspd+yearsproj)
  
)}

# Parameters monitored

parameters <- c("s", "mean.s", "sigma.s", "p", "sigma.p", "mean.p", "r", "alpha.r", "beta.r", "rho", "mean.rho", "sigma.rho", "mean.omega", "sigma.omega", "NB", "Nfloat", "Nlr", "N", "gam", "chi2", "lambda", "lambdaNB", "propNB", "propfloat", "propjuv", "Ntot", "Nfledglings")

# MCMC settings
ni <- 80000; nb <- 40000; nc <- 4; nt <- 40

#Run the model
out <- jags(jags.data, inits, parameters, "pva2b.jags", n.chains = nc, n.thin = nt, n.iter = ni, n.burnin = nb, parallel = T)

save(out, file = "pva_nohpai.RData")