# ============================================================
# KRONX :: mstep.R
# M-step parameter updates for Gaussian and Student-t HMMs.
# ============================================================

# Update transition matrix A and initial vector delta from posteriors.
update_A_delta <- function(fb) {
  gamma <- fb$gamma
  xi    <- fb$xi
  K     <- ncol(gamma)

  delta  <- safe_normalize(gamma[1L, ])
  A_num  <- apply(xi, c(2L, 3L), sum)
  A_den  <- colSums(gamma[-nrow(gamma), , drop = FALSE])

  A <- matrix(0, K, K)
  for (i in seq_len(K)) {
    if (A_den[i] <= 0) {
      A[i, ] <- rep(1 / K, K)
    } else {
      A[i, ] <- A_num[i, ] / A_den[i]
    }
  }
  A <- row_normalize(A)
  list(A = A, delta = delta)
}

# Update Gaussian emission parameters (mu, sigma) via weighted ML.
update_gaussian_params <- function(x, gamma) {
  K  <- ncol(gamma)
  mu <- sigma <- numeric(K)
  for (k in seq_len(K)) {
    w       <- gamma[, k]
    mu[k]   <- weighted_mean_safe(x, w)
    sigma[k] <- sqrt(weighted_var_safe(x, w, mu[k]))
  }
  list(mu = mu, sigma = pmax(sigma, 1e-6))
}

# Weighted log-likelihood for a single Student-t state (used in nu grid search).
student_state_objective <- function(x, gamma_k, mu, sigma, nu) {
  sum(gamma_k * log_dt_locscale_vec(x, mu, sigma, nu))
}

# Update Student-t emission parameters (mu, sigma, nu) via:
#   * latent-weight updates for mu and sigma,
#   * grid search over nu_grid for degrees-of-freedom.
update_student_params <- function(x, gamma, old_params, nu_grid) {
  K  <- ncol(gamma)
  mu <- sigma <- nu <- numeric(K)

  for (k in seq_len(K)) {
    gk         <- gamma[, k]
    old_mu     <- old_params$mu[k]
    old_sigma  <- max(old_params$sigma[k], 1e-6)
    old_nu     <- old_params$nu[k]

    # Robustification weights derived from the Student-t augmentation.
    z2  <- ((x - old_mu) / old_sigma)^2
    w   <- (old_nu + 1) / (old_nu + z2)

    mu[k]      <- weighted_mean_safe(x, gk * w)
    sigma2_num <- sum(gk * w * (x - mu[k])^2)
    sigma2_den <- sum(gk)
    sigma[k]   <- sqrt(max(sigma2_num / max(sigma2_den, 1e-12), 1e-10))

    obj_vals <- sapply(nu_grid, function(nu_cand) {
      student_state_objective(x, gk, mu[k], sigma[k], nu_cand)
    })
    nu[k] <- nu_grid[which.max(obj_vals)]
  }

  list(mu = mu, sigma = pmax(sigma, 1e-6), nu = pmax(nu, 2.1))
}
