################################################################################
## Multi-state occupancy dynamics code for:
## "Dynamics of tree-cavity occupancy support the hypothesis that low use of 
##  avian-excavated cavities is driven by resource suitability"
## D.J. Zavala, K.L. Cockle, M.R. Gomez, C.A. Ferreyra, E.B. Bonaparte, 
## F.G. DiSallo, and G. Ferraz
## DOI: 10.5281/zenodo.14262784
################################################################################

################################################################################
# 1. Clean up and load
################################################################################
rm(list = ls())
dat <- readRDS("CavNestMisiones.rds")
y <- dat$y
origen <- dat$origen

################################################################################
# 2. A little bit of data processing
################################################################################

## number of cavities
nsits <- dim(y)[1]
## number of years
nyears <- dim(y)[3]

## Compute nsurvey matrix: number of surveys per site/year
nsurvs <- array(1, dim = c(nsits, nyears)) # this is nsites x nyears
for(i in 1:nsits){
  for(t in 1:nyears){
    tmp <- which(!is.na(y[i,,t]))
    if(length(tmp) > 0){
      nsurvs[i,t] <- max(tmp)
    }
  }
}

## maximum state of each cavity per year
ymax <- mxstate <- apply(y, c(1,3), function(x) max(x, na.rm =TRUE))# maximum state
ymax[ymax == '-Inf'] <- 0
mxstate[mxstate == '-Inf'] <- NA
ysurv <-  apply(mxstate, c(1,2), function(x) as.numeric(x>0) )# Site*year evaluated
Nsurvs <- ysurv*nsurvs # nsurvs con 0 en los sitios sin visitas
Nsurvs[which(is.na(Nsurvs))] <- 0 
ysurv[which(is.na(ysurv))] <- 0 

## declare and fill vector with year of entry to database (first)
fyear <- apply(mxstate,1,function(x)min(which(!is.na(x))))
## Order rows of data according to year of entry (first)
r <- rank(fyear,ties.method = "first")
## f = year in which cavity entered the database, with values from 1 to 15
f <- rep(NA,nsits)
### fill array and vectors
for(i in 1:nsits){
  f[r[i]]<-fyear[i] # year of entry to the database
}
## number of cavities that were excavated by birds
nE <- length(which(origen==1))
## number of cavities that were not excavated by birds
nN <- length(which(origen==2))
## row index of excavated cavities in object ‘y’
Exc <- which(origen==1)
## row index of non-excavated cavities in object ‘y’
Nat <- which(origen==2)
## number of cavities known to be active (i.e. entered and not lost) per year
actsits <- cumsum(table(f))
actsits[nyears] <- actsits[nyears-1]
### number of excavated cavities known to be active per year.
actsitsE <- cumsum(table(f[Exc]))
actsitsE[nyears] <- actsitsE[nyears-1]
### number of non-excavated cavities known to be per year.
actsitsN <- cumsum(table(f[Nat]))
actsitsN <- c(actsitsN[1:3], actsitsN[3], actsitsN[4:14])
actsitsN[nyears] <- actsitsN[nyears-1]

################################################################################
# 3. Describe and save model
################################################################################
## Model 3: YEAR RANDOM EFFECTS, ORIGIN 
## This model is adapted from Chapter 6.4.3 of Applied hierarchical modeling 
# in ecology: Analysis of distribution, abundance and species richness in R 
# and BUGS: Volume 2

# Summarize and bundle data
str(bdata <- list(y = y, nsits = nsits, nsurvs = nsurvs, Nsurvs = Nsurvs, 
                  ymax=ymax, ysurv=ysurv,
                  nyears = dim(y)[3], origen = origen, 
                  actsits = as.vector(actsits), 
                  Nat = Nat, nN = nN, activeN = as.vector(actsitsN),
                  Exc = Exc, nE = nE, activeE = as.vector(actsitsE),
                  f = f, active = as.vector(actsits)))


