# ============================================================
# KRONX :: hmm_em.R
# EM fitting loop and volatility-based state ordering.
# ============================================================

#' Fit a Gaussian or Student-t HMM via EM with multiple random restarts.
#'
#' @param x        Numeric vector of observations (log returns).
#' @param K        Integer number of hidden states.
#' @param model    \code{"gaussian"} or \code{"student"}.
#' @param n_starts Number of random restarts.
#' @param max_iter Maximum EM iterations per restart.
#' @param tol      Log-likelihood convergence tolerance.
#' @param nu_grid  Candidate degrees-of-freedom for Student-t grid search.
#' @param verbose  Logical; if \code{TRUE} print iteration progress.
#'
#' @return A list with components:
#'   \describe{
#'     \item{model}{Model type string.}
#'     \item{A}{K x K transition matrix.}
#'     \item{delta}{Length-K initial state probability vector.}
#'     \item{params}{List with \code{mu}, \code{sigma}, and (Student-t) \code{nu}.}
#'     \item{gamma}{T x K posterior state probabilities.}
#'     \item{xi}{(T-1) x K x K posterior joint transition probabilities.}
#'     \item{logLik}{Scalar log-likelihood at convergence.}
#'     \item{n}{Number of observations.}
#'     \item{K}{Number of states.}
#'   }
#' @export

fit_hmm_em <- function(x, K = 3L,
                       model     = c("gaussian", "student"),
                       n_starts  = 5L,
                       max_iter  = 200L,
                       tol       = 1e-6,
                       nu_grid   = c(3:30, 40, 60, 100),
                       verbose   = TRUE) {

  # ---- Input validation ----
  if (!is.numeric(x)) {
    stop("'x' must be a numeric vector.", call. = FALSE)
  }

  x <- as.numeric(x)

  if (any(!is.finite(x))) {
    stop("'x' must contain only finite values.", call. = FALSE)
  }

  if (length(x) < 10L) {
    stop("'x' must contain at least 10 observations.", call. = FALSE)
  }

  model <- match.arg(model)

  if (!is.numeric(K) || length(K) != 1L || is.na(K) ||
      K < 1L || K != as.integer(K)) {
    stop("'K' must be a positive integer.", call. = FALSE)
  }
  K <- as.integer(K)

  if (!is.numeric(n_starts) || length(n_starts) != 1L || is.na(n_starts) ||
      n_starts < 1L || n_starts != as.integer(n_starts)) {
    stop("'n_starts' must be a positive integer.", call. = FALSE)
  }
  n_starts <- as.integer(n_starts)

  if (!is.numeric(max_iter) || length(max_iter) != 1L || is.na(max_iter) ||
      max_iter < 1L || max_iter != as.integer(max_iter)) {
    stop("'max_iter' must be a positive integer.", call. = FALSE)
  }
  max_iter <- as.integer(max_iter)

  if (!is.numeric(tol) || length(tol) != 1L || is.na(tol) ||
      !is.finite(tol) || tol <= 0) {
    stop("'tol' must be a positive finite number.", call. = FALSE)
  }

  if (!is.logical(verbose) || length(verbose) != 1L || is.na(verbose)) {
    stop("'verbose' must be TRUE or FALSE.", call. = FALSE)
  }

  if (model == "student") {
    if (!is.numeric(nu_grid) || length(nu_grid) < 1L) {
      stop("'nu_grid' must be a non-empty numeric vector.", call. = FALSE)
    }
    if (any(!is.finite(nu_grid))) {
      stop("'nu_grid' must contain only finite numeric values.", call. = FALSE)
    }
  }

  best <- NULL

  for (s in seq_len(n_starts)) {
    init   <- init_hmm_params(x, K, model)
    A      <- init$A
    delta  <- init$delta
    params <- init$params
    prev_ll <- -Inf

    for (iter in seq_len(max_iter)) {
      logB   <- emission_logdens_matrix(x, params, model)
      fb     <- forward_backward(logB, A, delta)
      ll     <- fb$logLik

      upd    <- update_A_delta(fb)
      A      <- upd$A
      delta  <- upd$delta

      if (model == "gaussian") {
        params <- update_gaussian_params(x, fb$gamma)
      } else {
        params <- update_student_params(x, fb$gamma, params, nu_grid)
      }

      if (verbose) {
        cat(sprintf("model=%s start=%d iter=%d logLik=%.10f\n",
                    model, s, iter, ll))
      }

      if (abs(ll - prev_ll) < tol) break
      prev_ll <- ll
    }

    logB <- emission_logdens_matrix(x, params, model)
    fb   <- forward_backward(logB, A, delta)
    ll   <- fb$logLik

    fit <- list(
      model  = model,
      A      = A,
      delta  = delta,
      params = params,
      gamma  = fb$gamma,
      xi     = fb$xi,
      logLik = ll,
      n      = length(x),
      K      = K
    )

    if (is.null(best) || fit$logLik > best$logLik) best <- fit
  }

  best
}

# Reorder states from lowest to highest volatility.
# Gaussian: ordered by sigma.
# Student-t: ordered by conditional variance = [nu/(nu-2)] * sigma^2.
reorder_fit_by_scale <- function(fit) {
  if (fit$model == "gaussian") {
    ord <- order(fit$params$sigma)
  } else {
    cond_var <- (fit$params$nu / (fit$params$nu - 2)) * fit$params$sigma^2
    ord <- order(cond_var)
  }

  fit$A              <- fit$A[ord, ord, drop = FALSE]
  fit$delta          <- fit$delta[ord]
  fit$gamma          <- fit$gamma[, ord, drop = FALSE]
  fit$params$mu      <- fit$params$mu[ord]
  fit$params$sigma   <- fit$params$sigma[ord]
  if (!is.null(fit$params$nu)) fit$params$nu <- fit$params$nu[ord]
  fit
}
