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')






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)

  ## 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.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])

    }
  }
}

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)

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