# ============================================================
# KRONX :: forward_backward.R
# Scaled forward-backward algorithm for HMM inference.
# ============================================================

# Run the scaled forward-backward algorithm.
#
# Arguments
#   logB  : T x K matrix of emission log-densities
#   A     : K x K transition matrix
#   delta : length-K initial state probability vector
#
# Returns a list with:
#   alpha  - T x K scaled forward probabilities
#   beta   - T x K scaled backward probabilities
#   gamma  - T x K posterior state probabilities P(S_t = k | data)
#   xi     - (T-1) x K x K posterior joint transition probabilities
#   logLik - scalar scaled log-likelihood
forward_backward <- function(logB, A, delta) {
  Tn <- nrow(logB)
  K  <- ncol(logB)

  B     <- exp(logB)
  alpha <- matrix(0, Tn, K)
  beta  <- matrix(0, Tn, K)
  scale <- numeric(Tn)

  # ---- Forward pass ----
  alpha[1L, ] <- delta * B[1L, ]
  scale[1L]   <- sum(alpha[1L, ])
  alpha[1L, ] <- alpha[1L, ] / scale[1L]

  for (t in 2:Tn) {
    alpha[t, ] <- (alpha[t - 1L, ] %*% A) * B[t, ]
    scale[t]   <- sum(alpha[t, ])
    alpha[t, ] <- alpha[t, ] / scale[t]
  }

  # ---- Backward pass ----
  beta[Tn, ] <- rep(1, K)

  if (Tn >= 2L) {
    for (t in (Tn - 1L):1L) {
      beta[t, ] <- A %*% (B[t + 1L, ] * beta[t + 1L, ])
      beta[t, ] <- beta[t, ] / scale[t + 1L]
    }
  }

  # ---- Posteriors ----
  gamma      <- alpha * beta
  gamma      <- gamma / rowSums(gamma)

  xi <- array(0, dim = c(Tn - 1L, K, K))
  if (Tn >= 2L) {
    for (t in seq_len(Tn - 1L)) {
      M        <- (alpha[t, ] * A) * rep(B[t + 1L, ] * beta[t + 1L, ], each = K)
      denom    <- sum(M)
      xi[t, , ] <- M / denom
    }
  }

  list(
    alpha  = alpha,
    beta   = beta,
    gamma  = gamma,
    xi     = xi,
    logLik = sum(log(scale))
  )
}
