test_that("fit_pHMM returns expected structure", {
  set.seed(123)
  dat <- simulate_stream(T0 = 50, TT = 150)
  y <- dat$y
  xlabeled <- dat$x
  xlabeled[sample(seq_along(xlabeled), 50)] <- NA

  d <- ncol(y)
  res <- fit_pHMM(
    y = y,
    xlabeled = xlabeled,
    nstates = 2,
    mean_start = list(rep(0, d), rep(1, d)),
    equal_covariance = TRUE
  )

  expect_type(res, "list")
  expect_true(all(c("log_lik", "iter", "AIC", "BIC") %in% names(res)))
  expect_equal(ncol(res$log_gamma), 2) # 2 states
  expect_true(is.numeric(res$AIC))
  expect_gt(res$log_lik, -1e5)
})

test_that("fit_pHMM handles diagonal covariance", {
  dat <- simulate_stream(T0 = 100, TT = 500)
  y <- dat$y
  xlabeled <- dat$x
  d <- ncol(y)

  res <- fit_pHMM(
    y = y,
    xlabeled = xlabeled,
    nstates = 3,
    mean_start = list(rep(0, d), rep(1, d), rep(2, d)),
    covariance_start = list(diag(d), diag(d), diag(d)),
    covariance_structure = "diagonal"
  )
  expect_true(is.list(res$covariance_hat))
})


test_that("fit_pHMM_auto runs and returns correct structure", {
  set.seed(123)
  dat <- simulate_stream(T0 = 50, TT = 120)
  y <- dat$y
  xlabeled <- dat$x
  xlabeled[sample(seq_along(xlabeled), 40)] <- NA

  res <- fit_pHMM_auto(y = y, xlabeled = xlabeled, max_nstates = 3, ntry = 3)

  expect_type(res, "list")
  expect_true(all(c("log_lik", "iter", "AIC", "BIC") %in% names(res)))
  expect_true(length(res$mean_hat) >= 1)
})

test_that("fit_pHMM_auto respects max_nstates", {
  dat <- simulate_stream(T0 = 20, TT = 60)
  y <- dat$y
  xlabeled <- dat$x

  expect_warning(
    res <- fit_pHMM_auto(y = y, xlabeled = xlabeled, max_nstates = 10),
    "max_nstates cannot exceed 5"
  )
  expect_lte(length(res$mean_hat), 5)
})