rm(list=ls())

library(coda)
library(rjags)
library(R2jags)

## --------------------------------------------------
## load the data for all species
load(url('https://www.sfu.ca/~lmgonigl/materials-qm/data/occ-data.RData'),
     verbose=TRUE)

## package data
my.data <- list(X=X,
                forest=forest,
                sp.uniq=sp.uniq,
                nsite=nsite,
                nvisit=nvisit,
                nsp=nsp)
  
## inits
my.inits <- function() {
  list(Z=Z.init)
}

## model
my.model <- function() {

  ## PRIORS

  ## Intercepts
  p.0   ~ dnorm(0,0.001)
  psi.0 ~ dnorm(0,0.001)
  
  ## Random effects of species:
  ## Detection
  sigma.p.sp ~ dunif(0,20)
  tau.p.sp <- 1/(sigma.p.sp*sigma.p.sp)
  for(sp in 1:nsp) {
    p.sp[sp] ~ dnorm(0, tau.p.sp) 
  }
  ## Occurrence
  sigma.psi.sp ~ dunif(0,20)
  tau.psi.sp <- 1/(sigma.psi.sp*sigma.psi.sp)
  for(sp in 1:nsp) {
    psi.sp[sp] ~ dnorm(0, tau.psi.sp) 
  }

  ## Model for detection
  for(sp in 1:nsp) {
    logit(p[sp]) <- p.0 + p.sp[sp]
  }
  ## Model for occurrence
  for(sp in 1:nsp) {
    logit(psi[sp]) <- psi.0 + psi.sp[sp]
  }
  
  ## LIKELIHOOD
  ##
  for(site in 1:nsite) {
    for(sp in 1:nsp) {
      ## Occurrence
      Z[site,sp] ~ dbern(psi[sp])
      p.eff[site,sp] <- Z[site,sp] * p[sp]
      for(visit in 1:nvisit) {
        ## Detection 
        X[site,visit,sp] ~ dbern(p.eff[site,sp])
      }
    }
  }
  ## Derived quantities
  for(sp in 1:nsp) {
    num.occ[sp] <- sum(Z[,sp]) ## Number of occupied sites
  }
}

## parameters to track
my.params <- c('psi',
               'p',
               'num.occ',
               'sigma.p.sp',
               'sigma.psi.sp')
  
res <- jags(data=my.data,
            inits=my.inits,
            parameters.to.save=my.params,
            model.file=my.model,
            n.iter=1100, ## runs for actual analyses should be much longer
            n.burnin=100,
            n.thin=10,
            n.chains=3,
            working.directory=NULL)

## to make life easier
bugs <- res$BUGSoutput
summ <- res$BUGSoutput$summary

## figure out which values correspond to Boat-billed Flycatcher
which(dimnames(X)$species=='bbfl')
## and White-Ruffed Manakin
which(dimnames(X)$species=='wrma')

## inspect output
columns <- c('mean','sd','2.5%','97.5%', 'Rhat')
## note - we now have >300 p and psi values.  Let's just look at a
## couple
summ['p[12]',columns]
summ['p[280]',columns] ## similar rates of detection
summ['psi[12]',columns]
summ['psi[280]',columns] ## very different rates of occupancy
summ['num.occ[12]',columns]
summ['num.occ[280]',columns]

## plot posteriors
posterior.hist <- function(param, lab, xlim) {
  hist(bugs$sims.array[,,param],
       prob=TRUE, main='', las=1,
       xlab=lab, col='red', xlim=xlim)
  abline(v=0, lty=2, col='gray')
}

posterior.hist('psi[12]',
               lab='Boat-billed Flycatcher occurrence',
               xlim=c(0,1))
posterior.hist('psi[280]',
               lab='White-ruffed Manakin occurrence',
               xlim=c(0,1))

posterior.hist('p[12]',
               lab='Boat-billed Flycatcher detectability',
               xlim=c(0,1))
posterior.hist('p[280]',
               lab='White-ruffed Manakin detectability',
               xlim=c(0,1))

posterior.hist('num.occ[12]',
               lab='Number of occupied sites,\nBoat-billed Flycatcher',
               xlim=c(0,44))
posterior.hist('num.occ[280]',
               lab='Number of occupied sites,\nWhite-ruffed Manakin',
               xlim=c(0,44))