test_that("ou_params returns correct alpha and tau", {
  kappa <- 0.05; sigma_psi <- 0.10
  p <- ou_params(kappa, sigma_psi)
  expect_equal(p$alpha, exp(-kappa), tolerance = 1e-12)
  expect_equal(p$tau^2,
               sigma_psi^2 / (2*kappa) * (1 - exp(-2*kappa)),
               tolerance = 1e-12)
})

test_that("seir_simulate returns correct dimensions", {
  fixed <- list(eps=1/5, gamma=1/10, delta=1/(70*365), b=1/(70*365),
                q=0.3, N=1e6, sigma_comp=0.01)
  x0 <- c(1e6-100, 50, 50, 0, log(0.25))
  set.seed(1)
  traj <- seir_simulate(50, 0.05, 0.10, log(0.25), x0, fixed)
  expect_s3_class(traj, "data.frame")
  expect_equal(nrow(traj), 51)
  expect_named(traj, c("day","S","E","I","R","Psi","beta","obs"))
  expect_true(all(traj$S >= 0))
  expect_true(all(traj$E >= 0))
  expect_true(all(traj$obs >= 0))
})

test_that("seir_simulate uses exact OU (stationarity check)", {
  # Mean of Psi over many steps should be near mu_psi
  fixed <- list(eps=1/5, gamma=1/10, delta=1/(70*365), b=1/(70*365),
                q=0.3, N=1e6, sigma_comp=0.01)
  mu_psi <- log(0.25)
  x0 <- c(1e6-100, 50, 50, 0, mu_psi)
  set.seed(99)
  traj <- seir_simulate(2000, 0.05, 0.10, mu_psi, x0, fixed)
  expect_equal(mean(traj$Psi[200:2001]), mu_psi, tolerance = 0.1)
})

test_that("npf_seir runs and returns expected structure", {
  fixed <- list(eps=1/5, gamma=1/10, delta=1/(70*365), b=1/(70*365),
                q=0.3, N=1e6, sigma_comp=0.01)
  x0 <- c(1e6-100, 50, 50, 0, log(0.25))
  set.seed(1)
  traj <- seir_simulate(60, 0.05, 0.10, log(0.25), x0, fixed)
  obs  <- traj$obs[-1]

  fit <- npf_seir(obs, fixed, K=10, M=20, seed=42)
  expect_s3_class(fit, "npf_seir")
  expect_length(fit$Rt_mean, 60)
  expect_length(fit$beta_mean, 60)
  expect_equal(dim(fit$final_theta), c(10, 3))
  expect_equal(length(fit$final_w), 10)
  expect_length(fit$final_inner, 10)
  expect_equal(sum(fit$final_w), 1, tolerance = 1e-10)
  # First obs used for init so Rt_mean[1] should be NA
  expect_true(is.na(fit$Rt_mean[1]))
  # Non-burn-in entries should be real
  expect_true(!any(is.na(fit$Rt_mean[5:60])))
})

test_that("summary.npf_seir validates burn", {
  fixed <- list(eps=1/5, gamma=1/10, delta=1/(70*365), b=1/(70*365),
                q=0.3, N=1e6, sigma_comp=0.01)
  x0 <- c(1e6-100, 50, 50, 0, log(0.25))
  set.seed(1)
  traj <- seir_simulate(20, 0.05, 0.10, log(0.25), x0, fixed)
  fit  <- npf_seir(traj$obs[-1], fixed, K=10, M=20, seed=42)

  expect_error(summary(fit, burn = 100), "burn must be smaller")
})

test_that("plot.npf_seir supports which selector", {
  fixed <- list(eps=1/5, gamma=1/10, delta=1/(70*365), b=1/(70*365),
                q=0.3, N=1e6, sigma_comp=0.01)
  x0 <- c(1e6-100, 50, 50, 0, log(0.25))
  set.seed(1)
  traj <- seir_simulate(20, 0.05, 0.10, log(0.25), x0, fixed)
  fit  <- npf_seir(traj$obs[-1], fixed, K=10, M=20, seed=42)

  expect_invisible(plot(fit, which = "Rt"))
  expect_invisible(plot(fit, which = "beta"))
  expect_invisible(plot(fit, which = "both"))
})

test_that("npf_seir with explicit x0 fills from n=1", {
  fixed <- list(eps=1/5, gamma=1/10, delta=1/(70*365), b=1/(70*365),
                q=0.3, N=1e6, sigma_comp=0.01)
  x0 <- c(1e6-100, 50, 50, 0, log(0.25))
  set.seed(1)
  traj <- seir_simulate(30, 0.05, 0.10, log(0.25), x0, fixed)
  obs  <- traj$obs[-1]
  fit  <- npf_seir(obs, fixed, K=10, M=20, seed=42, x0=x0)
  # With explicit x0, first step is filtered
  expect_false(is.na(fit$Rt_mean[1]))
})

test_that("cori_rt returns data frame with correct structure", {
  set.seed(1)
  inc <- rpois(30, 100)
  rt  <- cori_rt(inc, tau=7)
  expect_s3_class(rt, "data.frame")
  expect_equal(nrow(rt), 30)
  expect_named(rt, c("t","Rt_mean","Rt_lo","Rt_hi"))
  # First tau-1 rows should be NaN (no complete window)
  expect_true(all(is.nan(rt$Rt_mean[1:6])))
  # Later rows should be positive
  expect_true(all(rt$Rt_mean[7:30] > 0))
})

test_that("predict.npf_seir returns sensible forecasts", {
  fixed <- list(eps=1/5, gamma=1/10, delta=1/(70*365), b=1/(70*365),
                q=0.3, N=1e6, sigma_comp=0.01)
  x0 <- c(1e6-100, 50, 50, 0, log(0.25))
  set.seed(1)
  traj <- seir_simulate(40, 0.05, 0.10, log(0.25), x0, fixed)
  fit  <- npf_seir(traj$obs[-1], fixed, K=10, M=20, seed=42)
  fc   <- predict(fit, horizon=7, n_draws=100)
  expect_length(fc$mean, 7)
  expect_true(all(fc$mean >= 0))
  expect_true(all(fc$hi >= fc$lo))
})