cat(file = "dynMS.txt", "
model {

  ### (1) Priors for parameters
  ## (a) State process
  # Priors for parameters in initial state vector (Omega)
  for (i in 1:nsits){
    logit(psi[i]) <- alpha.lpsi[origen[i]]
    logit(r[i]) <- alpha.lr[origen[i]]
  }
  
  # Priors for parameters in the linear models of psi and r (Omega)
  # Region-specific intercepts (fixed effects)
  for (k in 1:2){
    alpha.lpsi[k] <- logit(mean.psi[k])
    mean.psi[k]  ~ dunif(0,1)
    alpha.lr[k] <- logit(mean.r[k])
    mean.r[k]  ~ dunif(0,1)
  }

  # Priors for parameters in state transition matrix (PhiMat; called 
  # Phi in text)
  # Transition probability matrix (PhiMat)
  for (i in 1:nsits){
    for(t in 1:(nyears-1)){
      logit(phi1[i,t]) <- alpha.lphi1[t] + beta.origen.lphi1[origen[i]]
      logit(rho1[i,t]) <- alpha.lrho1[t] + beta.origen.lrho1[origen[i]]
      logit(l1[i,t]) <- alpha.ll1[t] + beta.origen.ll1[origen[i]]
      logit(phi2[i,t]) <- alpha.lphi2[t] + beta.origen.lphi2[origen[i]]
      logit(rho2[i,t]) <- alpha.lrho2[t] + beta.origen.lrho2[origen[i]]
      logit(l2[i,t]) <- alpha.ll2[t] + beta.origen.ll2[origen[i]]
      logit(phi3[i,t]) <- alpha.lphi3[t] + beta.origen.lphi3[origen[i]]
      logit(rho3[i,t]) <- alpha.lrho3[t] + beta.origen.lrho3[origen[i]]
      logit(l3[i,t]) <- alpha.ll3[t] + beta.origen.ll3[origen[i]]
    }
  }

  # Priors for parameters in the linear models in PhiMat
  # Year-specific intercepts (random effects)
  for(t in 1:(nyears-1)){
    alpha.lphi1[t] ~ dnorm(mu.alpha.lphi1, tau.alpha.lphi1)
    alpha.lrho1[t] ~ dnorm(mu.alpha.lrho1, tau.alpha.lrho1)
    alpha.ll1[t] ~ dnorm(mu.alpha.ll1, tau.alpha.ll1)
    alpha.lphi2[t] ~ dnorm(mu.alpha.lphi2, tau.alpha.lphi2)
    alpha.lrho2[t] ~ dnorm(mu.alpha.lrho2, tau.alpha.lrho2)
    alpha.ll2[t] ~ dnorm(mu.alpha.ll2, tau.alpha.ll2)
    alpha.lphi3[t] ~ dnorm(mu.alpha.lphi3, tau.alpha.lphi3)
    alpha.lrho3[t] ~ dnorm(mu.alpha.lrho3, tau.alpha.lrho3)
    alpha.ll3[t] ~ dnorm(mu.alpha.ll3, tau.alpha.ll3)
  }

  # Hyperpriors for hyperparameters governing the random year effects

  # Hyperpriors for hyperparameters governing the random year effects
  mu.alpha.lphi1 <- logit(mean.phi1)
  mean.phi1 ~ dunif(0, 1)
  tau.alpha.lphi1 <- pow(sd.alpha.lphi1,-2)
  sd.alpha.lphi1 ~ dnt(0, 1/2.25^2, 1)I(0.01,)

  mu.alpha.lrho1 <- logit(mean.rho1)
  mean.rho1 ~ dunif(0, 1)
  tau.alpha.lrho1 <- pow(sd.alpha.lrho1,-2)
  sd.alpha.lrho1 ~ dnt(0, 1/2.25^2, 1)I(0.01,)

  mu.alpha.ll1 <- logit(mean.l1)
  mean.l1 ~ dunif(0, 1)
  tau.alpha.ll1 <- pow(sd.alpha.ll1,-2)
  sd.alpha.ll1 ~ dnt(0, 1/2.25^2, 1)I(0.01,)

  # -------------------------------
  mu.alpha.lphi2 <- logit(mean.phi2)
  mean.phi2 ~ dunif(0, 1)
  tau.alpha.lphi2 <- pow(sd.alpha.lphi2,-2)
  sd.alpha.lphi2 ~ dnt(0, 1/2.25^2, 1)I(0.01,)

  mu.alpha.lrho2 <- logit(mean.rho2)
  mean.rho2 ~ dunif(0, 1)
  tau.alpha.lrho2 <- pow(sd.alpha.lrho2,-2)
  sd.alpha.lrho2 ~ dnt(0, 1/2.25^2, 1)I(0.01,)

  mu.alpha.ll2 <- logit(mean.l2)
  mean.l2 ~ dunif(0, 1)
  tau.alpha.ll2 <- pow(sd.alpha.ll2,-2)
  sd.alpha.ll2 ~ dnt(0, 1/2.25^2, 1)I(0.01,)

  # -------------------------------
  mu.alpha.lphi3 <- logit(mean.phi3)
  mean.phi3 ~ dunif(0, 1)
  tau.alpha.lphi3 <- pow(sd.alpha.lphi3,-2)
  sd.alpha.lphi3 ~ dnt(0, 1/2.25^2, 1)I(0.01,)

  mu.alpha.lrho3 <- logit(mean.rho3)
  mean.rho3 ~ dunif(0, 1)
  tau.alpha.lrho3 <- pow(sd.alpha.lrho3,-2)
  sd.alpha.lrho3 ~ dnt(0, 1/2.25^2, 1)I(0.01,)

  mu.alpha.ll3 <- logit(mean.l3)
  mean.l3 ~ dunif(0, 1)
  tau.alpha.ll3 <- pow(sd.alpha.ll3,-2)
  sd.alpha.ll3 ~ dnt(0, 1/2.25^2, 1)I(0.01,)

  # Fixed effects of origen on the parameters in PhiMat
  beta.origen.lphi1[1] <- 0 # Avoid overparameterization
  beta.origen.lrho1[1] <- 0
  beta.origen.ll1[1] <- 0
  beta.origen.lphi2[1] <- 0 
  beta.origen.lrho2[1] <- 0
  beta.origen.ll2[1] <- 0
  beta.origen.lphi3[1] <- 0
  beta.origen.lrho3[1] <- 0
  beta.origen.ll3[1] <- 0
  
  beta.origen.lphi1[2] ~ dnorm(0, 0.1)
  beta.origen.lphi2[2] ~ dnorm(0, 0.1)
  beta.origen.lphi3[2] ~ dnorm(0, 0.1)
  beta.origen.lrho1[2] ~ dnorm(0, 0.1)
  beta.origen.lrho2[2] ~ dnorm(0, 0.1)
  beta.origen.lrho3[2] ~ dnorm(0, 0.1)
  beta.origen.ll1[2] ~ dnorm(0, 0.1)
  beta.origen.ll2[2] ~ dnorm(0, 0.1)
  beta.origen.ll3[2] ~ dnorm(0, 0.1)

  # (b) Observation process
  # Observation matrix (Theta)
  # Linear models
  for(t in 1:nyears){
    # Observation model for sites in occupied state 2 (= EXC)
      logit(p2[t]) <- alpha.lp2 
    # Observation model for sites in occupied state 3 (= SCN)
      logit(p3[t]) <- alpha.lp3 
  }
  
  # Priors for parameters in the linear models in Theta
  # Intercepts
  alpha.lp2 <- logit(mean.alpha.p2)
  mean.alpha.p2 ~ dunif(0, 1)

  alpha.lp3 <- logit(mean.alpha.p3)
  mean.alpha.p3 ~ dunif(0, 1)

  ### (2) Define relationships between basic model structure and parameters
  # Define initial state vector: Year 1
  for (i in 1:nsits){
    Omega[i,1] <- 1-psi[i]     # Prob empty in year 1
    Omega[i,2] <- psi[i]*(1-r[i]) # Prob occ EXC in year 1
    Omega[i,3] <- psi[i]*r[i]     # Prob occ SCN in year 1
    Omega[i,4] <- 0         # Prob cavity lost in year 1
  }

  # Define state transition probability matrix (PhiMat): years 2:nyears
  # Define probabilities of state S(t+1) given S(t)
  # For now, constant over sites and years
  # Note conditional Bernoulli parameterization of multinomial
  # Order of indices: Departing state, arrival state (state at t, state at t+1)
  for (i in 1:nsits){
    for(t in 1:(nyears-1)){
      PhiMat[1, 1, i, t] <- (1-l1[i,t]) * (1-phi1[i,t])
      PhiMat[1, 2, i, t] <- (1-l1[i,t]) * phi1[i,t] * (1-rho1[i,t])
      PhiMat[1, 3, i, t] <- (1-l1[i,t]) * phi1[i,t] * rho1[i,t]
      PhiMat[1, 4, i, t] <- l1[i,t]
    
      PhiMat[2, 1, i, t] <- (1-l2[i,t]) * (1-phi2[i,t])
      PhiMat[2, 2, i, t] <- (1-l2[i,t]) * phi2[i,t] * (1-rho2[i,t])
      PhiMat[2, 3, i, t] <- (1-l2[i,t]) * phi2[i,t] * rho2[i,t]
      PhiMat[2, 4, i, t] <- l2[i,t]
    
      PhiMat[3, 1, i, t] <- (1-l3[i,t]) * (1-phi3[i,t])
      PhiMat[3, 2, i, t] <- (1-l3[i,t]) * phi3[i,t] * (1-rho3[i,t])
      PhiMat[3, 3, i, t] <- (1-l3[i,t]) * phi3[i,t] * rho3[i,t]
      PhiMat[3, 4, i, t] <- l3[i,t]
    
      PhiMat[4, 1, i, t] <- 0
      PhiMat[4, 2, i, t] <- 0
      PhiMat[4, 3, i, t] <- 0
      PhiMat[4, 4, i, t] <- 1
    }
  }
  
  # Define observation probability matrix (Theta)
  # Order of indices: true state, observed state
  for(t in 1:nyears){
    Theta[1, 1, t] <- 1
    Theta[1, 2, t] <- 0
    Theta[1, 3, t] <- 0
    Theta[1, 4, t] <- 0
  
    Theta[2, 1, t] <- 1 - p2[t]
    Theta[2, 2, t] <- p2[t]
    Theta[2, 3, t] <- 0
    Theta[2, 4, t] <- 0
  
    Theta[3, 1, t] <- 1 - p3[t]
    Theta[3, 2, t] <- 0
    Theta[3, 3, t] <- p3[t]
    Theta[3, 4, t] <- 0
  
    Theta[4, 1, t] <- 0
    Theta[4, 2, t] <- 0
    Theta[4, 3, t] <- 0
    Theta[4, 4, t] <- 1
  }

  ### (3) Likelihood
  # Initial state: year 1
  for (i in 1:nsits){
    z[i,f[i]] ~ dcat(Omega[i,])
  }
  
  #State transitions from yearly interval 1:(nyears-1)
  for (i in 1:nsits){
    for(t in f[i]:(nyears-1)){
      z[i,t+1] ~ dcat(PhiMat[z[i,t], ,i,t])
    }
  }

  # Observation equation
  for (i in 1:nsits){
    for (t in f[i]:nyears){
      for (j in 1:nsurvs[i,t]){
        y[i,j,t] ~ dcat(Theta[z[i,t], ,t])
      }
    }
  }

  ### (4) Derived quantities

  # Annual average PhiMat By Origin
  # To Nat
  for(t in 1:(nyears-1)){
    PhiMat.annual.Nat[1,1,t] <- mean(PhiMat[1,1,Nat,t])
    PhiMat.annual.Nat[1,2,t] <- mean(PhiMat[1,2,Nat,t])
    PhiMat.annual.Nat[1,3,t] <- mean(PhiMat[1,3,Nat,t])
    PhiMat.annual.Nat[1,4,t] <- mean(PhiMat[1,4,Nat,t])
    PhiMat.annual.Nat[2,1,t] <- mean(PhiMat[2,1,Nat,t])
    PhiMat.annual.Nat[2,2,t] <- mean(PhiMat[2,2,Nat,t])
    PhiMat.annual.Nat[2,3,t] <- mean(PhiMat[2,3,Nat,t])
    PhiMat.annual.Nat[2,4,t] <- mean(PhiMat[2,4,Nat,t])
    PhiMat.annual.Nat[3,1,t] <- mean(PhiMat[3,1,Nat,t])
    PhiMat.annual.Nat[3,2,t] <- mean(PhiMat[3,2,Nat,t])
    PhiMat.annual.Nat[3,3,t] <- mean(PhiMat[3,3,Nat,t])
    PhiMat.annual.Nat[3,4,t] <- mean(PhiMat[3,4,Nat,t])
    PhiMat.annual.Nat[4,1,t] <- mean(PhiMat[4,1,Nat,t])
    PhiMat.annual.Nat[4,2,t] <- mean(PhiMat[4,2,Nat,t])
    PhiMat.annual.Nat[4,3,t] <- mean(PhiMat[4,3,Nat,t])
    PhiMat.annual.Nat[4,4,t] <- mean(PhiMat[4,4,Nat,t])
  }
  
  # Grand average NAT PhiMat (averaged over sites and years)
  PhiMat.avg.Nat[1,1] <- mean(PhiMat.annual.Nat[1,1,])
  PhiMat.avg.Nat[1,2] <- mean(PhiMat.annual.Nat[1,2,])
  PhiMat.avg.Nat[1,3] <- mean(PhiMat.annual.Nat[1,3,])
  PhiMat.avg.Nat[1,4] <- mean(PhiMat.annual.Nat[1,4,])
  PhiMat.avg.Nat[2,1] <- mean(PhiMat.annual.Nat[2,1,])
  PhiMat.avg.Nat[2,2] <- mean(PhiMat.annual.Nat[2,2,])
  PhiMat.avg.Nat[2,3] <- mean(PhiMat.annual.Nat[2,3,])
  PhiMat.avg.Nat[2,4] <- mean(PhiMat.annual.Nat[2,4,])
  PhiMat.avg.Nat[3,1] <- mean(PhiMat.annual.Nat[3,1,])
  PhiMat.avg.Nat[3,2] <- mean(PhiMat.annual.Nat[3,2,])
  PhiMat.avg.Nat[3,3] <- mean(PhiMat.annual.Nat[3,3,])
  PhiMat.avg.Nat[3,4] <- mean(PhiMat.annual.Nat[3,4,])
  PhiMat.avg.Nat[4,1] <- mean(PhiMat.annual.Nat[4,1,])
  PhiMat.avg.Nat[4,2] <- mean(PhiMat.annual.Nat[4,2,])
  PhiMat.avg.Nat[4,3] <- mean(PhiMat.annual.Nat[4,3,])
  PhiMat.avg.Nat[4,4] <- mean(PhiMat.annual.Nat[4,4,])

  
  # To Exc
  for(t in 1:(nyears-1)){
    PhiMat.annual.Exc[1,1,t] <- mean(PhiMat[1,1,Exc,t])
    PhiMat.annual.Exc[1,2,t] <- mean(PhiMat[1,2,Exc,t])
    PhiMat.annual.Exc[1,3,t] <- mean(PhiMat[1,3,Exc,t])
    PhiMat.annual.Exc[1,4,t] <- mean(PhiMat[1,4,Exc,t])
    PhiMat.annual.Exc[2,1,t] <- mean(PhiMat[2,1,Exc,t])
    PhiMat.annual.Exc[2,2,t] <- mean(PhiMat[2,2,Exc,t])
    PhiMat.annual.Exc[2,3,t] <- mean(PhiMat[2,3,Exc,t])
    PhiMat.annual.Exc[2,4,t] <- mean(PhiMat[2,4,Exc,t])
    PhiMat.annual.Exc[3,1,t] <- mean(PhiMat[3,1,Exc,t])
    PhiMat.annual.Exc[3,2,t] <- mean(PhiMat[3,2,Exc,t])
    PhiMat.annual.Exc[3,3,t] <- mean(PhiMat[3,3,Exc,t])
    PhiMat.annual.Exc[3,4,t] <- mean(PhiMat[3,4,Exc,t])
    PhiMat.annual.Exc[4,1,t] <- mean(PhiMat[4,1,Exc,t])
    PhiMat.annual.Exc[4,2,t] <- mean(PhiMat[4,2,Exc,t])
    PhiMat.annual.Exc[4,3,t] <- mean(PhiMat[4,3,Exc,t])
    PhiMat.annual.Exc[4,4,t] <- mean(PhiMat[4,4,Exc,t])
  }
  
  # Grand average EXC PhiMat (averaged over sites and years)
  PhiMat.avg.Exc[1,1] <- mean(PhiMat.annual.Exc[1,1,])
  PhiMat.avg.Exc[1,2] <- mean(PhiMat.annual.Exc[1,2,])
  PhiMat.avg.Exc[1,3] <- mean(PhiMat.annual.Exc[1,3,])
  PhiMat.avg.Exc[1,4] <- mean(PhiMat.annual.Exc[1,4,])
  PhiMat.avg.Exc[2,1] <- mean(PhiMat.annual.Exc[2,1,])
  PhiMat.avg.Exc[2,2] <- mean(PhiMat.annual.Exc[2,2,])
  PhiMat.avg.Exc[2,3] <- mean(PhiMat.annual.Exc[2,3,])
  PhiMat.avg.Exc[2,4] <- mean(PhiMat.annual.Exc[2,4,])
  PhiMat.avg.Exc[3,1] <- mean(PhiMat.annual.Exc[3,1,])
  PhiMat.avg.Exc[3,2] <- mean(PhiMat.annual.Exc[3,2,])
  PhiMat.avg.Exc[3,3] <- mean(PhiMat.annual.Exc[3,3,])
  PhiMat.avg.Exc[3,4] <- mean(PhiMat.annual.Exc[3,4,])
  PhiMat.avg.Exc[4,1] <- mean(PhiMat.annual.Exc[4,1,])
  PhiMat.avg.Exc[4,2] <- mean(PhiMat.annual.Exc[4,2,])
  PhiMat.avg.Exc[4,3] <- mean(PhiMat.annual.Exc[4,3,])
  PhiMat.avg.Exc[4,4] <- mean(PhiMat.annual.Exc[4,4,])

  
  # Annual average PhiMat
  for(t in 1:(nyears-1)){
    PhiMat.annual[1,1,t] <- mean(PhiMat[1,1,,t])
    PhiMat.annual[1,2,t] <- mean(PhiMat[1,2,,t])
    PhiMat.annual[1,3,t] <- mean(PhiMat[1,3,,t])
    PhiMat.annual[1,4,t] <- mean(PhiMat[1,4,,t])
    PhiMat.annual[2,1,t] <- mean(PhiMat[2,1,,t])
    PhiMat.annual[2,2,t] <- mean(PhiMat[2,2,,t])
    PhiMat.annual[2,3,t] <- mean(PhiMat[2,3,,t])
    PhiMat.annual[2,4,t] <- mean(PhiMat[2,4,,t])
    PhiMat.annual[3,1,t] <- mean(PhiMat[3,1,,t])
    PhiMat.annual[3,2,t] <- mean(PhiMat[3,2,,t])
    PhiMat.annual[3,3,t] <- mean(PhiMat[3,3,,t])
    PhiMat.annual[3,4,t] <- mean(PhiMat[3,4,,t])
    PhiMat.annual[4,1,t] <- mean(PhiMat[4,1,,t])
    PhiMat.annual[4,2,t] <- mean(PhiMat[4,2,,t])
    PhiMat.annual[4,3,t] <- mean(PhiMat[4,3,,t])
    PhiMat.annual[4,4,t] <- mean(PhiMat[4,4,,t])
  }
  
  # Grand average PhiMat (averaged over sites and years)
  PhiMat.avg[1,1] <- mean(PhiMat.annual[1,1,])
  PhiMat.avg[1,2] <- mean(PhiMat.annual[1,2,])
  PhiMat.avg[1,3] <- mean(PhiMat.annual[1,3,])
  PhiMat.avg[1,4] <- mean(PhiMat.annual[1,4,])
  PhiMat.avg[2,1] <- mean(PhiMat.annual[2,1,])
  PhiMat.avg[2,2] <- mean(PhiMat.annual[2,2,])
  PhiMat.avg[2,3] <- mean(PhiMat.annual[2,3,])
  PhiMat.avg[2,4] <- mean(PhiMat.annual[2,4,])
  PhiMat.avg[3,1] <- mean(PhiMat.annual[3,1,])
  PhiMat.avg[3,2] <- mean(PhiMat.annual[3,2,])
  PhiMat.avg[3,3] <- mean(PhiMat.annual[3,3,])
  PhiMat.avg[3,4] <- mean(PhiMat.annual[3,4,])
  PhiMat.avg[4,1] <- mean(PhiMat.annual[4,1,])
  PhiMat.avg[4,2] <- mean(PhiMat.annual[4,2,])
  PhiMat.avg[4,3] <- mean(PhiMat.annual[4,3,])
  PhiMat.avg[4,4] <- mean(PhiMat.annual[4,4,])

  
  # GoF part of the model
  # --------------------
  # Aggregate for observed data and draw 'replicate data' under the same model
  # and then aggregate these in the same way. In this way we then can compare
  # the observed frequencies with the frequencies that we would expect for 
  # a fitting model.  
  # (Note we change order of loops)
  
  for (i in 1:nsits){ # Loop over sites
    for (t in f[i]:nyears){ # Loop over years
      for (j in 1:nsurvs[i,t]){ # Loop over surveys
        # Create a new data set at each iteration. This data set is 'perfect' since
		    # we simulate it under the exact structural assumptions of our model, plus
		    # we use the exact values of the parameters right at each iteration of the 
		    # MCMC algorithm        
        yrep[i,j,t] ~ dcat(Theta[z[i,t], ,t])  # Replicated data
      }
      # Maximum observed state (cavity*year)
      # ymax[i,t] is given in the data bundle 
      # Maximum state of replicated data (cavity*year)
      yrepmax[i,t] <-  (max(yrep[i,1:nsurvs[i,t],t]))*ysurv[i,t]
      # Compute the probability that maximum observed state is 1 (eval1),
      # 2 (eval2), or 3 (eval3) for each site*year combination
      eval1[i,t] <- (((ysurv[i,t]*z[i,t])==1)*1)+
                    (((ysurv[i,t]*z[i,t])==2)*(1-p2[t])^Nsurvs[i,t])+
                    (((ysurv[i,t]*z[i,t])==3)*(1-p3[t])^Nsurvs[i,t])+
                    (((ysurv[i,t]*z[i,t])==4)*0)

      eval2[i,t] <- ((ysurv[i,t]*z[i,t])==2)*(1-(1-p2[t])^Nsurvs[i,t])+
                    (((ysurv[i,t]*z[i,t])==4)*0)

      eval3[i,t] <- ((ysurv[i,t]*z[i,t])==3)*(1-(1-p3[t])^Nsurvs[i,t])+
                    (((ysurv[i,t]*z[i,t])==4)*0)
      eval4[i,t] <- (((ysurv[i,t]*z[i,t])==4)*1)
      
    }

  }

  for (t in 1:nyears){
    # aggregate observed and simulated data by asking what is the number of
    # detections in state 1 (empty) on site i and year t 
    df1[t] <- sum(ymax[,t]==1)          # number of sites observed in max state 1
    df2[t] <- sum(ymax[,t]==2)          # number of sites observed in max state 2
    df3[t] <- sum(ymax[,t]==3)          # number of sites observed in max state 3
    df4[t] <- sum(ymax[,t]==4)          # number of sites observed in max state 4
    
    dfrep1[t] <- sum((yrepmax[1:actsits[t],t])==1)  # same for simulated data
    dfrep2[t] <- sum((yrepmax[1:actsits[t],t])==2)
    dfrep3[t] <- sum((yrepmax[1:actsits[t],t])==3)
    dfrep4[t] <- sum((yrepmax[1:actsits[t],t])==4)
    
    # Next tries to compute an expectation for df1 based on the average p per i, t and k
		# Then computes a Freeman-Tukey discrepancy for each i,t,k
		
		eval1sum[t] <- sum(eval1[1:actsits[t],t])
    eval2sum[t] <- sum(eval2[1:actsits[t],t])
    eval3sum[t] <- sum(eval3[1:actsits[t],t])
    eval4sum[t] <- sum(eval4[1:actsits[t],t])

    FTobs[t] <- sum((sqrt(df1[t])-sqrt(eval1sum[t]))^2, 
                    (sqrt(df2[t])-sqrt(eval2sum[t]))^2,
                    (sqrt(df3[t])-sqrt(eval3sum[t]))^2)#,
                    #(sqrt(df4[t])-sqrt(eval4sum[t]))^2)      # dist from pred to real
    
    FTobs1[t] <- sum((sqrt(df1[t])-sqrt(eval1sum[t]))^2)
    FTobs2[t] <- sum((sqrt(df2[t])-sqrt(eval2sum[t]))^2)      
    FTobs3[t] <- sum((sqrt(df3[t])-sqrt(eval3sum[t]))^2)      
    FTobs4[t] <- sum((sqrt(df4[t])-sqrt(eval4sum[t]))^2)      
    
    FTrep[t] <- sum((sqrt(dfrep1[t])-sqrt(eval1sum[t]))^2,   
                    (sqrt(dfrep2[t])-sqrt(eval2sum[t]))^2,
                    (sqrt(dfrep3[t])-sqrt(eval3sum[t]))^2) #,
                    #(sqrt(dfrep4[t])-sqrt(eval4sum[t]))^2)
    
    FTrep1[t] <- sum((sqrt(dfrep1[t])-sqrt(eval1sum[t]))^2)   
    FTrep2[t] <- sum((sqrt(dfrep2[t])-sqrt(eval2sum[t]))^2)   
    FTrep3[t] <- sum((sqrt(dfrep3[t])-sqrt(eval3sum[t]))^2)   
    FTrep4[t] <- sum((sqrt(dfrep4[t])-sqrt(eval4sum[t]))^2)   
    
    FTratio[t] <- FTobs[t] / (FTrep[t] + 0.0001)
    FTratio1[t] <- FTobs1[t] / (FTrep1[t] + 0.0001)
    FTratio2[t] <- FTobs2[t] / (FTrep2[t] + 0.0001)
    FTratio3[t] <- FTobs3[t] / (FTrep3[t] + 0.0001)
    FTratio4[t] <- FTobs4[t] / (FTrep4[t] + 0.0001)
  }

  FTobs1.tot <- sum(FTobs1)
  FTobs2.tot <- sum(FTobs2)
  FTobs3.tot <- sum(FTobs3)
  FTobs4.tot <- sum(FTobs4)
  
  FTrep1.tot <- sum(FTrep1)
  FTrep2.tot <- sum(FTrep2)
  FTrep3.tot <- sum(FTrep3)
  FTrep4.tot <- sum(FTrep4)
  
  FTobs.tot <- sum(FTobs)
  FTrep.tot <- sum(FTrep)

}
")


################################################################################
# 4. Fit the model and save results
################################################################################
# Initial values
#zst <- array(3, dim = c(bdata$nsits, bdata$nyears) )
zst <- array(1, dim = c(bdata$nsits, bdata$nyears))

for(i in 1:bdata$nsits) {
  zst[i,which(!is.na(mxstate[i,]))] <- mxstate[i,which(!is.na(mxstate[i,]))]
  if(any(mxstate[i,]==4,na.rm=TRUE)) {zst[i,which(mxstate[i,]==4):bdata$nyears]<-4}
  if(f[i]>1) {zst[i,1:(f[i]-1)]<-NA}
}
inits <- function(){list(z = zst)}

# Parameters monitored
params <- c("z","df1", "df2","df3",
            "dfrep1", "dfrep2", "dfrep3",
            "eval1sum", "eval2sum", "eval3sum",
            "eval1", "eval2", "eval3", "FTratio", "yrep","yrepmax",
            "FTobs.tot", "FTobs","FTobs1", "FTobs2", "FTobs3",  
            "FTrep.tot", "FTrep","FTrep1", "FTrep2", "FTrep3",
            "FTrep1.tot", "FTrep2.tot", "FTrep3.tot",
            "FTobs1.tot", "FTobs2.tot", "FTobs3.tot",
            "psi.Nat","psi.Exc","r.Nat","r.Exc",
            "beta.origen.lphi1","beta.origen.lphi1","beta.origen.lphi3",
            "beta.origen.lrho1","beta.origen.lrho2","beta.origen.lrho3",
            "beta.origen.ll1","beta.origen.ll2","beta.origen.ll3",
            "mu.alpha.lphi1", "mu.alpha.lphi2", "mu.alpha.lphi3",
            "mu.alpha.lrho1", "mu.alpha.lrho2", "mu.alpha.lrho3",
            "mu.alpha.ll1", "mu.alpha.ll2", "mu.alpha.ll3",
            "sd.alpha.lphi1", "sd.alpha.lphi2", "sd.alpha.lphi3",
            "sd.alpha.lrho1", "sd.alpha.lrho2", "sd.alpha.lrho3",
            "sd.alpha.ll1", "sd.alpha.ll2", "sd.alpha.ll3",
            "tau.alpha.lphi1", "tau.alpha.lphi2", "tau.alpha.lphi3", 
            "tau.alpha.lrho1", "tau.alpha.lrho2", "tau.alpha.lrho3",
            "tau.alpha.ll1", "tau.alpha.ll2", "tau.alpha.ll3",
            "alpha.lphi1", "alpha.lphi2", "alpha.lphi3",
            "alpha.lrho1", "alpha.lrho2", "alpha.lrho3",
            "alpha.ll1", "alpha.ll2", "alpha.ll3",
            "alpha.lpsi", "alpha.lr", 
            "PhiMat.annual.Nat", "PhiMat.avg.Nat",
            "PhiMat.annual.Exc", "PhiMat.avg.Exc",
            "PhiMat.annual","PhiMat.avg")


# MCMC settings
na <- 1000 ; ni <- 60000 ; nt <- 20 ; nb <- 40000 ; nc <- 3

# Call JAGS, check convergence and summarize posteriors
# odms stands for 'output dynamic multi-state'
library(jagsUI)
MSMO <- jags(bdata, inits, params, "dynMS.txt",
             n.adapt = na, n.chains = nc, n.thin = nt, n.iter = ni, 
             n.burnin = nb, parallel = TRUE)
print(MSMO, 3)
par(mfrow = c(3, 3)) ; traceplot(MSMO)
jags.View(MSMO) # not shown

################################################################################
# 5.Process results
################################################################################

# Load output from JAGS model
load("MSMO60k.RData")

# Transition probabilities
me <- MSMO$mean$PhiMat.avg 
q2.5 <- MSMO$q2.5$PhiMat.avg
q97.5 <- MSMO$q97.5$PhiMat.avg
# to Excavated cavities
meExc <- round(MSMO$mean$PhiMat.avg.Exc, 2)
q2.5Exc <- round(MSMO$q2.5$PhiMat.avg.Exc, 2)
q97.5Exc <- round(MSMO$q97.5$PhiMat.avg.Exc, 2)
# to Non-Excavated cavities
meNat <- round(MSMO$mean$PhiMat.avg.Nat, 2)
q2.5Nat <- round(MSMO$q2.5$PhiMat.avg.Nat, 2)
q97.5Nat <- round(MSMO$q97.5$PhiMat.avg.Nat, 2)

ProTra <- ProTraNat <- ProTraExc <- matrix(NA,nrow=4,ncol=4, byrow = T,
                                           dimnames = list(c("[1,]","[2,]","[3,]","[4,]"),
                                                           c("[,1]","[,2]","[,3]","[,4]")))
for (r in 1:4) {
  for (c in 1:4) {
    ProTra[r,c] <- paste(me[r,c],"(",q2.5[r,c],"-",q97.5[r,c],")")
    ProTraExc[r,c] <- paste(meExc[r,c],"(",q2.5Exc[r,c],"-",q97.5Exc[r,c],")")
    ProTraNat[r,c] <- paste(meNat[r,c],"(",q2.5Nat[r,c],"-",q97.5Nat[r,c],")")
  }
}
ProTra # values of the transition probabilities of all cavities
ProTraExc # values of the transition prob. of Excavated cavities (Fig. 3A)
ProTraNat # values of the transition prob. of No-excavated cavities (Fig. 3B)

################################################################################
# 6. Representation of Goodness of fit (GoF) Figure S2
################################################################################
par(mfrow=c(2,4), mar=c(3,4,3,0.5), oma=c(1,1,1,1))
# Bayesian p-value
Bp_v <- round(mean(MSMO$sims.list$FTrep.tot > MSMO$sims.list$FTobs.tot),2)
Bp_v1 <- round(mean(MSMO$sims.list$FTrep1.tot > MSMO$sims.list$FTobs1.tot),2)
Bp_v2 <- round(mean(MSMO$sims.list$FTrep2.tot > MSMO$sims.list$FTobs2.tot),2)
Bp_v3 <- round(mean(MSMO$sims.list$FTrep3.tot > MSMO$sims.list$FTobs3.tot),2)
# Graphics settings
x_lim <- c(0,3)
y_lim <- c(0,3)
#
df <- lm(as.data.frame(cbind(MSMO$sims.list$FTrep.tot, MSMO$sims.list$FTobs.tot)))
df1 <-lm(as.data.frame(cbind(MSMO$sims.list$FTrep1.tot, MSMO$sims.list$FTobs1.tot)))
df2 <-lm(as.data.frame(cbind(MSMO$sims.list$FTrep2.tot, MSMO$sims.list$FTobs2.tot))) 
df3 <-lm(as.data.frame(cbind(MSMO$sims.list$FTrep3.tot, MSMO$sims.list$FTobs3.tot)))
# Plot "State 1 - Empty"
plot(MSMO$sims.list$FTobs1.tot, MSMO$sims.list$FTrep1.tot, asp=1,
     xlab = "", ylab = "", las=1,
     pch = 16, col = rgb(0,0,0,0.3), xlim = x_lim, ylim = y_lim)
abline(0, 1)
abline(df1, lty = 5, col = "gray50")
points(mean(MSMO$sims.list$FTobs1.tot),mean(MSMO$sims.list$FTrep1.tot), 
       col="red", pch = 19)
title(main = "State 1 - Empty", 
      ylab="Discr. with sim. data", 
      xlab="Discrepancy with real data",
      cex.main=1, mgp = c(2,2,0),
      cex.lab=1)
text(1.5,2.7,paste("bpv = ",Bp_v1), cex = 0.8) # Bayesian p-value
# Plot "State 2 - EXC"
plot(MSMO$sims.list$FTobs2.tot, MSMO$sims.list$FTrep2.tot, asp=1,
     xlab = "", ylab = "",las=1,
     pch = 16, col = rgb(0,0,0,0.3), xlim = x_lim, ylim = y_lim)
abline(0, 1)
abline(df2, lty = 5, col = "gray50")
points(mean(MSMO$sims.list$FTobs2.tot),mean(MSMO$sims.list$FTrep2.tot), 
       col="red", pch = 19)
title(main = "State 2 - EXC",
      ylab="Discr. with sim. data", 
      xlab="Discrepancy with real data",
      cex.main=1, mgp = c(2,2,0),
      cex.lab=1)
text(1.5,2.7,paste("bpv = ",Bp_v2), cex = 0.8) # Bayesian p-value
# Plot "State 3 - SCN"
plot(MSMO$sims.list$FTobs3.tot, MSMO$sims.list$FTrep3.tot, asp=1,
     xlab = "", ylab = "",las=1,
     pch = 16, col = rgb(0,0,0,0.3), xlim = x_lim, ylim = y_lim)
abline(0, 1)
abline(df3, lty = 5, col = "gray50")
points(mean(MSMO$sims.list$FTobs3.tot),mean(MSMO$sims.list$FTrep3.tot), 
       col="red", pch = 19)
title(main = "State 3 - SCN",
      ylab="Discr. with sim. data", 
      xlab="Discrepancy with real data",
      cex.main=1, mgp = c(2,2,0),
      cex.lab=1)
text(1.5,2.7,paste("bpv = ",Bp_v3), cex = 0.8) # Bayesian p-value
# Plot "Total"
plot(MSMO$sims.list$FTobs.tot, MSMO$sims.list$FTrep.tot, asp=1,
     xlab = "", ylab = "",las=1,
     pch = 16, col = rgb(0,0,0,0.3), xlim = c(0,6), ylim = c(0,6))
abline(0, 1)
abline(df, lty = 5, col = "gray50")
points(mean(MSMO$sims.list$FTobs.tot),mean(MSMO$sims.list$FTrep.tot), 
       col="red", pch = 19)
title(main = "Total",
      ylab="Discr. with sim. data", 
      xlab="Discrepancy with real data",
      cex.main=1, mgp = c(2,2,0),
      cex.lab=1)
text(3,5.5,paste("bpv = ",Bp_v), cex = 0.8) # Bayesian p-value


# "Lack of fit ratio"
Ht <- round(mean(MSMO$sims.list$FTobs.tot/MSMO$sims.list$FTrep.tot),2)
H1 <- round(mean(MSMO$sims.list$FTobs1.tot/MSMO$sims.list$FTrep1.tot),2)
H2 <- round(mean(MSMO$sims.list$FTobs2.tot/MSMO$sims.list$FTrep2.tot),2)
H3 <- round(mean(MSMO$sims.list$FTobs3.tot/MSMO$sims.list$FTrep3.tot),2)
# Graphics settings
x_limH <- c(0,15)
y_limH <- c(0,320)
x_h <- 7.5
y_h <- 270

# to "State 1 - Empty"
hist(MSMO$sims.list$FTobs1.tot/MSMO$sims.list$FTrep1.tot, breaks = 100, col = 'gray50',
     border = "gray50",las=1,
     main = NULL, xlab = NULL,  ylab = NULL,
     xlim = x_limH ,ylim = y_limH)
abline(v = mean(MSMO$sims.list$FTobs1.tot/MSMO$sims.list$FTrep1.tot), 
       lwd = 1, col = 'blue')
title(ylab="Frequency", line=2.0, cex.lab=1,
      xlab="Lack of fit ratio")
text(x_h,y_h,paste("c-hat = ",H1), cex = 0.8)
# to "State 2 - EXC"
hist(MSMO$sims.list$FTobs2.tot/MSMO$sims.list$FTrep2.tot, breaks = 100, col = 'gray50',
     border = "gray50",las=1,
     main = NULL,  xlab = NULL, ylab = NULL,
     xlim = x_limH ,ylim = y_limH)
abline(v = mean(MSMO$sims.list$FTobs2.tot/MSMO$sims.list$FTrep2.tot), 
       lwd = 1, col = 'blue')
title(ylab="Frequency", line=2.0, cex.lab=1,
      xlab="Lack of fit ratio")
text(x_h,y_h,paste("c-hat = ",H2), cex = 0.8)
# to "State 3 - SCN"
hist(MSMO$sims.list$FTobs3.tot/MSMO$sims.list$FTrep3.tot, breaks = 100, col = 'gray50',
     border = "gray50",las=1,
     main = NULL,  xlab = NULL, ylab = NULL,
     xlim = x_limH ,ylim = y_limH)
abline(v = mean(MSMO$sims.list$FTobs3.tot/MSMO$sims.list$FTrep3.tot), 
       lwd = 1, col = 'blue')
title(ylab="Frequency", line=2.0, cex.lab=1,
      xlab="Lack of fit ratio")
text(x_h,y_h,paste("c-hat = ",H3), cex = 0.8)
# to "Total"
hist(MSMO$sims.list$FTobs.tot/MSMO$sims.list$FTrep.tot, breaks = 100, col = 'gray50',
     border = "gray50",las=1,
     main = NULL, xlab = NULL, ylab = NULL,
     xlim = x_limH ,ylim = y_limH)
abline(v = mean(MSMO$sims.list$FTobs.tot/MSMO$sims.list$FTrep.tot), 
       lwd = 1, col = 'blue')
title(ylab="Frequency", line=2.0, cex.lab=1,
      xlab="Lack of fit ratio")
text(x_h,y_h,paste("c-hat = ",Ht), cex = 0.8)


