## Jags script 2. FULL MODEL WITH STEP LENGTH ALLOWED TO INCREASE WITH SQRT(DEPTH)
##################################################################################

# INPUT DATA FOR EACH TIME STEP
#
#   startP		start depth at time j
# 	endP 		start depth at time j+1
#	  RC		presence/absence of echolocation
#	  odba		overall dynamic body acceleration
#	  pitch		absolute value of mean pitch angle, scaled to [0,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], RC.beta0[stateN], mu[N], T.beta31, a.mean[stateN], a.rate[stateN], p.beta0[6], p.gamma[6], p.beta1, tau.beta0, temp.sd[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
  T.beta31 ~ dnorm(0,0.1)T(,1.0E-6)   # transit from LRS to LRS
  
  # 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))
    }
  }
  # transit from LRS to LRS
  for (k in 1:stateN) {     
    for (s in 1:stateN) {
      T.beta31.temp[k,s] <- T.beta31*(1-min(abs(k-s),1))*(1-min(abs(3-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])
  temp.sd[1] <- sqrt(tau[1])
  
  for(k in 2:stateN) {
    tau[k] ~ dgamma(tau.shape[2], tau.rate[2])
    temp.sd[k] <- sqrt(tau[k])
  } # all other taus
  
  tau.beta0 ~ dgamma(3, 2)
  
  ## 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]) 
  }
  
  ## 4. ODBA 
  
  for(k in 1:4) {
    a.mean.p[k,1] <- 5
    a.mean.p[k,2] <- 0.2
    a.rate.p[k,1] <- 3
    a.rate.p[k,2] <- 0.5
  }
  a.mean.p[5,1] <- 1
  a.mean.p[5,2] <- 1
  a.rate.p[5,1] <- 1
  a.rate.p[5,2] <- 2
  a.mean.p[6,1] <- 5
  a.mean.p[6,2] <- 0.2
  a.rate.p[6,1] <- 3
  a.rate.p[6,2] <- 0.5
  
  for(k in 1:stateN){
    a.mean[k] ~ dgamma(a.mean.p[k,1],a.mean.p[k,2])
    a.rate[k] ~ dgamma(a.rate.p[k,1],a.rate.p[k,2])
    a.shape[k] <- a.mean[k]*a.rate[k]
  }
  
  ## 5. PITCH REGRESSION
  
  ## Intercept, variance
  
  for(k in 1:3){ # 1: surface, 2: descend/bottom/ascend/other, 3: resting
    p.beta0.temp[k] ~ dnorm(0,0.1)
    p.gamma.temp[k] ~ dgamma(1,0.1)T(1.0E-6,)
  }
  p.beta0[1] <- p.beta0.temp[1] # surface
  p.gamma[1] <- p.gamma.temp[1]
  for(k in 2:4) {
    p.beta0[k] <- p.beta0.temp[2] # descend, LRS, ascend
    p.gamma[k] <- p.gamma.temp[2]
  }
  p.beta0[5] <- p.beta0.temp[3] # resting
  p.gamma[5] <- p.gamma.temp[3]
  p.beta0[6] <- p.beta0.temp[2] # other
  p.gamma[6] <- p.gamma.temp[2]
  
  ## Coefficient for step
  
  p.beta1 ~ dnorm(0,1.0E-6)
  p.beta1.temp[1] <- 0           	# surface
  p.beta1.temp[2] <- p.beta1     # descend
  p.beta1.temp[3] <- p.beta1     # LRS
  p.beta1.temp[4] <- p.beta1     # ascend
  p.beta1.temp[5] <- 0           	# rest
  p.beta1.temp[6] <- p.beta1     # other
  
  
  
  ## 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] 
      + T.beta31.temp[bd[j-1],k]*minFromSurf[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
    
    mu.t[j,1] <- startP[j]
    mu.t[j,2] <- startP[j] + drift[1] + tau.beta0*sqrt(startP[j])
    mu.t[j,3] <- startP[j]
    mu.t[j,4] <- startP[j] - drift[2] - tau.beta0*sqrt(startP[j])
    for(k in 5:stateN) {
      mu.t[j,k] <- startP[j]}
    
    sd.t[j,1] <- temp.sd[1]
    sd.t[j,2] <- temp.sd[2]
    sd.t[j,3] <- temp.sd[3] + tau.beta0*sqrt(startP[j])
    sd.t[j,4] <- temp.sd[4]
    for(k in 5:stateN) {
      sd.t[j,k] <- temp.sd[k]}
    
    mu[j] <- max(mu.t[j,bd[j]],0)
    endP[j] ~ dnorm(mu[j], 1/(sd.t[j,bd[j]])^2)  
    
    ## 3. CLICKING
    
    RC[j] ~ dbern(RC.beta0[bd[j]])
    
    ## 4. ODBA
    
    odba[j] ~ dgamma(a.shape[bd[j]], a.rate[bd[j]])
    
    ## 5. PITCH REGRESSION
    
    pitch.linpred[j] <- p.beta0[bd[j]] + p.beta1.temp[bd[j]]*vertStep[j]
    logit(pitch.mu[j]) <- pitch.linpred[j] 
    
    pitch.a[j] <- pitch.mu[j]*p.gamma[bd[j]]    
    pitch.b[j] <- (1-pitch.mu[j])*p.gamma[bd[j]]
    pitch[j] ~ dbeta(pitch.a[j], pitch.b[j])
    
  } # loop over each data point
  
} # end of model