rm(list=ls())

library(rjags)
library(R2jags)

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

## inspect each item that you just loaded
head(captured)
head(alive) ## this is constructed from 'captured'
head(phenotype)
head(soil)
head(site)

## package data for JAGS
my.data <- list(captured=captured,
                alive=alive,
                phenotype=as.integer(phenotype)-1,
                site=as.integer(as.factor(site)),
                soil=as.integer(soil)-1,
                nind=nrow(captured),
                nrep=ncol(captured),
                nsite=length(unique(site)))
  
## create inits (empty list - JAGS will initialize from priors)
my.inits <- function() {
  list()
}
  
## specify the parameters to be monitored
my.params <- c('p.0',
               'p.phenotype', 
               'p.soil',
               'p.soil.phenotype',
               'phi.0',
               'phi.phenotype', 
               'phi.soil',
               'phi.soil.phenotype',
               'sigma.phi.site',
               'phi.lm.ls',
               'phi.lm.ds',
               'phi.dm.ls',
               'phi.dm.ds')

my.model <- function() {

  ## priors (recapture)
  p.0                ~ dnorm(0,0.001)
  p.soil             ~ dnorm(0,0.001)
  p.phenotype        ~ dnorm(0,0.001)
  p.soil.phenotype   ~ dnorm(0,0.001)
  ## priors (survival)
  phi.0              ~ dnorm(0,0.001)
  phi.soil           ~ dnorm(0,0.001)
  phi.phenotype      ~ dnorm(0,0.001)
  phi.soil.phenotype ~ dnorm(0,0.001)

  ## random effect of site
  sigma.phi.site ~ dunif(0,20)
  tau.phi.site <- 1/(sigma.phi.site*sigma.phi.site)
  for(ss in 1:nsite) {
    phi.site[ss] ~ dnorm(0, tau.phi.site)
  }

  ## for each individual
  for(ind in 1:nind) {

    ## model for recapture
    logit(p[ind]) <-
      p.0 +
      p.phenotype*phenotype[ind] +
      p.soil*soil[ind] +
      p.soil.phenotype*soil[ind]*phenotype[ind]

    ## model for survival
    logit(phi[ind]) <-
      phi.0 +
      phi.site[site[ind]] +
      phi.phenotype*phenotype[ind] +
      phi.soil*soil[ind] +
      phi.soil.phenotype*soil[ind]*phenotype[ind]

    ## likelihood
    for(rep in 2:nrep) {

      ## state equation
      alivep[ind,rep] <- phi[ind] * alive[ind,rep-1] 
      alive[ind,rep] ~ dbern(alivep[ind,rep])

      ## observation equation
      sightp[ind,rep] <- p[ind] * alive[ind, rep] 
      captured[ind,rep] ~ dbern(sightp[ind,rep])

    }
  }

  logit(phi.dm.ls) <-
    phi.0 +
    phi.phenotype*0 +
    phi.soil*1 +
    phi.soil.phenotype*1*0

  logit(phi.dm.ds) <-
    phi.0 +
    phi.phenotype*0 +
    phi.soil*0 +
    phi.soil.phenotype*0*0

  logit(phi.lm.ls) <-
    phi.0 +
    phi.phenotype*1 +
    phi.soil*1 +
    phi.soil.phenotype*1*1

  logit(phi.lm.ds) <-
    phi.0 +
    phi.phenotype*1 +
    phi.soil*0 +
    phi.soil.phenotype*0*1
  
}

res <- jags(data=my.data,
            inits=my.inits,
            parameters.to.save=my.params,
            model.file=my.model,
            n.iter=11000,
            n.burnin=1000,
            n.thin=10,
            n.chains=3,
            working.directory=NULL)

## inspect output
columns <- c('mean','sd','2.5%','97.5%')
res$BUGSoutput$summary[,columns]

## I like to make an object called 'bugs' and 'summ' to save me typing
## 'res$BUGSoutput' and res$BUGSoutput$summary all the time
bugs <- res$BUGSoutput
summ <- res$BUGSoutput$summary

## 1. Posterior for 'phi.soil.phenotype'
hist(bugs$sims.array[,,'phi.soil.phenotype'], col='red',
     prob=TRUE, main='', breaks=40, las=1,
     xlab='phi.soil.phenotype', xlim=c(0,4))

## 2. Posterior for 'phi.soil.phenotype' with 95% BCI

h <- hist(bugs$sims.array[,,'phi.soil.phenotype'],
          breaks=40, plot=FALSE)
cuts <- cut(h$breaks, c(-Inf,
                        summ['phi.soil.phenotype','2.5%'],
                        summ['phi.soil.phenotype','97.5%'],
                        Inf))
cols <- rep('blue', length(cuts))
cols[as.integer(cuts)!=2] <- 'red'
hist(bugs$sims.array[,,'phi.soil.phenotype'],
     prob=TRUE, main='', breaks=40, las=1,
     xlab='phi.soil.phenotype', xlim=c(0,4), col=cols)

## 3. The four values have been added to the model.  Their means and
## BCI can be accessed with
vars <- c('phi.dm.ds', 'phi.dm.ls', 'phi.lm.ds', 'phi.lm.ls')
summ[vars,c('mean', '2.5%', '97.5%')]
## I'll leave it up to you to plot (arrows is handy command for the
## BCI)  

## 4. The above model has been updated.  'phi.site' includes
## site-level effects, and 'sigma.phi.site' is the width of the random
## effect.
summ['sigma.phi.site',c('mean', '2.5%', '97.5%')]

## 5. assess convergence
matplot(res$BUGSoutput$sims.array[,,'sigma.phi.site'],
        type='l', lty=1, xlab='Iteration', ylab='sigma.phi.site')
## hmm... convergence was not very good for me, so I ran it for 10x
## longer and it looked much better