# functions to extract parameters, 
# get the Viterbi decoded states, 
# t.p.ms,
# stationary distributions 
# and confidence intervals from a fitted model

# get the Viterbi decoded states
getVit <- function(mod, data, N = 2, cycle = 24){
  newdata_trigon <- todtrans(x = 1:cycle)
  
  Gamma <- array(NA, dim = c(N,N,cycle))
  probs12 <- predict.gam(mod$mods$gammod1, newdata_trigon, type = "response")
  probs21 <- predict.gam(mod$mods$gammod2, newdata_trigon, type = "response")
  Gamma[1,1,] <- (1-probs12)
  Gamma[1,2,] <- probs12
  Gamma[2,1,] <- probs21
  Gamma[2,2,] <- (1-probs21)
  
  vit_states <- Viterbi(tod = data$tod, step = data$step, angle = data$angle, 
                               trackID = data$ID, ID = data$animalID, N = N, 
                               mod$stepShape, mod$stepRate,
                               mod$anglemean, mod$anglecon, Gamma,
                               mod$deltas)
  return(vit_states)
}

# get the Gamma array
getGamma <- function(mod, cycle = 24, N = 2, no_tp = 240){
  var_x <- seq(cycle/no_tp, cycle, length = no_tp)
  newdata_trigon <- todtrans(x=var_x)
  
  Gamma <- array(NA, dim = c(N,N,no_tp))
  # get Gamma and delta for trigon -------------------------------------------------------------
  if (length(mod$mods$gammod1$coefficients)==1)
  {probs12 <- fitted(mod$mods$gammod1)[1:cycle]
  probs21 <- fitted(mod$mods$gammod2)[1:cycle]}
  else 
  {probs12 <- predict.gam(mod$mods$gammod1, newdata_trigon, type = "response")
   probs21 <- predict.gam(mod$mods$gammod2, newdata_trigon, type = "response")
  }
  # using the results of the GLMMs, define the t.p.m. (for each hour of the day)
  Gamma[1,1,] <- (1-probs12)
  Gamma[1,2,] <- probs12
  Gamma[2,1,] <- probs21
  Gamma[2,2,] <- (1-probs21)
  
  return(Gamma)
}

# get the state probabilities implied by periodic stationarity
getDelta <- function(Gamma, cycle = 24, N = 2, no_tp = 240){

  deltas <- matrix(NA, nrow = no_tp, ncol = N)
  index <- c(1:no_tp,1:no_tp) # index makes it easier to choose the correct Gamma in the matrix multiplications
  
  # get Gamma_star for periodic stationary distributions
  Gamma_star <- Gamma # Gamma_star will be the multiplication of the next 24 hourly(or 48 half-hourly) Gamma matrices
  for (t in 1:no_tp){
    # loopi determines which gamma matrices will be multiplied for which t
    # e.g. for hourly data, only 24 instead of 240 matrices multiplied
    loopi <- index[seq(t+no_tp/cycle,no_tp+t-no_tp/cycle, by = no_tp/cycle)]
    for (i in loopi){
      Gamma_star[, , t] <- Gamma_star[, , t]%*%Gamma[,,i]
    }
  }
  # get periodic stationary distributions
  for (t in 1:no_tp) {
    deltas[t, ] <- solve(t(diag(2) - Gamma_star[, , t] + 1), c(1, 1))
  }
  return(deltas)
}

# get confidence intervals for the t.p.m.
getConfint <- function(gammod, cycle = 24, N = 2, no_tp = 240){
  var_x <- seq(cycle/no_tp, cycle, length = no_tp)
  newdata_trigon <- todtrans(x=var_x)
  
  # confidence intervals for gammas:
  probsConf <- matrix(NA, nrow = no_tp, ncol = 3)
  probs <- predict.gam(gammod, newdata_trigon, se = TRUE)
  probsConf[,1] <- plogis(as.vector(probs$fit) + qnorm(0.025)*as.vector(probs$se.fit))
  probsConf[,2] <- plogis(as.vector(probs$fit))
  probsConf[,3] <- plogis(as.vector(probs$fit) + qnorm(0.975)*as.vector(probs$se.fit))
  
  return(probsConf)
}  

# get confidence intervals for the stationary state distributions
getDeltaConfint <- function(mod, N = 2, cycle = 48, no_tp = 240, n = 1000){
  var_x <- seq(cycle/no_tp, cycle, length = no_tp)
  newdata_trigon <- todtrans(x=var_x)

  Xp1 <- predict(mod$mods$gammod1, newdata = newdata_trigon, type="lpmatrix") # some sort of design matrix
  beta1 <- coef(mod$mods$gammod1) ## posterior mean of coefs
  Vb1   <- vcov(mod$mods$gammod1) ## posterior cov of coefs
  Xp2 <- predict(mod$mods$gammod2, newdata = newdata_trigon, type="lpmatrix") # same for mod2
  beta2 <- coef(mod$mods$gammod2)
  Vb2   <- vcov(mod$mods$gammod2) 
  mrand1 <- mvrnorm(n, beta1, Vb1) ## simulate n rep coef vectors from posterior
  mrand2 <- mvrnorm(n, beta2, Vb2) ## same for mod2
  
  deltas <- matrix(NA, nrow = n, ncol = no_tp)
  Gamma <- array(NA, dim = c(N,N,no_tp))
  for (i in 1:n) { 
    probs12   <- plogis(Xp1 %*% mrand1[i, ])
    probs21   <- plogis(Xp2 %*% mrand2[i, ])
    Gamma[1,1,] <- (1-probs12)
    Gamma[1,2,] <- probs12
    Gamma[2,1,] <- probs21
    Gamma[2,2,] <- (1-probs21)
    delta <- getDelta(Gamma, no_tp = no_tp)
    deltas[i, ] <- delta[,2]
    if (i%%(n/10)==0) print(paste(i, "of", n, "Monte Carlo draws done!"))
  }
  confints <- apply(deltas, 2, quantile, probs = c(0.025, 0.975))
  return(confints)
}

# get the Viterbi decoded states for the homogeneous HMM
getVit0 <- function(mod, data, N = 2, cycle = 24){

  Gamma <- array(NA, dim = c(N,N,cycle))
  probs12 <- fitted(mod$mods$gammod1)[1:cycle]
  probs21 <- fitted(mod$mods$gammod2)[1:cycle]
  Gamma[1,1,] <- (1-probs12)
  Gamma[1,2,] <- probs12
  Gamma[2,1,] <- probs21
  Gamma[2,2,] <- (1-probs21)
  
  vit_states <- Viterbi(tod = data$tod, step = data$step, angle = data$angle, 
                        trackID = data$ID, ID = data$animalID, N = N, 
                        mod$stepShape, mod$stepRate,
                        mod$anglemean, mod$anglecon, Gamma,
                        mod$deltas)
  return(vit_states)
}