# ============================================================
# KRONX :: viterbi.R
# Viterbi algorithm for most-likely hidden-state path decoding.
# ============================================================

#' Decode the most likely hidden-state path via the Viterbi algorithm.
#'
#' @param x   Numeric vector of observations (log returns).
#' @param fit A fitted HMM object returned by \code{\link{fit_hmm_em}}.
#'
#' @return An integer vector of length \code{length(x)} with state labels
#'   (1-indexed).
#' @export
viterbi_decode <- function(x, fit) {
  logB      <- emission_logdens_matrix(x, fit$params, fit$model)
  A         <- fit$A
  delta     <- fit$delta
  Tn        <- nrow(logB)
  K         <- ncol(logB)

  delta_log <- matrix(-Inf, Tn, K)
  psi       <- matrix(0L,   Tn, K)

  delta_log[1L, ] <- log(delta) + logB[1L, ]

  for (t in 2:Tn) {
    for (j in seq_len(K)) {
      vals          <- delta_log[t - 1L, ] + log(A[, j])
      psi[t, j]     <- which.max(vals)
      delta_log[t, j] <- max(vals) + logB[t, j]
    }
  }

  path      <- integer(Tn)
  path[Tn]  <- which.max(delta_log[Tn, ])
  if (Tn >= 2L) {
    for (t in (Tn - 1L):1L) {
      path[t] <- psi[t + 1L, path[t + 1L]]
    }
  }
  path
}
