## Jags script 1. BASE MODEL WITH FIXED VERTICAL STEP ##############
####################################################################

# INPUT DATA FOR EACH TIME STEP
#
# 	startP		start depth at time j
# 	endP 		start depth at time j+1

# ESTIMATED FOR EACH TIME STEP
#
# 	mu		process depth
# 	bd		hidden state


## DECLARE VARIABLES

var 

T.beta0[stateN,stateN], T.beta11, bd[N], tau[stateN], drift[2], mu[N], RC.beta0[stateN];

model{


 ## FIRST DATA POINT

    bd[1] ~ dcat(T.first[1:stateN])
    mu[1] <- endP[1]


## PRIORS

    ## 1. TRANSITION PROBABILITIES

    ## Intercepts

    for (k in 1:stateN) {
    T.first[k] <- 1/stateN
    for (s in 1:stateN) {
    T.beta0.temp[k,s] ~ dbeta(T.beta0.p1[k,s],T.beta0.p2[k,s])
    T.beta0[k,s] <- logit(T.beta0.temp[k,s])
    }
    }

     ## Coefficients

    T.beta11 ~ dnorm(0,0.1)T(-15,-0.2)  # transit from any to surface

    # transit from any to surface
    for (k in 1:stateN) {     
    for (s in 1:stateN) {
    T.beta11.temp[k,s] <- T.beta11*(1-min(abs(1-s),1))
    }
    }

    ## 2. DEPTH

    ## Variances

    tau.mean[1] <- 1          # surface
    tau.var[1] <- 1
    tau.mean[2] <- 30^2       # diving
    tau.var[2] <- 25^2^2
    tau.shape <- (tau.mean^2)/tau.var 
    tau.rate <- tau.mean/tau.var

    tau[1] ~ dgamma(tau.shape[1], tau.rate[1])        # surface
    for(k in 2:stateN) {
    tau[k] ~ dgamma(tau.shape[2], tau.rate[2])} 	    # diving

    ## Drift

    drift.mean <- 78 # m/min, Watwood et al 2006
    drift.var <- 25^2
    drift.shape <- (drift.mean^2)/drift.var 
    drift.rate <- drift.mean/drift.var
    for (k in 1:2) { 
    drift[k] ~ dgamma(drift.shape, drift.rate)        # 1- descend, 2- ascend
    }
    drift.t[1] <- 0
    drift.t[2] <- drift[1]
    drift.t[3] <- 0
    drift.t[4] <- -drift[2]
    for(k in 5:stateN) {
    drift.t[k] <- 0
    }

    ## 3. CLICKING
    
    RC.p[1,1] <- 1
    RC.p[1,2] <- 10
    for(k in 2:4) {
      RC.p[k,1] <- 2
      RC.p[k,2] <- 1
    }
    for(k in 5:6) {
      RC.p[k,1] <- 1
      RC.p[k,2] <- 10
    }
    
    for(k in 1:stateN) {
      RC.beta0[k] ~ dbeta(RC.p[k,1],RC.p[k,2]) 
    }
    


## MODEL FOR DATA

      for(j in 2:N) { # loop over each data point

      # 1. TRANSITION PROBABILITIES

      for(k in 1:stateN) {
      logit(t.mu[j,k]) <- T.beta0[bd[j-1],k] + T.beta11.temp[bd[j-1],k]*startP[j] 
      }
   
      # scale probabilities
      tP.t[j, 1:stateN, 1] <- T.first[1:stateN]
      tP.t[j, 1:stateN, 2] <- t.mu[j,1:stateN]/sum(t.mu[j,1:stateN])
      
      # draw hidden state
      bd[j] ~ dcat(tP.t[j,1:stateN,firstStep[j]])
      
      ## 2. DEPTH

      for(k in 1:stateN) {
        mu.t[j,k] <- startP[j] + drift.t[bd[j]]
      }
      
      mu[j] <- max(mu.t[j,bd[j]],0)
      endP[j] ~ dnorm(mu[j], 1/tau[bd[j]])

      ## 3. CLICKING
    
      RC[j] ~ dbern(RC.beta0[bd[j]])

      }

    } # end of jags model