##    Echological Archieves Supplement for Ecology
##
##     Isojunno SI, Miller PJO (2014) Sperm whale response to tag boat presence: 
##     biologically informed hidden state models quantify lost feeding opportunities 
##


##    Gibbs sample 2 hidden state models in jags from within R
##    ########################################################

##    1. Model structure 2 (depth + clicking) with 5 states and fixed step length random walk
##    2. Model structure 5 (minFromSurf + ODBA + pitch ) with 6 states and variable step length

##    In order to run script, download jags at: http://mcmc-jags.sourceforge.net/ 
##    and install R2jags package in r



library(R2jags)
set.seed(0)


# load data and initial values

load("jags_data.Rd") # contains:

  str(base.data) # list - data & priors for transition matrix to fit base model (1)
  str(full.data) # list - data & priors for transition matrix to fit full model (2)
    
  # note: the data are formatted for each model from the same dataset (2 tag records)

  str(base.inits) # function - fixed initial values for base model (1)
  str(full.inits) # function - fixed initial values for full model (2)

  str(base.fit$BUGSoutput$mean) # R2jags fit object - example run for base model (posterior mean values shown)
  str(full.fit$BUGSoutput$mean) # R2jags fit object - example run for full model (posterior mean values shown)


# set parameters to monitor

  base.params <- c("T.beta0","T.beta11","bd","tau", "drift", "RC.beta0")
  full.params <- c(base.params, "T.beta31", "a.mean", "a.rate", "p.beta0", "p.beta1","p.gamma")


## FITTING MODELS

# sampling / updating takes easily hours;
# example fit objects base.fit and full.fit supplied

if(FALSE) {

  ## 1. base model
  
    base.fit = jags(data=base.data, inits=base.inits, parameters.to.save=base.params,
               model.file="base_model_clicking.txt",
               n.chains=3, n.iter=12000, n.burnin=6000, n.thin=18)
  
  
  
  ## 2. full model
  
    full.fit = jags(data=full.data, inits=full.inits, parameters.to.save=full.params,
               model.file="full_model.txt",
               n.chains=3, n.iter=12000, n.burnin=6000, n.thin=18)
  }


## plot iteration histories

  traceplot(base.fit, varname=base.params[-3]) 
  traceplot(full.fit, varname=full.params[-3]) 


## calculate posterior probabilities for states in each time step

  base.props <- matrix(0, base.data$stateN, base.data$N)
  full.props <- matrix(0, full.data$stateN, full.data$N)
  for(s in 1:5) {
    base.props[s,] <- apply(base.fit$BUGSoutput$sims.list$bd==s, 2, mean)
    full.props[s,] <- apply(full.fit$BUGSoutput$sims.list$bd==s, 2, mean)
  }
  s <- 6
  full.props[s,] <- apply(full.fit$BUGSoutput$sims.list$bd==s, 2, mean)


## plot the two time series of posterior state probabilities

myPlot <- function(tBool, mySeries) {
  
    plot(tBool, -full.data$startP[tBool], type="l", ylim=c(-700, 700), yaxt="n", 
         ylab="depth + pitch", xlim=c(min(tBool), max(tBool)), bty="n", xlab=" ",
         main="DATA: depth (black), clicking (blue), pitch (grey)")
    grid()
    axis(2, at=seq(-800, 0, 200))
    points(tBool, -full.data$startP[tBool], cex=0.1, pch=full.data$RC-1, col="blue")
    lines(full.data$pitch[tBool]*700, col="grey")
    barplot(base.props[,tBool], col=1:6, border=1:6, main=paste("BASE model", mySeries), ylab="posterior P")
    barplot(full.props[,tBool], col=1:6, border=1:6, main=paste("FULL model", mySeries), ylab="posterior P")
  }

  tempI <- which(full.data$firstStep==1) # index for the first time step in each time series

  par(mfrow=c(3,1), mar=c(1, 4, 4, 1), cex=0.9, ask=F)
  myPlot(1:(tempI[2]-1), "series 1")
  
  par(mfrow=c(3,1), mar=c(1, 4, 4, 1), cex=0.9, ask=F)
  myPlot(tempI[1]:base.data$N, "series 2")