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)

  psi.sp.uniq        ~ dnorm(0,0.001)
  psi.forest         ~ dnorm(0,0.001)
  psi.sp.uniq.forest ~ 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(site in 1:nsite) {
    for(sp in 1:nsp) {
      logit(psi[site,sp]) <-
        psi.0 +
        psi.sp[sp] +
        psi.sp.uniq * sp.uniq[sp] +
        psi.forest * forest[site] +
        psi.sp.uniq.forest * sp.uniq[sp] * forest[site]
    }
  }
  
  ## LIKELIHOOD
  ##
  for(site in 1:nsite) {
    for(sp in 1:nsp) {
      ## Occurrence
      Z[site,sp] ~ dbern(psi[site,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',
               'psi.sp.uniq',
               'psi.forest',
               'psi.sp.uniq.forest',
               '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,
            n.burnin=100,
            n.thin=10,
            n.chains=3,
            working.directory=NULL)

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

## look at parameters of interest
columns <- c('mean','sd','2.5%','97.5%', 'Rhat')
summ[c('psi.sp.uniq',
       'psi.forest',
       'psi.sp.uniq.forest'),columns]

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

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

## 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 for psi.sp.uniq
posterior.hist('psi.sp.uniq',
               lab='Effect of sp.uniq on occurrence',
               xlim=c(-1,1))
## negative values imply more 'uniqe' birds have lower rates of
## occurrence

## posterior for psi.forest
posterior.hist('psi.forest',
               lab='Effect of forest on occurrence',
               xlim=c(-1,1))
## positive values imply rates of occurrence are higher in 'forest'

## interaction
posterior.hist('psi.sp.uniq.forest',
               lab='Interaction effect on occurrence',
               xlim=c(-1,1))
## positive interaction means unique birds

values imply rates of occurrence are higher in 'forest'