ilogit <- function(x) {
  return(exp(x)/(1+exp(x)))
}

logit <- function(p) {
  return(log(p)-log(1-p))
}

getET <- function(fitvals, preddata=NULL, sp.data) {
  
  # INPUTS
  #
  # fitvals		list of posterior values for each parameter (1 posterior sample each)
  #	preddata	name vector of data to be included in the emission probability
  #	sp.data		list of data used to fit the model
  
  # OUTPUT FOR EACH TIME STEP
  #
  # logE		log probability density for emission of data for each state
  # tP.t		probability of state transition
  
  
  # SET UP PARAMETERS
  
  if(is.null(preddata)) {
    preddata <- "RC"
    if(length(fitvals$a.mean)>0) {preddata <- c(preddata, "odba")}
    if(length(fitvals$p.beta0)>0) {preddata <- c(preddata, "pitch")} 
  }
  
  N <- sp.data$N
  stateN <- dim(fitvals$T.beta0)[1]
  
  
  ## 1. TRANSITION PROBABILITIES
  
  ## Intercepts
  
  T.first <- rep(1/stateN, stateN)
  T.beta0 <- fitvals$T.beta0
  
  ## Coefficients
  
  T.beta11 <- fitvals$T.beta11
  
  T.beta11.temp <- matrix(0, stateN, stateN) 	# 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))
    }
  }
  
  if(length(fitvals$T.beta31)>0) {  			# transit from LRS to LRS
    T.beta31 <- fitvals$T.beta31
  } else {
    T.beta31 <- 0
  }
  T.beta31.temp <- matrix(0, stateN, stateN)
  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 <- fitvals$tau
  
  ## Drift
  
  drift <- fitvals$drift
  
  drift.t <- 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.beta0 <- fitvals$RC.beta0
  
  
  ## 4. ODBA
  
  if(length(fitvals$a.mean)>0) { 
    
    a.mean <- fitvals$a.mean
    a.rate <- fitvals$a.rate
    a.shape <- a.mean*a.rate
  }
  
  
  ## 5. PITCH REGRESSION
  
  if(length(fitvals$p.beta0)>0) {
    
    ## Intercept, variance
    
    p.beta0 <- fitvals$p.beta0
    p.gamma <- fitvals$p.gamma
    
    ## Coefficient for step
    
    p.beta1 <- fitvals$p.beta1
    
    p.betAtemp <- 0           		# surface
    p.betAtemp[2] <- p.beta1     	# descend
    p.betAtemp[3] <- p.beta1     	# bottom
    p.betAtemp[4] <- p.beta1     	# ascend
    p.betAtemp[5] <- 0           	# rest
    p.betAtemp[6] <- p.beta1     	# silent active
    
  }
  
  
  ## ARRAYS TO STORE PROBABILITIES 
  
  tP.t <- array(0, dim=c(N, stateN, stateN))
  t.d <- matrix(0, N, stateN) 
  RC.d <- matrix(0, N, stateN)
  pitch.d <- matrix(0, N, stateN)
  depth.d <- matrix(0, N, stateN)
  odba.d <- matrix(0, N, stateN)
  
  
  
  ## THE MODEL
  
  for(j in 1:N) { 				# loop over data
    
    t.mu <- matrix(0, stateN, stateN) 	# unscaled transition probabilities for each time step
    
    
    for(s in 1:stateN) { 			# loop over potential states
      
      
      # 1. TRANSITION PROBABILITIES
      
      for(k in 1:stateN) {
        t.mu[s, k] <- ilogit(T.beta0[s,k] + T.beta11.temp[s,k]*sp.data$startP[j]  + T.beta31.temp[s,k]*sp.data$minFromSurf[j])
      }
      
      # scale probabilities
      if(sp.data$firstStep[j]==1) {
        tP.t[j, s, 1:stateN] <- T.first[1:stateN]
      } else {
        tP.t[j, s, 1:stateN] <- t.mu[s, 1:stateN]/sum(t.mu[s, 1:stateN]) 
      }
      
      
      # 2. DEPTH
      mu.t <- sp.data$startP[j] + drift.t[s]    
      mu <- max(mu.t,0)
      depth.d[j,s] <- dnorm(sp.data$endP[j], mean=mu, sd=sqrt(tau[s]))
      
      
      # 3. CLICKING 
      
      RC.d[j,s] <- dbinom(x=sp.data$RC[j], size=1, prob=RC.beta0[s])
      
      
      # 4. ODBA
      if(length(fitvals$a.mean)>0) { 
        odba.d[j,s] <- dgamma(sp.data$odba[j], shape=a.shape[s], rate=a.rate[s])
      }
      
      
      # PITCH REGRESSION
      
      if(length(fitvals$p.beta0)>0) { 
        pitch.linpred <- p.beta0[s] + p.betAtemp[s]*sp.data$vertStep[j]
        pitch.mu <- ilogit(pitch.linpred)
        
        pitch.a <- pitch.mu*p.gamma[s]    
        pitch.b <- (1-pitch.mu)*p.gamma[s]
        
        pitch.d[j,s] <- dbeta(sp.data$pitch[j], pitch.a, pitch.b)
      }
      
    }    
  } 
  
  ## SUM LOG-LIKELIHOODS ACROSS DATA
  
  logE <- log(depth.d)
  
  if(any(preddata=="RC")) {
    logE <- logE + log(RC.d)
  }
  if(any(preddata=="odba")) {
    logE <- logE + log(odba.d)
  }
  if(any(preddata=="pitch")) {
    logE <- logE + log(pitch.d)  
  }
  
  return(list(logE, tP.t))
} 

viterbi2state <- function(trans, logE) { 
  
  # INPUTS
  #
  # 	trans		probability of state transition at each time step
  #	  logE		log emission probability for each state at each time step
  #	  stateN	list of data used to fit the model
  
  # OUTPUT FOR EACH TIME STEP
  #
  # 	vitseq		predicted Viterbi sequence
  
  
  # SET UP PARAMETERS
  
  N <- dim(logE)[1]
  stateN <- dim(trans[1,,])[1]
  logT <- log(trans) 
  delta <- matrix(0, N, stateN)		# matrix to store max probabilities of transiting to state
  delta[1,] <- (log(1/stateN) + logE[1,])  	# transition probability in 1st time step is unknown
  
  
  # PROBABILITY OF BEING IN A STATE 
  
  for (i in 2:N) { 	# loop over data
    
    for(sTo in 1:stateN) {  
      sFrom <- 1
      temp <- delta[i-1,sFrom] + logT[i, sFrom,sTo] + logE[i,sTo]
      for(sFrom in 2:stateN) {        
        temp1 <- delta[i-1,sFrom] + logT[i, sFrom,sTo] + logE[i,sTo]
        temp <- c(temp, temp1)
      }  
      delta[i,sTo] <- max(temp)
    }}
  
  
  # BACK-CALCULATE MOST LIKELY PATH
  
  pointers <- matrix(0, N-1, stateN) 
  
  for(i in 1:(N - 1)) {  # loop over data
    
    for(sTo in 1:stateN) {
      sFrom <- 1
      temp <- delta[i,sFrom] + logT[i,sFrom,sTo]
      for(sFrom in 2:stateN) {
        temp <- c(temp, delta[i,sFrom] + logT[i,sFrom,sTo])
      }
      
      pointers[i,sTo] <- which.max(temp)
    } }
  vitseq <- rep(0,N) 
  vitseq[N] <- which.max(delta[N,]) 
  
  for (i in (N-1):1) {
    vitseq[i] <- pointers[i,vitseq[i+1]] 
  }
  
  return(vitseq)
}