## 'augment' ---------------------------------------------------------------

test_that("'augment' works with Poisson, disp - has data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod_fitted <- fit(mod)
    ans <- augment(mod_fitted, quiet = TRUE)
    expect_true(is.data.frame(ans))
    expect_identical(names(ans),
                     c(names(data), c(".observed", ".fitted", ".expected")))
})

test_that("'augment' calculates fitted in cells with missing outcome or offset -  Poisson", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2010, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    data$popn[1] <- NA
    data$deaths[2] <- NA
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    suppressWarnings(ans_unfit <- augment(mod, quiet = TRUE)) ## high value for lambda
    mod_fitted <- fit(mod)
    ans_fit <- augment(mod_fitted, quiet = TRUE)
    expect_false(any(rvec::draws_any(is.na(ans_fit$.fitted))))
    expect_true(".deaths" %in% names(ans_fit))
    expect_true(all(rvec::draws_all(ans_fit$.deaths[-2] == ans_fit$deaths[-2])))
    expect_equal(rvec::draws_mean(ans_fit$.deaths[2]),
                 rvec::draws_mean(ans_fit$.fitted[2] * ans_fit$popn[2]),
                 tolerance = 0.02)
    expect_false(".deaths" %in% names(ans_unfit))
})

test_that("'augment' works with Poisson, disp - no data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    aug_notfitted <- suppressWarnings(augment(mod, quiet = TRUE)) ## high values for lambda
    mod_fitted <- fit(mod)
    aug_fitted <- augment(mod_fitted, quiet = TRUE)
    expect_identical(names(aug_fitted), names(aug_notfitted))
})

test_that("'augment' works with binomial, no disp - has data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.5)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_disp(mod, mean = 0)
    mod_fitted <- fit(mod)
    ans <- augment(mod_fitted, quiet = TRUE)
    expect_true(is.data.frame(ans))
    expect_identical(names(ans),
                     c(names(data), c(".observed", ".fitted")))
})

test_that("'augment' works with normal - with data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$deaths <- rpois(n = nrow(data), lambda = 100)
    data$deaths[1] <- NA
    formula <- deaths ~ age + sex + time
    mod <- mod_norm(formula = formula,
                    data = data,
                    weights = 1)
    mod_fitted <- fit(mod)
    ans <- augment(mod_fitted, quiet = TRUE)
    expect_true(is.data.frame(ans))
    expect_identical(names(ans),
                     c(names(data), ".deaths", ".fitted"))
})

test_that("'augment' works with normal - no data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$deaths <- rpois(n = nrow(data), lambda = 100)
    formula <- deaths ~ age + sex + time
    mod <- mod_norm(formula = formula,
                    data = data,
                    weights = 1)
    ans <- augment(mod, quiet = TRUE)
    expect_true(is.data.frame(ans))
    expect_identical(names(ans),
                     c(names(data), ".fitted"))
})

test_that("'augment' gives same answer when run twice - with data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.5)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod_fitted <- fit(mod)
    ans1 <- augment(mod_fitted, quiet = TRUE)
    ans2 <- augment(mod_fitted, quiet = TRUE)
    expect_identical(ans1, ans2)
})

test_that("'augment' gives same answer when run twice - no data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.5)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    ans1 <- suppressWarnings(augment(mod, quiet = TRUE)) ## high values for lambda
    ans2 <- suppressWarnings(augment(mod, quiet = TRUE)) ## high values for lambda
    expect_identical(ans1, ans2)
})

test_that("'augment' gives message when used unfitted and quiet is FALSE", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                      KEEP.OUT.ATTRS = FALSE)
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_prior(age ~ NFix(sd = 0.1)) |>
    set_prior(time ~ NFix(sd = 0.1)) |>
    set_prior(sex ~ NFix(sd = 0.1))
  suppressMessages(
    expect_message(augment(mod),
                   "Model not fitted, so drawing values straight from prior distribution.")
    )
})


## 'components' ---------------------------------------------------------------

test_that("'components' works with no disp", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_disp(mod, mean = 0)
    mod_fitted <- fit(mod)
    ans <- components(mod_fitted)
    expect_true(is.data.frame(ans))
    expect_identical(unique(ans$component), c("effect", "hyper"))
})

test_that("'components' works with no data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    comp_nodata <- components(mod, quiet = TRUE)
    mod_data <- fit(mod)
    comp_data <- components(mod_data)
    comp_merge <- merge(comp_nodata, comp_data, by = c("term", "component", "level"))
    expect_identical(nrow(comp_merge), nrow(comp_data))
})

test_that("'components' gives same answer when run twice - with data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_disp(mod, mean = 0)
    mod_fitted <- fit(mod)
    ans1 <- components(mod_fitted, quiet = TRUE)
    ans2 <- components(mod_fitted, quiet = TRUE)
    expect_identical(ans1, ans2)
})

test_that("'components' gives same answer when run twice - no data", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    ans1 <- components(mod, quiet = TRUE)
    ans2 <- components(mod, quiet = TRUE)
    expect_identical(ans1, ans2)
})

test_that("'components' gives message when used unfitted and quiet is FALSE", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    expect_message(components(mod),
                   "Model not fitted, so values drawn straight from prior distribution.")
})

test_that("'disp' estimates not affected by weights in normal model", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$income <- rnorm(n = nrow(data), mean = 10, sd = 3)
    data$wt1 <- runif(n = nrow(data), max = 5)
    data$wt2 <- 20 * data$wt1
    formula <- income ~ age + sex + time
    mod1 <- mod_norm(formula = formula,
                    data = data,
                    weights = wt1)
    mod1_fitted <- fit(mod1)
    comp1 <- components(mod1_fitted, quiet = TRUE)
    disp1 <- comp1$.fitted[comp1$term == "disp"]
    mod2 <- mod_norm(formula = formula,
                    data = data,
                    weights = wt2)
    mod2_fitted <- fit(mod2)
    comp2 <- components(mod1_fitted, quiet = TRUE)
    disp2 <- comp2$.fitted[comp1$term == "disp"]
    expect_equal(rvec::draws_mean(disp1), rvec::draws_mean(disp2))
})

test_that("'components' gives expected message when 'original_scale' is FALSE and model is normal", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$income <- rnorm(n = nrow(data), mean = 10, sd = 3)
    data$wt <- runif(n = nrow(data), max = 5)
    formula <- income ~ age + sex + time
    mod <- mod_norm(formula = formula,
                    data = data,
                    weights = wt) |>
      fit()
    expect_message(components(mod),
                   "Values for `.fitted`")
})


## 'computations' -------------------------------------------------------------

test_that("'computations' returns NULL (with warning) if applied to unfitted model", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    expect_message({x <- computations(mod)},
                   "Model not fitted..")
    expect_true(is.null(x))
})

test_that("'computations' returns tibble if applied to fitted model", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn) |>
      fit()
    expect_true(tibble::is_tibble(computations(mod)))
})


## 'draw_vals_augment_fitted' -------------------------------------------------

test_that("'draw_vals_augment_fitted' works with Poisson, has disp", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod_fitted <- fit(mod)
    ans <- draw_vals_augment_fitted(mod_fitted, quiet = TRUE)
    expect_true(is.data.frame(ans))
    expect_identical(names(ans),
                     c(names(data), c(".observed", ".fitted", ".expected")))
})

test_that("'draw_vals_augment_fitted' works with Poisson, has disp, rr3", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2002, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 3) * 3
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn) |>
      set_confidential_rr3()
    mod_fitted <- fit(mod)
    expect_message(
      ans <- draw_vals_augment_fitted(mod_fitted, quiet = FALSE),
      "Adding variable `\\.deaths` with true values for `deaths`."
    )
    expect_true(is.data.frame(ans))
    expect_identical(names(ans),
                     c(names(data), c(".deaths",
                                      ".observed",
                                      ".fitted",
                                      ".expected")))
})

test_that("'draw_vals_augment_fitted' calculates fitted in cells with missing outcome or offset -  Poisson", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    data$popn[1] <- NA
    data$deaths[2] <- NA
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod_fitted <- fit(mod)
    ans <- draw_vals_augment_fitted(mod_fitted, quiet = TRUE)
    expect_false(any(rvec::draws_any(is.na(ans$.fitted))))
})

test_that("'draw_vals_augment_fitted' works with binomial, no disp", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.5)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_disp(mod, mean = 0)
    mod_fitted <- fit(mod)
    ans <- draw_vals_augment_fitted(mod_fitted, quiet = TRUE)
    expect_true(is.data.frame(ans))
    expect_identical(names(ans),
                     c(names(data), c(".observed", ".fitted")))
})

test_that("'draw_vals_augment_fitted' works with normal", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                        KEEP.OUT.ATTRS = FALSE)
    data$deaths <- rpois(n = nrow(data), lambda = 100)
    data$deaths[1] <- NA
    data$wt <- runif(n = nrow(data), max = 1000)
    formula <- deaths ~ age + sex + time
    mod <- mod_norm(formula = formula,
                    data = data,
                    weights = wt)
    mod_fitted <- fit(mod)
    expect_message(
      ans <- draw_vals_augment_fitted(mod_fitted, quiet = FALSE),
      "Adding variable `.deaths` with true values for `deaths`."
    )
    expect_true(is.data.frame(ans))
    expect_setequal(names(ans),
                    c(names(data), ".deaths", ".fitted"))
    expect_equal(rvec::draws_mean(mean(ans$.fitted)), mean(data$deaths, na.rm = TRUE), tolerance = 0.1)
    expect_equal(rvec::draws_mean(mean(ans$.deaths)), mean(data$deaths, na.rm = TRUE), tolerance = 0.1)
})


test_that("'draw_vals_augment_fitted' gives message if imputes exposure", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                      KEEP.OUT.ATTRS = FALSE)
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_disp(mean = 0) |>
    set_datamod_exposure(cv = 0.01)
  mod_fitted <- fit(mod)
  expect_message(draw_vals_augment_fitted(mod_fitted, quiet = FALSE),
                 "Adding variable `.popn` with true values for `popn`.")
})



## 'draw_vals_augment_unfitted' -----------------------------------------------

test_that("'draw_vals_augment_unfitted' works with 'bage_mod_pois' - has disp, no exposure", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 20)
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = 1)
  mod <- set_n_draw(mod, 10)
  ans_obtained <- draw_vals_augment_unfitted(mod, quiet = TRUE)
  set.seed(mod$seed_augment)
  components <- draw_vals_components_unfitted(mod, n_sim = 10)
  disp <- components$.fitted[components$component == "disp"]
  expected <- exp(make_linpred_from_components(mod = mod,
                                               components = components,
                                               data = mod$data,
                                               dimnames_terms = mod$dimnames_terms))
  fitted <- draw_fitted(nm_distn = "pois",
                        expected = expected,
                        disp = disp)
  outcome <- draw_outcome_true(nm_distn = "pois",
                               fitted = fitted,
                               disp = disp,
                               offset = mod$offset)
  ans_expected <- tibble::as_tibble(data)
  ans_expected$deaths <- outcome
  ans_expected$.fitted <- fitted
  ans_expected$.expected <- expected
  expect_equal(ans_obtained, ans_expected)
  expect_identical(names(augment(fit(mod, quiet = TRUE))), names(ans_obtained))
})

test_that("'draw_vals_augment_unfitted' works with 'bage_mod_pois' - no disp, has exposure", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 20)
  data$popn <- rpois(n = nrow(data), lambda = 30)
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_disp(mod, mean = 0)
  mod <- set_n_draw(mod, 10)
  ans_obtained <- draw_vals_augment_unfitted(mod, quiet = TRUE)
  set.seed(mod$seed_augment)
  components <- draw_vals_components_unfitted(mod = mod,
                                              n_sim = n_sim)
  fitted <- exp(make_linpred_from_components(mod = mod,
                                             components = components,
                                             data = mod$data,
                                             dimnames_term = mod$dimnames_terms))
  outcome <- draw_outcome_true(nm_distn = "pois",
                               fitted = fitted,
                               disp = NULL,
                               offset = mod$offset)
  ans_expected <- tibble::as_tibble(data)
  ans_expected$deaths <- outcome
  ans_expected$.observed <- outcome / data$popn
  ans_expected$.fitted <- fitted
  expect_equal(ans_obtained, ans_expected)
  expect_identical(names(augment(fit(mod), quiet = TRUE)), names(ans_obtained))
})

test_that("'draw_vals_augment_unfitted' works with 'bage_mod_pois' - has disp, has exposure, has NA", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 20)
  data$popn <- rpois(n = nrow(data), lambda = 30)
  data$popn[1] <- NA
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_disp(mod, mean = 0)
  mod <- set_n_draw(mod, 10)
  ans_obtained <- draw_vals_augment_unfitted(mod, quiet = TRUE)
  set.seed(mod$seed_augment)
  components <- draw_vals_components_unfitted(mod = mod,
                                              n_sim = n_sim)
  fitted <- exp(make_linpred_from_components(mod = mod,
                                             components = components,
                                             data = mod$data,
                                             dimnames_term = mod$dimnames_terms))
  outcome <- draw_outcome_true(nm_distn = "pois",
                               fitted = fitted,
                               disp = get_disp(mod),
                               offset = mod$offset)
  ans_expected <- tibble::as_tibble(data)
  ans_expected$deaths <- outcome
  ans_expected$.observed <- outcome / data$popn
  ans_expected$.fitted <- fitted
  expect_equal(ans_obtained, ans_expected)
  expect_identical(names(augment(fit(mod), quiet = TRUE)), names(ans_obtained))
})

test_that("'draw_vals_augment_unfitted' works with 'bage_mod_pois' - has disp, has exposure, has RR3", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                      KEEP.OUT.ATTRS = FALSE)
  data$deaths <- rpois(n = nrow(data), lambda = 3) * 3
  data$popn <- rpois(n = nrow(data), lambda = 30)
  data$popn[1] <- NA
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_disp(mod, mean = 0)
  mod <- set_n_draw(mod, 10)
  mod <- set_confidential_rr3(mod)
  suppressMessages(
    expect_message(
      ans_obtained <- draw_vals_augment_unfitted(mod, quiet = FALSE),
      "Overwriting existing values for `deaths`."
    )
  )
  set.seed(mod$seed_augment)
  components <- draw_vals_components_unfitted(mod = mod,
                                              n_sim = n_sim)
  fitted <- exp(make_linpred_from_components(mod = mod,
                                             components = components,
                                             data = mod$data,
                                             dimnames_term = mod$dimnames_terms))
  outcome_true <- draw_outcome_true(nm_distn = "pois",
                               fitted = fitted,
                               disp = get_disp(mod),
                               offset = mod$offset)
  outcome_obs <- poputils::rr3(outcome_true)
  ans_expected <- tibble::as_tibble(data)
  ans_expected$.deaths <- outcome_true
  ans_expected$deaths <- outcome_obs
  ans_expected <- ans_expected[c("age", "time", "sex", "deaths",
                                 ".deaths", "popn")]
  ans_expected$.observed <- outcome_obs / data$popn
  ans_expected$.fitted <- fitted
  expect_equal(ans_obtained, ans_expected)
  expect_identical(names(augment(fit(mod), quiet = TRUE)), names(ans_obtained))
})

test_that("'draw_vals_augment_unfitted' works with 'bage_mod_pois' - no disp, has exposure, has datamod_noise", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                      KEEP.OUT.ATTRS = FALSE)
  data$deaths <- rpois(n = nrow(data), lambda = 3) * 3
  data$popn <- rpois(n = nrow(data), lambda = 30)
  data$popn[1] <- NA
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_disp(mod, mean = 0)
  mod <- set_n_draw(mod, 10)
  mod <- set_datamod_noise(mod, sd = 0.2)
  suppressMessages(
    expect_message(
      ans_obtained <- draw_vals_augment_unfitted(mod, quiet = FALSE),
      "Overwriting existing values for `deaths`."
    )
  )
  set.seed(mod$seed_augment)
  components <- draw_vals_components_unfitted(mod = mod,
                                              n_sim = n_sim)
  fitted <- exp(make_linpred_from_components(mod = mod,
                                             components = components,
                                             data = mod$data,
                                             dimnames_term = mod$dimnames_terms))
  outcome_true <- draw_outcome_true(nm_distn = "pois",
                                    fitted = fitted,
                                    disp = get_disp(mod),
                                    offset = mod$offset)
  outcome_obs <- draw_outcome_obs_given_true(datamod = mod$datamod,
                                             components = components,
                                             outcome_true = outcome_true,
                                             offset = mod$offset,
                                             fitted = fitted)
  ans_expected <- tibble::as_tibble(data)
  ans_expected$.deaths <- outcome_true
  ans_expected$deaths <- outcome_obs
  ans_expected <- ans_expected[c("age", "time", "sex", "deaths",
                                 ".deaths", "popn")]
  ans_expected$.observed <- outcome_obs / data$popn
  ans_expected$.fitted <- fitted
  expect_equal(ans_obtained, ans_expected)
  expect_identical(names(augment(fit(mod), quiet = TRUE)), names(ans_obtained))
})

test_that("'draw_vals_augment_unfitted' works with 'bage_mod_pois' - has datamod_exposure", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                      KEEP.OUT.ATTRS = FALSE)
  data$deaths <- rpois(n = nrow(data), lambda = 3) * 3
  data$popn <- rpois(n = nrow(data), lambda = 30)
  data$popn[1] <- NA
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_disp(mod, mean = 0)
  mod <- set_n_draw(mod, 10)
  mod <- set_datamod_exposure(mod, cv = 0.2)
  suppressMessages(
    expect_message(
      ans_obtained <- draw_vals_augment_unfitted(mod, quiet = FALSE),
      "Overwriting existing values for `deaths`."
    )
  )
  set.seed(mod$seed_augment)
  components <- draw_vals_components_unfitted(mod = mod,
                                              n_sim = n_sim)
  fitted <- exp(make_linpred_from_components(mod = mod,
                                             components = components,
                                             data = mod$data,
                                             dimnames_term = mod$dimnames_terms))
  outcome_true <- draw_outcome_true(nm_distn = "pois",
                                    fitted = fitted,
                                    disp = get_disp(mod),
                                    offset = mod$offset)
  offset_obs <- draw_offset_obs_given_true(datamod = mod$datamod,
                                           components = components,
                                           offset_true = mod$offset)
  ans_expected <- tibble::as_tibble(data)
  ans_expected$deaths <- outcome_true
  ans_expected$popn <- offset_obs
  ans_expected$.popn <- mod$offset
  ans_expected <- ans_expected[c("age", "time", "sex", "deaths",
                                 "popn", ".popn")]
  ans_expected$.observed <- outcome_true / offset_obs
  ans_expected$.fitted <- fitted
  expect_equal(ans_obtained, ans_expected)
  expect_identical(names(augment(fit(mod), quiet = TRUE)), names(ans_obtained))
})

test_that("'draw_vals_augment_unfitted' works with bage_mod_norm", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                      KEEP.OUT.ATTRS = FALSE)
  data$income <- rnorm(n = nrow(data), mean = 100, sd = 5)
  data$wt <- runif(n = nrow(data), max = 1000)
  formula <- income ~ age + sex + time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  expect_message(
    ans <- draw_vals_augment_unfitted(mod, quiet = FALSE),
    "Overwriting existing values for `income`."
  )
  expect_true(is.data.frame(ans))
  expect_identical(names(ans), c(names(data), ".fitted"))
  expect_true(abs(rvec::draws_mean(mean(ans$.fitted)) - 100) < 2)
  expect_true(abs(rvec::draws_mean(mean(ans$.fitted)) - 100) < 2)
})

test_that("'draw_vals_augment_unfitted' works with 'bage_mod_norm'", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$income <- rnorm(n = nrow(data), mean = 200, sd = 30)
  data$wt <- rpois(n = nrow(data), lambda = 100)
  formula <- income ~ age + sex + time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  mod <- set_n_draw(mod, n_sim)
  ans_obtained <- draw_vals_augment_unfitted(mod, quiet = TRUE)
  set.seed(mod$seed_augment)
  components <- draw_vals_components_unfitted(mod = mod,
                                              n_sim = n_sim)
  disp <- components$.fitted[components$component == "disp"]
  fitted <- make_linpred_from_components(mod = mod,
                                         components = components,
                                         data = mod$data,
                                         dimnames_term = mod$dimnames_terms)
  fitted <- fitted * mod$outcome_sd + mod$outcome_mean
  disp <- disp * mod$outcome_sd * sqrt(mod$offset_mean)
  outcome <- draw_outcome_true(nm_distn = "norm",
                               fitted = fitted,
                               disp = disp,
                               offset = data$wt)
  ans_expected <- tibble::as_tibble(mod$data)
  ans_expected$income <- outcome
  ans_expected$.fitted <- fitted
  expect_equal(ans_obtained, ans_expected)
  expect_identical(names(augment(fit(mod), quiet = TRUE)), names(ans_obtained))
  expect_equal(rvec::draws_mean(mean(ans_obtained$.fitted)),
               mean(rvec::draws_mean(ans_obtained$income)), tolerance = 0.1)
})

test_that("'draw_vals_augment_unfitted' works with bage_mod_norm, noise data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                      KEEP.OUT.ATTRS = FALSE)
  data$income <- rnorm(n = nrow(data), mean = 100, sd = 5)
  data$wt <- runif(n = nrow(data), max = 1000)
  formula <- income ~ age + sex + time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt) |>
    set_datamod_noise(sd = 1)
  suppressMessages(
    expect_message(
      ans <- draw_vals_augment_unfitted(mod, quiet = FALSE),
      "Overwriting existing values for `income`.",
      )
  )
  expect_true(is.data.frame(ans))
  expect_setequal(names(ans), c(names(data), ".fitted", ".income"))
  expect_true(abs(rvec::draws_mean(mean(ans$.fitted)) - 100) < 2)
  expect_true(abs(rvec::draws_mean(mean(ans$.fitted)) - 100) < 2)
})


## 'draw_fitted_given_outcome' ------------------------------------------------

test_that("'draw_fitted_given_outcome' works with 'bage_mod_pois'", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  data$popn[1] <- NA
  data$deaths[2] <- NA
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  expected <- exp(rvec::rnorm_rvec(n = nrow(data), n_draw = 10))
  disp <- rvec::runif_rvec(n = 1, n_draw = 10)
  offset <- rvec::rpois_rvec(n = nrow(data), lambda = 100, n_draw = 10)
  offset[1] <- NA
  set.seed(1)
  ans_obtained <- draw_fitted_given_outcome(mod,
                                            outcome = data$deaths,
                                            offset = offset,
                                            expected = expected,
                                            disp = disp)
  set.seed(1)
  d <- data$deaths
  p <- offset
  d[1:2] <- 0
  p[1:2] <- 0
  ans_expected <- rvec::rgamma_rvec(n = nrow(data),
                                    shape = d + 1 / disp,
                                    rate = p  + 1 / (disp * expected))
  expect_equal(ans_obtained, ans_expected)
})

test_that("'draw_fitted_given_outcome' works with 'bage_mod_binom'", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  data$popn[1] <- NA
  data$deaths[2] <- NA
  formula <- deaths ~ age + sex + time
  mod <- mod_binom(formula = formula,
                  data = data,
                  size = popn)
  expected <- rvec::runif_rvec(n = nrow(data), n_draw = 10)
  disp <- rvec::runif_rvec(n = 1, n_draw = 10)
  offset <- rvec::rpois_rvec(n = nrow(data), lambda = 100, n_draw = 10)
  offset[1] <- NA
  set.seed(1)
  ans_obtained <- draw_fitted_given_outcome(mod,
                                            outcome = data$deaths,
                                            offset = offset,
                                            expected = expected,
                                            disp = disp)
  set.seed(1)
  d <- data$deaths
  p <- offset
  d[1:2] <- 0
  p[1:2] <- 0
  ans_expected <- rvec::rbeta_rvec(n = nrow(data),
                                   shape1 = d + expected / disp,
                                   shape2 = p - d  + (1 - expected) / disp)
  expect_equal(ans_obtained, ans_expected)
})


## 'fit' -----------------------------------------------------------------

test_that("'fit' works with valid inputs - pois has exposure", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with valid inputs - pois has exposure", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$deaths <- rpois(n = nrow(data), lambda = 10000)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = 1)
    mod <- set_prior(mod, age ~ RW2(sd = 0))
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with valid inputs - binom", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
    formula <- deaths ~ age + sex + time
    mod <- mod_binom(formula = formula,
                    data = data,
                    size = popn)
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with valid inputs - binom, disp is 0", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
    formula <- deaths ~ age + sex + time
    mod <- mod_binom(formula = formula,
                    data = data,
                    size = popn) |>
                    set_disp(mean = 0)
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with valid inputs - norm", {
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data[] <- lapply(data, factor)
    data$wt <- rpois(n = nrow(data), lambda = 100)
    data$val <- rnorm(n = nrow(data), mean = (data$sex == "F"))
    formula <- val ~ age + sex + time
    mod <- mod_norm(formula = formula,
                    data = data,
                    weights = wt)
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with known intercept and sex effect", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_prior(mod, `(Intercept)` ~ Known(values = -2))
    mod <- set_prior(mod, sex ~ Known(values = c(-0.1, 0.1)))
    ans_obtained <- fit(mod)
    expect_equal(ans_obtained$draws_effectfree[1,1], -2)
    expect_equal(ans_obtained$draws_effectfree[12:13,1], c(-0.1, 0.1))
})

test_that("'fit' works with AR1", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age:sex + age:time + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_prior(mod, time ~ AR1())
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' gives the similar imputed rate when outcome is NA and offset is NA", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2002, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age:sex + age:time + time
    ## outcome NA
    set.seed(0)
    data_outcome <- data
    data_outcome$deaths[5] <- NA
    mod_outcome <- mod_pois(formula = formula,
                            data = data_outcome,
                            exposure = popn)
    mod_outcome <- fit(mod_outcome)
    ans_outcome <- rvec::draws_median(augment(mod_outcome, quiet = TRUE)$.fitted[5])
    ## offset NA
    set.seed(0)
    data_offset <- data
    data_offset$popn[5] <- NA
    mod_offset <- mod_pois(formula = formula,
                           data = data_offset,
                           exposure = popn)
    mod_offset <- fit(mod_offset)
    ans_offset <- rvec::draws_median(augment(mod_offset, quiet = TRUE)$.fitted[5])
    ## compare
    expect_equal(ans_outcome, ans_offset, tolerance = 0.05)
})

test_that("'fit' works when all observed values for one year are NA", {
    data <- data.frame(deaths = c(NA, 2:10),
                       age = rep(1:2, each = 5),
                       time = 2001:2010)
    mod <- mod_pois(deaths ~ age + time,
                    data = data,
                    exposure = 1)
    mod_fitted <- fit(mod)
    expect_true(is_fitted(mod_fitted))
})

test_that("'fit' works when model consists of intercept only", {
    data <- data.frame(deaths = 1:10,
                       age = rep(1:2, each = 5),
                       time = 2001:2010)
    mod <- mod_pois(deaths ~ 1,
                    data = data,
                    exposure = 1)
    mod_fitted <- fit(mod)
    expect_true(is_fitted(mod_fitted))
})

test_that("'fit' works when model has no hyper-parameters", {
    data <- data.frame(deaths = 1:10,
                       age = rep(1:2, each = 5),
                       time = 2001:2010)
    mod <- mod_pois(deaths ~ 1,
                    data = data,
                    exposure = 1) |>
                    set_disp(mean = 0)
    mod_fitted <- fit(mod)
    expect_true(is_fitted(mod_fitted))
})

test_that("'fit' works when single dimension", {
    data <- data.frame(deaths = 1:10,
                       time = 2001:2010)
    mod <- mod_pois(deaths ~ time,
                    data = data,
                    exposure = 1) |>
      set_prior(time ~ RW(sd = 0))
    mod_fitted <- fit(mod)
    expect_identical(nrow(mod_fitted$draws_effectfree), nrow(data))
})

test_that("'fit' works with SVD", {
    set.seed(0)
    data <- expand.grid(age = poputils::age_labels(type = "five", max = 60),
                        time = 2000:2005,
                        sex = c("Female", "Male"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age:sex + age:time + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_prior(mod, age:sex ~ SVD(HMD))
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with Lin", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ sex + age + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_prior(mod, time ~ Lin())
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with Lin, n_by > 1", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ sex * time + age
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_prior(mod, sex:time ~ Lin())
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with RW - n_by > 0", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ sex * time + age
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_prior(mod, sex:time ~ RW())
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with RW2, n_by > 1", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2019, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ sex * time + age
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_prior(mod, sex:time ~ RW2())
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with no hyper", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ sex
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with AR", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ sex * time + age
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_prior(mod, time ~ AR(n = 2))
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with SVD, n_by > 1", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(age = c(0:59, "60+"), time = 2000:2005, reg = c("a", "b"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age:reg + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_prior(mod, age:reg ~ SVD(HMD))
  ans_obtained <- fit(mod)
  expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with SVD, n_by > 1", {
  set.seed(0)
  data <- expand.grid(age = c(0:59, "60+"), time = 2000:2001, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age:sex:time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_prior(mod, age:sex:time ~ SVD(HMD))
  ans_obtained <- fit(mod)
  expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with Lin_AR", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(age = c(0:59, "60+"), time = 2000:2005, reg = c("a", "b"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age:reg + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_prior(mod, time ~ Lin_AR(s = 0.2))
  ans_obtained <- fit(mod)
  expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with Lin(s = 0)", {
    set.seed(0)
    data <- expand.grid(age = c(0:59, "60+"), time = 2000:2005, reg = c("a", "b"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age:reg + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_prior(mod, time ~ Lin(s = 0))
    ans_obtained <- fit(mod)
    expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with RW2_Infant()", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(age = c(0:59, "60+"), time = 2000:2005, reg = c("a", "b"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age:reg + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_prior(mod, age:reg ~ RW2_Infant())
  ans_obtained <- fit(mod)
  mod <- set_prior(mod, age:reg ~ RW2_Infant(con = "by"))
  ans_obtained <- fit(mod)
  expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' and 'forecast' work with SVD_AR", {
  data <- expand.grid(age = poputils::age_labels(type = "five", min = 15, max = 60),
                      time = 2001:2010)
  data$population <- runif(n = nrow(data), min = 100, max = 300)
  data$deaths <- rpois(n = nrow(data), lambda = 0.05 * data$population)
  data$deaths[1] <- NA
  mod <- mod_pois(deaths ~ age:time,
                  data = data,
                  exposure = population) |>
                  set_prior(age:time ~ SVD_AR1(LFP))
  mod <- fit(mod)
  expect_true(is_fitted(mod))
  f <- forecast(mod, labels = 2011:2012)
  expect_setequal(c(names(f), ".deaths"),
                  names(augment(mod, quiet = TRUE)))
})

test_that("'fit' and 'forecast' work with SVD_RW", {
  data <- expand.grid(age = poputils::age_labels(type = "five", min = 15, max = 60),
                      time = 2001:2010)
  data$population <- runif(n = nrow(data), min = 100, max = 300)
  data$deaths <- rpois(n = nrow(data), lambda = 0.05 * data$population)
  data$deaths[1] <- NA
  mod <- mod_pois(deaths ~ age:time,
                  data = data,
                  exposure = population) |>
                  set_prior(age:time ~ SVD_RW(LFP))
  mod <- fit(mod)
  expect_true(is_fitted(mod))
  f <- forecast(mod, labels = 2011:2012)
  expect_setequal(c(names(f), ".deaths"),
                  names(augment(mod, quiet = TRUE)))
  mod <- mod |>
    set_prior(age:time ~ SVD_RW(LFP, sd = 0))
  mod <- fit(mod)
  expect_true(is_fitted(mod))
})

test_that("'fit' and 'forecast' work with SVD_RW2", {
  data <- expand.grid(age = poputils::age_labels(type = "five", min = 15, max = 60),
                      time = 2001:2010)
  data$population <- runif(n = nrow(data), min = 100, max = 300)
  data$deaths <- rpois(n = nrow(data), lambda = 0.05 * data$population)
  data$deaths[1] <- NA
  mod <- mod_pois(deaths ~ age:time,
                  data = data,
                  exposure = population) |>
                  set_prior(age:time ~ SVD_RW2(LFP))
  mod <- fit(mod)
  expect_true(is_fitted(mod))
  f <- forecast(mod, labels = 2011:2012)
  expect_setequal(c(names(f), ".deaths"),
                  names(augment(mod, quiet = TRUE)))
  mod <- mod |>
    set_prior(age:time ~ SVD_RW2(LFP, sd = 0))
  mod <- fit(mod)
  expect_true(is_fitted(mod))
})

test_that("'fit' works inner-outer", {
  set.seed(0)
  ## structure of data
  data <- expand.grid(age = poputils::age_labels(type = "lt"),
                      sex = c("Female", "Male"),
                      time = 2011:2015,
                      region = 1:10)
  data$population <- runif(n = nrow(data), min = 10, max = 1000)
  data$deaths <- NA
  ## generate single dataset
  mod_sim <- mod_pois(deaths ~ age * sex + region + time,
                      data = data,
                      exposure = population) |>
                      set_prior(`(Intercept)` ~ NFix(s = 0.1)) |>
                      set_prior(age ~ RW(s = 0.02)) |>
                      set_prior(sex ~ NFix(sd = 0.1)) |>
                      set_prior(age:sex ~ SVD(HMD)) |>
                      set_prior(time ~ Lin_AR1(s = 0.05, sd = 0.02)) |>
                      set_prior(region ~ NFix(sd = 0.05)) |>
                      set_disp(mean = 0.005)
  data_sim <- mod_sim |>
  set_n_draw(n = 1) |>
  augment(quiet = TRUE) |>
  dplyr::mutate(deaths = rvec::draws_median(deaths)) |>
  dplyr::select(age, sex, time, region, population, deaths)
  ## specify estimation model
  mod_est <- mod_pois(deaths ~ age * sex + region + time,
                      data = data_sim,
                      exposure = population) |>
                      set_prior(age:sex ~ SVD(HMD)) |>
                      set_prior(time ~ Lin_AR())
  ## fit estimation model
  mod_est <- mod_est |>
  fit(method = "inner-outer")
  expect_true(is_fitted(mod_est))
})

test_that("'fit' works when optimizer is 'optim'", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(age = c(0:59, "60+"), time = 2000:2005, reg = c("a", "b"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age:reg + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  ans_obtained <- fit(mod)
  expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works when 'quiet' is FALSE and optimizer is 'nlminb'", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(age = c(0:59, "60+"), time = 2000:2005, reg = c("a", "b"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age:reg + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  suppressMessages(
    capture.output(ans_obtained <- fit(mod, quiet = FALSE), file = NULL)
  )
  expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works when 'quiet' is FALSE and optimizer is 'BFGS'", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(age = c(0:59, "60+"), time = 2000:2005, reg = c("a", "b"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age:reg + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  suppressMessages(
    capture.output(ans_obtained <- fit(mod, optimizer = "BFGS", quiet = FALSE), file = NULL)
  )
  expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works when 'quiet' is FALSE and optimizer is 'CG'", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(age = c(0:59, "60+"), time = 2000:2005, reg = c("a", "b"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age:reg + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  suppressMessages(
    capture.output(ans_obtained <- fit(mod, optimizer = "CG", quiet = FALSE), file = NULL)
  )
  expect_s3_class(ans_obtained, "bage_mod")
})

test_that("'fit' works with covariates - no shrinkage", {
  set.seed(0)
  data <- expand.grid(age = 0:9,
                      region = c("a", "b"),
                      sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  data$income <- runif(n = nrow(data))
  data$distance <- runif(n = nrow(data))
  mod <- mod_pois(formula = deaths ~ age * sex ,
                  data = data,
                  exposure = popn)
  mod <- set_covariates(mod, ~ income + distance)
  ans_obtained <- fit(mod)
  expect_s3_class(ans_obtained, "bage_mod")
})



## 'forecast' -----------------------------------------------------------------

test_that("'forecast' works with fitted model - output is 'augment'", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2004, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age * sex + sex * time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- fit(mod)
    ans_no_est <- forecast(mod, labels = 2005:2006)
    ans_est <- forecast(mod, labels = 2005:2006, include_estimates = TRUE)
    expect_identical(names(ans_est), names(ans_no_est))
})

test_that("'forecast' works with fitted model - uses newdata", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2004, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age * sex + sex * time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- fit(mod)
    newdata <- make_data_forecast_labels(mod, labels = 2005)
    newdata$deaths <- rpois(n = nrow(newdata), lambda = 10)
    ans_est <- forecast(mod, newdata = newdata, include_estimates = TRUE)
    expect_identical(names(ans_est), names(augment(mod, quiet = TRUE)))
})


test_that("'forecast' gives same answer when run twice - output is 'augment'", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2004, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age * sex + sex * time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_n_draw(mod, n = 10)
    mod <- fit(mod)
    ans1 <- forecast(mod, labels = 2005:2006)
    ans2 <- forecast(mod, labels = 2005:2006)
    ans3 <- forecast(mod, labels = 2005:2006,
                     output = "aug",
                     include_estimates = TRUE)
    ans3 <- ans3[ans3$time >= 2005,]
    expect_identical(ans1, ans2)
    expect_identical(ans1, ans3)
})

test_that("'forecast' works with fitted model - output is 'components'", {
  set.seed(0)
  data <- expand.grid(age = 0:4, time = 2000:2004, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- fit(mod)
  ans_no_est <- forecast(mod,
                         labels = 2005:2006,
                         output = "comp")
  ans_est <- forecast(mod,
                      labels = 2005:2006,
                      output = "comp",
                      include_estimates = TRUE)
  expect_identical(names(ans_est), names(ans_no_est))
})

test_that("'forecast' gives same answer when run twice - output is 'components'", {
  set.seed(0)
  data <- expand.grid(age = 0:4, time = 2000:2004, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  ans1 <- forecast(mod,
                   labels = 2005:2006,
                   output = "comp")
  ans2 <- forecast(mod,
                   labels = 2005:2006,
                   output = "components")
  expect_identical(ans1, ans2)
})

test_that("'forecast' gives same answer when run twice - output is 'components'", {
  set.seed(0)
  data <- expand.grid(age = 0:4, time = 2000:2004, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  ans1 <- forecast(mod,
                   labels = 2005:2006,
                   output = "comp")
  ans2 <- forecast(mod,
                   labels = 2005:2006,
                   output = "components")
  expect_identical(ans1, ans2)
})

test_that("'forecast' throws correct error when time var not identified'", {
    set.seed(0)
    data <- expand.grid(age = 0:4, epoch = 2000:2004, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age * sex
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- fit(mod)
    expect_error(forecast(mod,
                          labels = 2005:2006,
                          output = "comp"),
                 "Can't forecast when time variable not identified.")
})

test_that("'forecast' throws correct error when labels and newdata both supplied'", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2004, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age * sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- fit(mod)
    labels <- 2007:2008
    newdata <- make_data_forecast_labels(mod = mod, labels = labels)
    expect_error(forecast(mod,
                          labels = labels,
                          newdata = newdata),
                 "Values supplied for `newdata` and for `labels`.")
})

test_that("'forecast' throws error when neither labels nor newdata supplied", {
    set.seed(0)
    data <- expand.grid(age = 0:4, time = 2000:2004, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 10)
    formula <- deaths ~ age * sex + sex * time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- fit(mod)
    newdata <- make_data_forecast_labels(mod, 2005:2006)
    newdata$deaths <- rpois(n = nrow(newdata), lambda = 10)
    expect_error(forecast(mod, include_estimates = TRUE),
                 "No value supplied for `newdata` or for `labels`.")
})

test_that("'forecast' throws error when exposure specified via formula", {
  set.seed(0)
  data <- expand.grid(age = 0:4, time = 2000:2004, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 10)
  formula <- deaths ~ age * sex + sex * time
  suppressWarnings(
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = ~popn + 1)
  )
  mod <- fit(mod)
  newdata <- make_data_forecast_labels(mod, 2005:2006)
  newdata$deaths <- rpois(n = nrow(newdata), lambda = 10)
  expect_error(forecast(mod, newdata = newdata),
               "`forecast\\(\\)` cannot be used with models where exposure specified using formula.")
})


## 'forecast_augment' --------------------------------------------------------

test_that("'forecast_augment' works - Poisson, has disp, no forecasted offset", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 100)
  data$exposure <- rpois(n = nrow(data), lambda = 1000)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = exposure)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod)
  labels_forecast <- 2006:2008
  data_forecast <- make_data_forecast_labels(mod = mod, labels_forecast = labels_forecast)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = FALSE,
                          has_newdata = FALSE)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
})


test_that("'forecast_augment' works - Poisson, no offset", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 100)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = 1)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod)
  data_forecast <- make_data_forecast_labels(mod = mod, labels_forecast = 2006:2008)
  labels_forecast <- 2006:2008
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = FALSE,
                          has_newdata = FALSE)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
  expect_false(rvec::is_rvec(ans$deaths))
  expect_false(".observed" %in% names(ans))
})

test_that("'forecast_augment' works - Poisson, has disp, has forecasted offset", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 100)
  data$exposure <- rpois(n = nrow(data), lambda = 1000)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = exposure)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod)
  data_forecast <- make_data_forecast_labels(mod = mod, labels_forecast = 2006:2008)
  labels_forecast <- 2006:2008
  data_forecast$exposure <- rpois(n = nrow(data_forecast), lambda = 1000)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = TRUE,
                          has_newdata = FALSE)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
  expect_true(rvec::is_rvec(ans$deaths))
  expect_true(all(is.na(ans$.observed)))
})

test_that("'forecast_augment' works - Poisson, has disp, has forecasted offset, imputed historical est", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 100)
  data$deaths[1] <- NA
  data$exposure <- rpois(n = nrow(data), lambda = 1000)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = exposure)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod)
  data_forecast <- make_data_forecast_labels(mod = mod, labels_forecast = 2006:2008)
  labels_forecast <- 2006:2008
  data_forecast$exposure <- rpois(n = nrow(data_forecast), lambda = 1000)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = TRUE,
                          has_newdata = FALSE)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
  expect_true(all(is.na(ans$deaths)))
  expect_true(all(is.na(ans$.observed)))
  expect_true(rvec::is_rvec(ans$.deaths))
})

test_that("'forecast_augment' works - Poisson, has disp, has forecasted offset, confidentialization", {
  set.seed(0)
  data <- expand.grid(age = 0:5, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 30) * 3
  data$deaths[1] <- NA
  data$exposure <- rpois(n = nrow(data), lambda = 1000)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = exposure)
  mod <- set_n_draw(mod, n = 10)
  mod <- set_confidential_rr3(mod)
  mod <- fit(mod)
  components_est <- components(mod)
  data_forecast <- make_data_forecast_labels(mod = mod, labels_forecast = 2006:2008)
  labels_forecast <- 2006:2008
  data_forecast$exposure <- rpois(n = nrow(data_forecast), lambda = 1000)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = TRUE,
                          has_newdata = FALSE)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
  expect_true(rvec::is_rvec(ans$deaths))
  expect_true(all(is.na(ans$.observed)))
  expect_true(rvec::is_rvec(ans$.deaths))
  expect_true(all(rvec::extract_draw(ans$deaths) %% 3 == 0))
})

test_that("'forecast_augment' works - binomial, no disp", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
  formula <- deaths ~ age + sex + time
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn)
  mod <- set_n_draw(mod, n = 10)
  mod <- set_disp(mod, mean = 0)
  mod <- fit(mod)
  components_est <- components(mod)
  labels_forecast <- 2006:2008
  data_forecast <- make_data_forecast_labels(mod= mod, labels_forecast = labels_forecast)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = FALSE,
                          has_newdata = has_newdata)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
})

test_that("'forecast_augment' works - normal", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data))
  formula <- income ~ age + sex * time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = popn)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod, quiet = TRUE)
  labels_forecast <- 2006:2008
  data_forecast <- make_data_forecast_labels(mod= mod, labels_forecast = labels_forecast)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = FALSE,
                          has_newdata = has_newdata)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
})

test_that("'forecast_augment' works - normal, has forecasted offset", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data), mean = 1000, sd = 100)
  formula <- income ~ age + sex * time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = popn)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod, quiet = TRUE)
  labels_forecast <- 2006:2008
  data_forecast <- make_data_forecast_labels(mod= mod, labels_forecast = labels_forecast)
  data_forecast$popn <- rpois(n = nrow(data_forecast), lambda = 1000)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = TRUE,
                          has_newdata = has_newdata)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
  expect_true(rvec::is_rvec(ans$income))
})

test_that("'forecast_augment' works - normal, no offset, no NA in outcome", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$income <- rnorm(n = nrow(data), mean = 10000, sd = 500)
  formula <- income ~ age + sex * time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = 1)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod, quiet = TRUE)
  labels_forecast <- 2006:2008
  data_forecast <- make_data_forecast_labels(mod= mod,
                                             labels_forecast = labels_forecast)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = FALSE,
                          has_newdata = has_newdata)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
  expect_true(is.numeric(ans$income))
})

test_that("'forecast_augment' works - normal, has offset, offset not in forecast data", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$income <- rnorm(n = nrow(data))
  data$wt <- runif(n = nrow(data))
  formula <- income ~ age + sex * time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod, quiet = TRUE)
  labels_forecast <- 2006:2008
  data_forecast <- make_data_forecast_labels(mod= mod,
                                             labels_forecast = labels_forecast)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = FALSE,
                          has_newdata = has_newdata)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
  expect_identical(ans$income, rep(NA_real_, nrow(ans)))
})

test_that("'forecast_augment' works - normal, estimated has imputed, has offset", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data))
  data$income[1] <- NA
  formula <- income ~ age + sex * time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = popn)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod, quiet = TRUE)
  labels_forecast <- 2006:2008
  data_forecast <- make_data_forecast_labels(mod= mod, labels_forecast = labels_forecast)
  data_forecast$popn <- rpois(n = nrow(data_forecast), lambda = 1000)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = TRUE,
                          has_newdata = has_newdata)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
  expect_true(rvec::is_rvec(ans$.income))
  expect_true(all(is.na(ans$income)))
})


test_that("'forecast_augment' works - normal, estimated has imputed, no offset", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$income <- rnorm(n = nrow(data))
  data$income[1] <- NA
  formula <- income ~ age + sex * time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = 1)
  mod <- set_n_draw(mod, n = 10)
  mod <- fit(mod)
  components_est <- components(mod, quiet = TRUE)
  labels_forecast <- 2006:2008
  data_forecast <- make_data_forecast_labels(mod= mod, labels_forecast = labels_forecast)
  set.seed(1)
  components_forecast <- forecast_components(mod = mod,
                                             components_est = components_est,
                                             labels_forecast = labels_forecast)
  components <- vctrs::vec_rbind(components_est, components_forecast)
  dimnames_forecast <- make_dimnames_terms_forecast(dimnames_terms = mod$dimnames_terms,
                                                    var_time = mod$var_time,
                                                    labels_forecast = labels_forecast,
                                                    time_only = FALSE)
  linpred_forecast <- make_linpred_from_components(mod = mod,
                                                   components = components,
                                                   data = data_forecast,
                                                   dimnames_terms = dimnames_forecast)
  ans <- forecast_augment(mod = mod,
                          data_forecast = data_forecast,
                          components_forecast = components_forecast,
                          linpred_forecast = linpred_forecast,
                          has_offset_forecast = FALSE,
                          has_newdata = has_newdata)
  aug_est <- augment(mod, quiet = TRUE)
  expect_setequal(ans$age, aug_est$age)
  expect_setequal(ans$sex, aug_est$sex)
  expect_setequal(ans$time, 2006:2008)
  expect_identical(names(ans), names(aug_est))
  expect_true(all(is.na(ans$.income)))
  expect_true(all(is.na(ans$income)))
})


## 'get_fun_ag_offset' --------------------------------------------------------

test_that("'get_fun_ag_offset' works with Poisson - varying offset", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
  formula <- deaths ~ age + sex
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  f <- get_fun_ag_offset(mod)
  ans_obtained <- f(mod$offset)
  ans_expected <- sum(mod$offset)
  expect_identical(ans_obtained, ans_expected)
})

test_that("'get_fun_ag_offset' works with Poisson - offset = 1", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
  formula <- deaths ~ age + sex
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = 1)
  f <- get_fun_ag_offset(mod)
  ans_obtained <- f(mod$offset)
  ans_expected <- 1
  expect_identical(ans_obtained, ans_expected)
})

test_that("'get_fun_ag_offset' works with binom", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
  formula <- deaths ~ age + sex
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn)
  f <- get_fun_ag_offset(mod)
  ans_obtained <- f(mod$offset)
  ans_expected <- sum(mod$offset)
  expect_identical(ans_obtained, ans_expected)
})

test_that("'get_fun_ag_offset' works with norm, varying offset", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$wt <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data))
  formula <- income ~ age + sex
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  f <- get_fun_ag_offset(mod)
  ans_obtained <- f(mod$offset)
  n <- length(mod$offset)
  ans_expected <- (n^2) / sum(1 / mod$offset)
  expect_identical(ans_obtained, ans_expected)
})

test_that("'get_fun_ag_offset' works with norm, offset = 1", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$income <- rnorm(n = nrow(data))
  formula <- income ~ age + sex
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = 1)
  f <- get_fun_ag_offset(mod)
  ans_obtained <- f(mod$offset)
  n <- length(mod$offset)
  ans_expected <- (n^2) / n
  expect_identical(ans_obtained, ans_expected)
})


## 'get_fun_ag_outcome' --------------------------------------------------------

test_that("'get_fun_ag_outcome' works with binom", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
  formula <- deaths ~ age + sex + time
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn)
  f <- get_fun_ag_outcome(mod)
  ans_obtained <- f(mod$outcome)
  ans_expected <- sum(mod$outcome)
  expect_identical(ans_obtained, ans_expected)
})

test_that("'get_fun_ag_outcome' works with norm", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$wt <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data))
  formula <- income ~ age + sex + time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  f <- get_fun_ag_outcome(mod)
  ans_obtained <- f(mod$outcome)
  ans_expected <- mean(mod$outcome)
  expect_identical(ans_obtained, ans_expected)
})

## 'get_fun_inv_transform' ----------------------------------------------------

test_that("'get_fun_inv_transform' works with valid inputs", {
    set.seed(0)
    x <- runif(100)
    logit <- function(x) log(x) - log(1 - x)
    expect_equal(get_fun_inv_transform(structure(1, class = "bage_mod_pois"))(log(x)), x)
    expect_equal(get_fun_inv_transform(structure(1, class = "bage_mod_binom"))(logit(x)), x)
    expect_equal(get_fun_inv_transform(structure(1, class = "bage_mod_norm"))(x), x)
})


## 'get_fun_orig_scale_disp' --------------------------------------------------

test_that("'get_fun_orig_scale_disp' works with valid inputs", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                      KEEP.OUT.ATTRS = FALSE)
  data$income <- rnorm(n = nrow(data))
  data$wt <- runif(n = nrow(data))
  formula <- income ~ age + sex + time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  f <- get_fun_orig_scale_disp(mod)
  x <- runif(1)
  ans_obtained <- f(x)
  ans_expected <- x * sd(data$income) * sqrt(mean(data$wt))
  expect_equal(ans_obtained, ans_expected)
})


## 'get_fun_orig_scale_linpred' -----------------------------------------------

test_that("'get_fun_orig_scale_linpred' works with valid inputs", {
  expect_equal(get_fun_orig_scale_linpred(
    structure(list(outcome_mean = 3, outcome_sd = 2),
              class = c("bage_mod_norm", "bage_mod")))(1), 5)
})


## 'get_fun_orig_scale_offset' ------------------------------------------------

test_that("'get_fun_orig_scale_offset' works with valid inputs", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"),
                      KEEP.OUT.ATTRS = FALSE)
  data$income <- rnorm(n = nrow(data))
  data$wt <- runif(n = nrow(data))
  formula <- income ~ age + sex + time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  f <- get_fun_orig_scale_offset(mod)
  x <- runif(n = nrow(data))
  ans_obtained <- f(x)
  ans_expected <- x * mean(data$wt)
  expect_equal(ans_obtained, ans_expected)
})


## 'get_nm_offset_data' -------------------------------------------------------

test_that("'get_nm_offset_data' works with 'bage_mod_pois' - variable name", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 20)
  data$popn <- rpois(n = nrow(data), lambda = 30)
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  expect_identical(get_nm_offset_data(mod), "popn")
})


## 'get_nm_offset_mod' --------------------------------------------------------

test_that("'get_nm_offset_mod' works with 'bage_mod_pois'", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 20)
  data$popn <- rpois(n = nrow(data), lambda = 30)
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  expect_identical(get_nm_offset_mod(mod), "exposure")
})

test_that("'get_nm_offset_mod' works with 'bage_mod_binom'", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 20)
  data$deaths <- rbinom(n = nrow(data), prob = 0.5, size = data$popn)
  formula <- deaths ~ age + sex + time
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn)
  expect_identical(get_nm_offset_mod(mod), "size")
})

test_that("'get_nm_offset_mod' works with 'bage_mod_norm'", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 20)
  data$wt <- runif(n = nrow(data))
  formula <- deaths ~ age + sex + time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  expect_identical(get_nm_offset_mod(mod), "weights")
})


## 'get_nm_outcome_data' ------------------------------------------------------

test_that("'get_nm_outcome_data' works with 'bage_mod_pois'", {
  set.seed(0)
  n_sim <- 10
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$deaths <- rpois(n = nrow(data), lambda = 20)
  data$popn <- rpois(n = nrow(data), lambda = 30)
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  expect_identical(get_nm_outcome_data(mod), "deaths")
})


## 'has_confidential' ---------------------------------------------------------

test_that("'has_confidential' works with valid inputs", {
    data <- data.frame(deaths = 1:10 * 3,
                       time = 2001:2010,
                       income = rnorm(10))
    mod <- mod_pois(deaths ~ time,
                    data = data,
                    exposure = 1)
    expect_false(has_confidential(mod))
    mod <- set_confidential_rr3(mod)
    expect_true(has_confidential(mod))
})


## 'has_covariates' -----------------------------------------------------------

test_that("'has_covariates' works with valid inputs", {
    data <- data.frame(deaths = 1:10,
                       time = 2001:2010,
                       income = rnorm(10))
    mod <- mod_pois(deaths ~ time,
                    data = data,
                    exposure = 1)
    expect_false(has_covariates(mod))
    mod <- set_covariates(mod, ~ income)
    expect_true(has_covariates(mod))
})


## 'has_datamod' --------------------------------------------------------------

test_that("'has_datamod' works with valid inputs", {
    data <- data.frame(deaths = 1:10 * 3,
                       time = 2001:2010,
                       income = rnorm(10))
    mod <- mod_pois(deaths ~ time,
                    data = data,
                    exposure = 1)
    expect_false(has_datamod(mod))
    mod <- mod |>
      set_datamod_undercount(prob = data.frame(mean = 0.9, disp = 0.2))
    expect_true(has_datamod(mod))
})


## 'has_datamod_outcome' ------------------------------------------------------

test_that("'has_datamod_outcome' works with valid inputs", {
    data <- data.frame(deaths = 1:10 * 3,
                       time = 2001:2010,
                       popn = rep(100, 10),
                       income = rnorm(10))
    mod <- mod_pois(deaths ~ time,
                    data = data,
                    exposure = popn) |>
      set_disp(mean = 0)
    expect_false(has_datamod_outcome(mod))
    mod <-
      set_datamod_exposure(mod, cv = 0.1)
    expect_true(has_datamod(mod))
    suppressMessages(
    mod <- set_datamod_undercount(mod,
                                  prob = data.frame(mean = 0.9, disp = 0.2))
    )
    expect_true(has_datamod(mod))
})


## 'has_datamod_param' --------------------------------------------------------

test_that("'has_datamod_param' works with Poisson", {
  data <- data.frame(deaths = 1:10 * 3,
                     time = 2001:2010)
  mod <- mod_pois(deaths ~ time,
                  data = data,
                  exposure = 1)
  expect_false(has_datamod_param(mod))
  mod <- mod |>
    set_datamod_overcount(rate = data.frame(mean = 0.1, disp = 0.2))
  expect_true(has_datamod_param(mod))
})

test_that("'has_datamod_param' works with normal", {
  data <- data.frame(time = 2001:2010,
                     income = rnorm(10))
  mod <- mod_norm(income ~ time,
                  data = data,
                  weights = 1)
    expect_false(has_datamod_param(mod))
  mod <- mod |>
    set_datamod_noise(sd = 0.2)
  expect_false(has_datamod_param(mod))
})


## 'has_disp' -----------------------------------------------------------------

test_that("'has_disp' works with valid inputs", {
    data <- data.frame(deaths = 1:10,
                       time = 2001:2010)
    mod <- mod_pois(deaths ~ time,
                    data = data,
                    exposure = 1)
    expect_true(has_disp(mod))
    mod <- set_disp(mod, mean = 0)
    expect_false(has_disp(mod))
})


## 'has_varying_offset' -------------------------------------------------------

test_that("'has_varying_offset' works with Poisson, valid inputs", {
  data <- data.frame(deaths = 1:10,
                     popn = 21:30,
                     time = 2001:2010)
  mod <- mod_pois(deaths ~ time,
                  data = data,
                  exposure = popn)
  expect_true(has_varying_offset(mod))
  mod <- mod_pois(deaths ~ time,
                  data = data,
                  exposure = 1)
  expect_false(has_varying_offset(mod))
})

test_that("'has_varying_offset' raises correct errors with Poisson", {
  data <- data.frame(deaths = 1:10,
                     popn = 21:30,
                     time = 2001:2010)
  mod <- mod_pois(deaths ~ time,
                  data = data,
                  exposure = popn)
  mod_wrong <- mod
  mod_wrong$offset <- NULL
  expect_error(has_varying_offset(mod_wrong),
               "Internal error: offset is NULL")
  mod_wrong <- mod
  mod_wrong$offset <- rep(NA, 10)
  expect_error(has_varying_offset(mod_wrong),
               "Internal error: offset all NA")
  mod_wrong <- mod
  mod_wrong$nm_offset_data <- NULL
  expect_error(has_varying_offset(mod_wrong),
               "Internal error: offset not all ones, but no nm_offset_data")
  mod_wrong <- mod
  mod_wrong$offset <- rep(1, 10)
  mod_wrong$nm_offset_data <- "popn"
  expect_error(has_varying_offset(mod_wrong),
               "Internal error: offset all ones, but has nm_offset_data")
})

test_that("'has_varying_offset' works with binomial, valid inputs", {
  data <- data.frame(deaths = 1:10,
                     popn = 21:30,
                     time = 2001:2010)
  mod <- mod_binom(deaths ~ time,
                  data = data,
                  size = popn)
  expect_true(has_varying_offset(mod))
})

test_that("'has_varying_offset' raises correct errors with Poisson", {
  data <- data.frame(deaths = 1:10,
                     popn = 21:30,
                     time = 2001:2010)
  mod <- mod_binom(deaths ~ time,
                  data = data,
                  size = popn)
  mod_wrong <- mod
  mod_wrong$offset <- NULL
  expect_error(has_varying_offset(mod_wrong),
               "Internal error: offset is NULL")
  mod_wrong <- mod
  mod_wrong$offset <- rep(NA, 10)
  expect_error(has_varying_offset(mod_wrong),
               "Internal error: offset all NA")
  mod_wrong <- mod
  mod_wrong$nm_offset_data <- NULL
  expect_error(has_varying_offset(mod_wrong),
               "Internal error: nm_offset_data is NULL")
})


## 'is_fitted' ----------------------------------------------------------------

test_that("'is_fitted' works with valid inputs", {
  data <- data.frame(deaths = 1:10,
                     time = 2001:2010)
  mod <- mod_pois(deaths ~ time,
                  data = data,
                  exposure = 1)
  expect_false(is_fitted(mod))
  mod <- fit(mod)
  expect_true(is_fitted(mod))
})


## 'make_disp_obs' ------------------------------------------------------------

test_that("'make_disp_obs' works with Poisson, exposure data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  cv <- data.frame(time = 2000:2006, cv = 2)
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_disp(mean = 0) |>
    set_datamod_exposure(cv = cv)
  ans_obtained <- make_disp_obs(mod = mod,
                                components = NULL)
  d <- get_datamod_disp(mod$datamod)
  ans_expected <- 1 / (3 + d^{-1})
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_disp_obs' works with Poisson", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  rate <- data.frame(sex = c("F", "M"), mean = c(1.1, 1.2), disp = c(0.5, 0.3))
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    fit()
  ans_obtained <- make_disp_obs(mod = mod,
                                components = NULL)
  ans_expected <- rep(get_disp(mod), times = nrow(data))
  expect_equal(ans_obtained, ans_expected)
  mod <- mod |>
    set_disp(mean = 0) |>
    fit()
  ans_obtained <- make_disp_obs(mod = mod,
                                components = NULL)
  ans_expected <- NULL
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_disp_obs' works with Poisson, overcount data model, mean_disp is 0", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  rate <- data.frame(sex = c("F", "M"), mean = c(1.1, 1.2), disp = c(0.5, 0.3))
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_disp(mean = 0) |>
    set_datamod_overcount(rate = rate)
  ans_obtained <- make_disp_obs(mod = mod,
                                components = NULL)
  ans_expected <- NULL
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_disp_obs' works with binomial, undercount data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  prob <- data.frame(sex = c("F", "M"), mean = c(0.1, 0.2), disp = c(0.5, 0.3))
  mod <- mod_binom(formula = formula,
                  data = data,
                  size = popn) |>
    set_datamod_undercount(prob = prob) |>
    fit()
  ans_obtained <- make_disp_obs(mod = mod,
                                components = NULL)
  ans_expected <- rep(get_disp(mod), times = nrow(data))
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_disp_obs' works with binomial, mean_disp is 0", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  prob <- data.frame(sex = c("F", "M"), mean = c(0.1, 0.2), disp = c(0.5, 0.3))
  mod <- mod_binom(formula = formula,
                  data = data,
                  size = popn) |>
    set_n_draw(10) |>
    fit()
  ans_obtained <- make_disp_obs(mod = mod,
                                components = NULL)
  ans_expected <- rep(get_disp(mod), times = nrow(data))
  expect_equal(ans_obtained, ans_expected)
  mod <- mod |>
    set_disp(mean = 0) |>
    fit()
  ans_obtained <- make_disp_obs(mod = mod,
                                components = NULL)
  ans_expected <- NULL
  expect_equal(ans_obtained, ans_expected)
})


## 'make_expected_obs' --------------------------------------------------------

test_that("'make_expected_obs' works with Poisson, no data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  components <- data.frame(term = c("(Intercept)", rep("datamod", 6)),
                           component = c("(Intercept)", rep("disp", 6)),
                           level = c("(Intercept)", 2000:2005),
                           .fitted = rvec::runif_rvec(7, n_draw = 10))
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  ans_obtained <- make_expected_obs(mod = mod,
                                    components = components,
                                    expected = expected)
  ans_expected <- expected
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_expected_obs' works with Poisson, exposure data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  cv <- data.frame(time = 2000:2006, cv = 2)
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_disp(mean = 0) |>
    set_datamod_exposure(cv)
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  ans_obtained <- make_expected_obs(mod = mod,
                                    components = NULL,
                                    expected = expected)
  d <- get_datamod_disp(mod$datamod)
  ans_expected <- ((3 + d^{-1})/(1 + d^{-1})) * expected
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_expected_obs' works with Poisson, miscount data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  prob <- data.frame(sex = c("F", "M"), mean = c(0.1, 0.2), disp = c(1, 3))
  rate <- data.frame(time = 2000:2006, mean = 2, disp = c(0.5, 0.2, 0.3, 0.4, 0.2, 0.1, 0.1))
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_datamod_miscount(prob = prob, rate = rate)
  components <- data.frame(term = c("(Intercept)", rep("datamod", 8)),
                           component = c("(Intercept)", rep(c("prob", "rate"), times = c(2, 6))),
                           level = c("(Intercept)", c("F", "M", 2000:2005)),
                           .fitted = rvec::runif_rvec(9, n_draw = 10))
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  ans_obtained <- make_expected_obs(mod = mod,
                                    components = components,
                                    expected = expected)
  p <- get_datamod_prob(mod$datamod, components)
  r <- get_datamod_rate(mod$datamod, components)
  ans_expected <- (p + r) * expected
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_expected_obs' works with Poisson, overcount data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  rate <- data.frame(time = 2000:2006, mean = 2, disp = c(0.5, 0.2, 0.3, 0.4, 0.2, 0.1, 0.1))
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_datamod_overcount(rate = rate)
  components <- data.frame(term = c("(Intercept)", rep("datamod", 6)),
                           component = c("(Intercept)", rep("rate", times = 6)),
                           level = c("(Intercept)", 2000:2005),
                           .fitted = rvec::runif_rvec(7, n_draw = 10))
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  ans_obtained <- make_expected_obs(mod = mod,
                                    components = components,
                                    expected = expected)
  r <- get_datamod_rate(mod$datamod, components)
  ans_expected <- (1 + r) * expected
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_expected_obs' works with Poisson throws error with invalid data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  rate <- data.frame(time = 2000:2006, mean = 2, disp = c(0.5, 0.2, 0.3, 0.4, 0.2, 0.1, 0.1))
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_datamod_overcount(rate = rate)
  components <- data.frame(term = c("(Intercept)", rep("datamod", 6)),
                           component = c("(Intercept)", rep("rate", times = 6)),
                           level = c("(Intercept)", 2000:2005),
                           .fitted = rvec::runif_rvec(7, n_draw = 10))
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  class(mod$datamod) <- "wrong"
  expect_error(make_expected_obs(mod = mod,
                                 components = components,
                                 expected = expected),
               "Internal error: Can't handle data model with class")
})


test_that("'make_expected_obs' works with Poisson, undercount data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  prob <- data.frame(sex = c("F", "M"), mean = c(0.1, 0.2), disp = c(1, 3))
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_datamod_undercount(prob = prob)
  components <- data.frame(term = c("(Intercept)", rep("datamod", 2)),
                           component = c("(Intercept)", rep("prob", times = 2)),
                           level = c("(Intercept)", c("F", "M")),
                           .fitted = rvec::runif_rvec(3, n_draw = 10))
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  ans_obtained <- make_expected_obs(mod = mod,
                                    components = components,
                                    expected = expected)
  p <- get_datamod_prob(mod$datamod, components)
  ans_expected <- p * expected
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_expected_obs'  with binom throws expected error with invalid data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- round(0.5 * data$popn)
  formula <- deaths ~ age + time + sex
  mean <- data.frame(sex = c("F", "M"), mean = c(0.1, 0.2))
  sd <- data.frame(sex = c("F", "M"), sd = c(0.1, 0.2))
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn)
  datamod <- new_bage_datamod_noise(sd_sd = 0.3,
                                    sd_levels = "sd",
                                    sd_matrix_outcome = Matrix::sparseMatrix(x = rep(1, 120),
                                                                             i = seq_len(120),
                                                                             j = rep(1, 120)),
                                    sd_arg = sd,
                                    nms_by = c("age", "sex"),
                                    outcome_sd = 2)
  mod$datamod <- datamod
  components <- data.frame(term = c("(Intercept)", rep("datamod", 2)),
                           component = c("(Intercept)", rep("prob", times = 2)),
                           level = c("(Intercept)", c("F", "M")),
                           .fitted = rvec::runif_rvec(3, n_draw = 10))
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  expect_error(make_expected_obs(mod = mod,
                                 components = components,
                                 expected = expected),
               "Internal error: Can't handle data model")
})

test_that("'make_expected_obs' works with binomial, no data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn)
  components <- data.frame(term = c("(Intercept)", rep("datamod", 6)),
                           component = c("(Intercept)", rep("disp", 6)),
                           level = c("(Intercept)", 2000:2005),
                           .fitted = rvec::runif_rvec(7, n_draw = 10))
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  ans_obtained <- make_expected_obs(mod = mod,
                                    components = components,
                                    expected = expected)
  ans_expected <- expected
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_expected_obs' works with Binomial, undercount data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  prob <- data.frame(sex = c("F", "M"), mean = c(0.1, 0.2), disp = c(1, 3))
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn) |>
    set_datamod_undercount(prob = prob)
  components <- data.frame(term = c("(Intercept)", rep("datamod", 2)),
                           component = c("(Intercept)", rep("prob", times = 2)),
                           level = c("(Intercept)", c("F", "M")),
                           .fitted = rvec::runif_rvec(3, n_draw = 10))
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  ans_obtained <- make_expected_obs(mod = mod,
                                    components = components,
                                    expected = expected)
  p <- get_datamod_prob(mod$datamod, components)
  ans_expected <- p * expected
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_expected_obs' with binomial throws expected error with invalid data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  mean <- data.frame(sex = c("F", "M"), mean = c(0.1, 0.2))
  sd <- data.frame(sex = c("F", "M"), sd = c(0.1, 0.2))
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn)
  datamod <- new_bage_datamod_noise(sd_sd = 0.3,
                                    sd_levels = "sd",
                                    sd_matrix_outcome = Matrix::sparseMatrix(x = rep(1, 120),
                                                                             i = seq_len(120),
                                                                             j = rep(1, 120)),
                                    sd_arg = sd,
                                    nms_by = c("age", "sex"),
                                    outcome_sd = 2)
  mod$datamod <- datamod
  components <- data.frame(term = c("(Intercept)", rep("datamod", 2)),
                           component = c("(Intercept)", rep("prob", times = 2)),
                           level = c("(Intercept)", c("F", "M")),
                           .fitted = rvec::runif_rvec(3, n_draw = 10))
  expected <- rvec::runif_rvec(n = 120, n_draw = 10)
  expect_error(make_expected_obs(mod = mod,
                                 components = components,
                                 expected = expected),
               "Internal error: Can't handle data model")
})



## 'make_i_lik' ---------------------------------------------------------------

test_that("'make_i_lik' works with bage_mod_pois", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- 3 * rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  expect_identical(make_i_lik(mod), 100000L)
  mod <- mod |>
    set_datamod_overcount(rate = data.frame(mean = 0.2, disp = 3))
  expect_identical(make_i_lik(mod), 104000L)
  mod <- mod |>
    set_confidential_rr3()
  expect_identical(make_i_lik(mod), 104010L)
})


## 'make_i_lik_part' ----------------------------------------------------------

test_that("'make_i_lik_part' works with bage_mod_pois", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  expect_identical(make_i_lik_part(mod), 100000L)
  mod0 <- set_disp(mod, mean = 0)
  expect_identical(make_i_lik_part(mod0), 200000L)
})

test_that("'make_i_lik_part' works with bage_mod_binom", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.02)
  formula <- deaths ~ age + time + sex
  mod <- mod_binom(formula = formula,
                  data = data,
                  size = popn)
  expect_identical(make_i_lik_part(mod), 300000L)
  mod0 <- set_disp(mod, mean = 0)
  expect_identical(make_i_lik_part(mod0), 400000L)
})

test_that("'make_i_lik_part' works with bage_mod_norm", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$wt <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data), sd = 10)
  formula <- income ~ age + time + sex
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  expect_identical(make_i_lik_part(mod), 500000L)
})


## 'make_mod_disp' -----------------------------------------------------------

test_that("'make_mod_disp' works with pois", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rpois(n = nrow(data), lambda = 0.3 * data$popn)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- fit(mod)
  mod_disp <- make_mod_disp(mod)
  expect_setequal(names(mod_disp$priors), "(Intercept)")
  mu <- exp(make_linpred_from_stored_draws(mod, point = TRUE))
  expect_equal(mod_disp$offset, mod$offset * mu)
  expect_true(mod_disp$mean_disp > 0)
  expect_identical(length(mod_disp$dimnames_terms), 1L)
})

test_that("'make_mod_disp' works with pois - large dataset", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(time = 1:6000, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rpois(n = nrow(data), lambda = 0.3 * data$popn)
  formula <- deaths ~ sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_prior(time ~ RW(sd = 0))
  mod$point_effectfree <- rnorm(1 + 5999 + 2)
  mod_disp <- make_mod_disp(mod)
  expect_true(identical(nrow(mod_disp$data), 10000L))
})

test_that("'make_mod_disp' works with binom", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), size =  data$popn, prob = 0.3)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn)
  mod <- fit(mod)
  mod_disp <- make_mod_disp(mod)
  expect_setequal(names(mod_disp$priors), names(mod$priors))
  expect_true(mod_disp$mean_disp > 0)
  expect_s3_class(mod_disp$priors[["age"]], "bage_prior_known")
  expect_equal(mod_disp$offset, mod$offset)
})

test_that("'make_mod_disp' works with binom", {
  set.seed(0)
  data <- expand.grid(time = 1:6, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), prob = 0.2, size = data$popn)
  formula <- deaths ~ sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  mod <- fit(mod)
  mod_disp <- make_mod_disp(mod)
  expect_setequal(names(mod_disp$priors), "(Intercept)")
  mu <- exp(make_linpred_from_stored_draws(mod, point = TRUE))
  expect_true(mod_disp$mean_disp > 0)
})

test_that("'make_mod_disp' works with binom - large dataset", {
  set.seed(0)
  data <- expand.grid(time = 1:6000, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), prob = 0.3, size = data$popn)
  formula <- deaths ~ sex + time
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn) |>
    set_prior(time ~ RW2(sd = 0))
  mod$point_effectfree <- rnorm(1 + 5999 + 2)
  mod$draws_effectfree  <- 1 ## to fool 'is_fitted'
  mod_disp <- make_mod_disp(mod)
  expect_true(identical(nrow(mod_disp$data), 10000L))
  expect_identical(mod_disp$nm_offset_data, "offset_inner_outer")
})

test_that("'make_mod_disp' works with norm", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$wt <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data))
  formula <- income ~ age * sex + sex * time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  mod <- fit(mod)
  mod_disp <- make_mod_disp(mod)
  expect_setequal(names(mod_disp$priors), "(Intercept)")
  mu <- make_linpred_from_stored_draws(mod, point = TRUE)
  expect_equal(mod_disp$outcome, mod$outcome - mu)
  expect_true(mod_disp$mean_disp > 0)
  expect_identical(length(mod_disp$dimnames_terms), 1L)
  expect_identical(mod_disp$nm_offset_data, "offset_inner_outer")
})

test_that("'make_mod_disp' works with norm - large dataset", {
  testthat::skip_on_cran()
  set.seed(0)
  data <- expand.grid(time = 2000:8000, sex = c("F", "M"))
  data$wt <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data))
  formula <- income ~ time + sex
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  mod$point_effectfree <- rnorm(1 + 6001 + 2)
  mod_disp <- make_mod_disp(mod)
  expect_true(identical(nrow(mod_disp$data), 10000L))
  expect_identical(mod_disp$nm_offset_data, "offset_inner_outer")
})


## 'make_mod_inner' -----------------------------------------------------------

test_that("'make_mod_inner' works with pois", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rpois(n = nrow(data), lambda = 0.3 * data$popn)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  use_term <- make_use_term(mod, vars_inner = c("age", "sex"))
  ans_obtained <- make_mod_inner(mod, use_term)
  ans_expected <- reduce_model_terms(mod, use_term = use_term)
  ans_expected$mean_disp <- 0
  expect_identical(ans_obtained, ans_expected)
})

test_that("'make_mod_inner' works with norm", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$wt <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data))
  formula <- income ~ age * sex + sex * time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  use_term <- make_use_term(mod, vars_inner = c("age", "sex"))
  ans_obtained <- make_mod_inner(mod, use_term = use_term)
  ans_expected <- reduce_model_terms(mod, use_term = use_term)
  expect_identical(ans_obtained, ans_expected)
})


## 'make_mod_outer' -----------------------------------------------------------

test_that("'make_mod_outer' works with pois", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rpois(n = nrow(data), lambda = 0.3 * data$popn)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  use_term <- make_use_term(mod, vars_inner = c("age", "sex"))
  mod_inner <- reduce_model_terms(mod, use_term = use_term)
  mod_inner <- fit(mod_inner)
  mod_outer <- make_mod_outer(mod,
                              mod_inner = mod_inner,
                              use_term = use_term)
  expect_setequal(names(mod_outer$priors), c("time", "sex:time"))
  mu <- exp(make_linpred_from_stored_draws(mod_inner, point = TRUE))
  expect_equal(mod_outer$offset, mod$offset * mu)
  expect_equal(mod_outer$mean_disp, 0)
  expect_identical(mod_outer$nm_offset_data, "offset_inner_outer")
})

test_that("'make_mod_outer' works with binom", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 1000)
  data$deaths <- rbinom(n = nrow(data), size =  data$popn, prob = 0.3)
  formula <- deaths ~ age * sex + sex * time
  mod <- mod_binom(formula = formula,
                   data = data,
                   size = popn)
  use_term <- make_use_term(mod, vars_inner = c("age", "sex"))
  mod_inner <- reduce_model_terms(mod, use_term = use_term)
  mod_inner <- fit(mod_inner)
  mod_outer <- make_mod_outer(mod,
                              mod_inner = mod_inner,
                              use_term = use_term)
  expect_setequal(names(mod_outer$priors), names(mod$priors))
  expect_s3_class(mod_outer$priors[["age"]], "bage_prior_known")
  expect_equal(mod_outer$offset, mod$offset)
  expect_equal(mod_outer$mean_disp, 0)
  expect_identical(mod_outer$nm_offset_data, "offset_inner_outer")
})

test_that("'make_mod_outer' works with norm", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$wt <- rpois(n = nrow(data), lambda = 1000)
  data$income <- rnorm(n = nrow(data))
  formula <- income ~ age * sex + sex * time
  mod <- mod_norm(formula = formula,
                  data = data,
                  weights = wt)
  use_term <- make_use_term(mod, vars_inner = c("age", "sex"))
  mod_inner <- reduce_model_terms(mod, use_term = use_term)
  mod_inner <- fit(mod_inner)
  mod_outer <- make_mod_outer(mod,
                              mod_inner = mod_inner,
                              use_term = use_term)
  expect_setequal(names(mod_outer$priors), c("time", "sex:time"))
  mu <- make_linpred_from_stored_draws(mod_inner, point = TRUE)
  expect_equal(mod_outer$outcome, mod$outcome - mu)
  expect_true(mod_outer$mean_disp > 0)
  expect_identical(mod_outer$nm_offset_data, "offset_inner_outer")
})


## 'make_sd_obs' --------------------------------------------------------------

test_that("'make_sd_obs' works with Poisson, noise data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_disp(mean = 0) |>
    set_datamod_noise(sd = 1)
  ans_obtained <- make_sd_obs(mod)
  ans_expected <- rep(1, times = nrow(data))
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_sd_obs' works with Poisson, no data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- rpois(n = nrow(data), lambda = 5)
  formula <- deaths ~ age + time + sex
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn)
  ans_obtained <- make_sd_obs(mod)
  ans_expected <- NULL
  expect_equal(ans_obtained, ans_expected)
})

test_that("'make_sd_obs' works with binomial, no data model", {
  set.seed(0)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- round(0.5 * data$popn)
  formula <- deaths ~ age + time + sex
  mod <- mod_binom(formula = formula,
                  data = data,
                  size = popn)
  ans_obtained <- make_sd_obs(mod)
  ans_expected <- NULL
  expect_equal(ans_obtained, ans_expected)
})


## 'model_descr' --------------------------------------------------------------

test_that("'model_descr' works with valid inputs", {
    expect_identical(model_descr(structure(1, class = "bage_mod_pois")), "Poisson")
    expect_identical(model_descr(structure(1, class = "bage_mod_binom")), "binomial")
    expect_identical(model_descr(structure(1, class = "bage_mod_norm")), "normal")
})


## 'n_draw' -------------------------------------------------------------------

test_that("'n_draw' works with mod_pois", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- 3 * rpois(n = nrow(data), lambda = 0.4 * data$popn)
    data$income <- rnorm(n = nrow(data))
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    expect_identical(n_draw(mod), 1000L)
    mod <- set_n_draw(mod, n_draw = 10)
    expect_identical(n_draw(mod), 10L)
})


## 'nm_distn' -----------------------------------------------------------------

test_that("'nm_distn' works with valid inputs", {
    expect_identical(nm_distn(structure(1, class = "bage_mod_pois")), "pois")
    expect_identical(nm_distn(structure(1, class = "bage_mod_binom")), "binom")
    expect_identical(nm_distn(structure(1, class = "bage_mod_norm")), "norm")
})


## 'print' --------------------------------------------------------------------

test_that("'print' works with mod_pois", {
    set.seed(0)
    data <- expand.grid(age = 0:1, time = 2000:2001, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- 3 * rpois(n = nrow(data), lambda = 0.4 * data$popn)
    data$income <- rnorm(n = nrow(data))
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn) |>
      set_confidential_rr3() |>
      set_covariates(~ income) |>
      set_datamod_undercount(prob = data.frame(mean = 0.9, disp = 0.2))
    expect_snapshot(print(mod))
    ## don't use snapshot, since printed option includes timings, which can change
    capture.output(print(fit(mod)), file = NULL)
})

test_that("'print' works with mod_pois - inner-outer fitting method", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- 3 * rpois(n = nrow(data), lambda = 0.4 * data$popn)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    expect_snapshot(print(mod))
    ## don't use snapshot, since printed option includes timings, which can change
    capture.output(print(fit(mod,
                             method = "inner-outer",
                             vars_inner = "age")),
                   file = NULL)
})


## 'replicate_data' -----------------------------------------------------------

test_that("'replicate_data' works with mod_pois - has disp", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 0.4 * data$popn)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- fit(mod)
    set.seed(1)
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    tab <- tapply(ans$deaths, ans$.replicate, sd)
    expect_true(var(tab) > 0)
    ans_fit <- replicate_data(mod, condition_on = "fitted")
    expect_equal(mean(ans_fit$deaths), mean(ans$deaths), tolerance = 0.01)
    set.seed(1)
    ans2 <- replicate_data(mod, condition_on = "expected")
    expect_equal(ans, ans2)
})

test_that("'replicate_data' works with mod_pois - no disp", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rpois(n = nrow(data), lambda = 0.4 * data$popn)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn) |>
      set_disp(mean = 0)
    mod <- fit(mod)
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    tab <- tapply(ans$deaths, ans$.replicate, sd)
    expect_true(var(tab) > 0)
    ans_fit <- replicate_data(mod, condition_on = "fitted")
    expect_equal(mean(ans_fit$deaths), mean(ans$deaths), tolerance = 0.01)
})

test_that("'replicate_data' works with mod_pois, rr3 confidential", {
    set.seed(10)
    data <- expand.grid(age = 0:2, time = 2000:2001, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- 3 * rpois(n = nrow(data), lambda = 0.1 * data$popn)
    formula <- deaths ~ age + sex + time
    mod <- mod_pois(formula = formula,
                    data = data,
                    exposure = popn)
    mod <- set_confidential_rr3(mod)
    mod <- fit(mod)
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    tab <- tapply(ans$deaths, ans$.replicate, sd)
    expect_true(var(tab) > 0)
    ans_fit <- replicate_data(mod, condition_on = "fitted")
    expect_equal(mean(ans_fit$deaths), mean(ans$deaths), tolerance = 0.02)
    expect_true(all(ans_fit$deaths %% 3 == 0))
})

test_that("'replicate_data' works with mod_pois, exposure datamodel", {
  set.seed(10)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- 3 * rpois(n = nrow(data), lambda = 0.1 * data$popn)
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_disp(mean = 0) |>
    set_datamod_exposure(cv = 0.02) |>
    fit()
  ans <- replicate_data(mod)
  expect_identical(names(ans), c(".replicate", names(data)))
  expect_identical(nrow(ans), nrow(data) * 20L)
  tab <- tapply(ans$deaths, ans$.replicate, sd)
  expect_true(var(tab) > 0)
  ans_fit <- replicate_data(mod, condition_on = "fitted")
  expect_equal(mean(ans_fit$deaths), mean(ans$deaths), tolerance = 0.02)
})


test_that("'replicate_data' works with mod_pois, noise datamodel", {
  set.seed(10)
  data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
  data$popn <- rpois(n = nrow(data), lambda = 100)
  data$deaths <- 3 * rpois(n = nrow(data), lambda = 0.1 * data$popn)
  formula <- deaths ~ age + sex + time
  mod <- mod_pois(formula = formula,
                  data = data,
                  exposure = popn) |>
    set_disp(mean = 0) |>
    set_datamod_noise(sd = 2) |>
    fit()
  ans <- replicate_data(mod)
  expect_identical(names(ans), c(".replicate", names(data)))
  expect_identical(nrow(ans), nrow(data) * 20L)
  tab <- tapply(ans$deaths, ans$.replicate, sd)
  expect_true(var(tab) > 0)
  ans_fit <- replicate_data(mod, condition_on = "fitted")
  expect_equal(mean(ans_fit$deaths), mean(ans$deaths), tolerance = 0.02)
})


test_that("'replicate_data' works with mod_binom - has disp", {
    set.seed(0)
    data <- expand.grid(age = 0:29, time = 2000:2002, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
    formula <- deaths ~ age + sex + time
    mod <- mod_binom(formula = formula,
                    data = data,
                    size = popn)
    mod <- fit(mod)
    set.seed(1)
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    tab <- tapply(ans$deaths, ans$.replicate, mean)
    expect_true(var(tab) > 0)
    ans_fit <- replicate_data(mod, condition_on = "fitted")
    expect_equal(mean(ans_fit$deaths), mean(ans$deaths), tolerance = 0.01)
    set.seed(1)
    ans2 <- replicate_data(mod, condition_on = "expected")
    expect_equal(ans, ans2)
})

test_that("'replicate_data' works with mod_binom - no disp", {
    set.seed(0)
    data <- expand.grid(age = 0:29, time = 2000:2002, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
    formula <- deaths ~ age + sex + time
    mod <- mod_binom(formula = formula,
                    data = data,
                    size = popn) |>
      set_disp(mean = 0)
    mod <- fit(mod)
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    tab <- tapply(ans$deaths, ans$.replicate, mean)
    expect_true(var(tab) > 0)
})

test_that("'replicate_data' works with mod_binom, rr3 ", {
    set.seed(0)
    data <- expand.grid(age = 0:29, time = 2000:2002, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- 3 * rbinom(n = nrow(data), size = data$popn, prob = 0.1)
    formula <- deaths ~ age + sex + time
    mod <- mod_binom(formula = formula,
                    data = data,
                    size = popn)
    mod <- set_confidential_rr3(mod)
    mod <- fit(mod)
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    tab <- tapply(ans$deaths, ans$.replicate, mean)
    expect_true(var(tab) > 0)
    ans_fit <- replicate_data(mod, condition_on = "fitted")
    expect_equal(mean(ans_fit$deaths), mean(ans$deaths), tolerance = 0.03)
    expect_true(all(ans_fit$deaths %% 3 == 0))
})

test_that("'replicate_data' works with mod_binom, undercount data model ", {
    set.seed(0)
    data <- expand.grid(age = 0:29, time = 2000:2002, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
    formula <- deaths ~ age + sex + time
    mod <- mod_binom(formula = formula,
                    data = data,
                    size = popn) |>
      set_datamod_undercount(prob = data.frame(mean = 0.9, disp = 0.001)) |>
      fit()
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    tab <- tapply(ans$deaths, ans$.replicate, mean)
    expect_true(var(tab) > 0)
    ans_fit <- replicate_data(mod, condition_on = "fitted")
    expect_equal(mean(ans_fit$deaths), mean(ans$deaths), tolerance = 0.03)
})

test_that("'replicate_data' works with mod_binom, rr3, disp 0", {
    set.seed(0)
    data <- expand.grid(age = 0:2, time = 2000:2002, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- 3 * rbinom(n = nrow(data), size = data$popn, prob = 0.1)
    formula <- deaths ~ age + sex + time
    mod <- mod_binom(formula = formula,
                    data = data,
                    size = popn)
    mod <- set_disp(mod, mean = 0)
    mod <- set_confidential_rr3(mod)
    mod <- fit(mod)
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    tab <- tapply(ans$deaths, ans$.replicate, mean)
    expect_true(var(tab) > 0)
    ans_fit <- replicate_data(mod, condition_on = "fitted")
    expect_equal(mean(ans_fit$deaths), mean(ans$deaths), tolerance = 0.03)
    expect_true(all(ans_fit$deaths %% 3 == 0))
})

test_that("'replicate_data' works with mod_norm", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$income <- rnorm(n = nrow(data), mean = 100, sd = 2)
    data$w <- runif(n = nrow(data), min = 500, max = 1000)
    formula <- income ~ age + sex + time
    mod <- mod_norm(formula = formula,
                    data = data,
                    weights = w)
    mod <- set_prior(mod, age ~ N())
    mod <- set_prior(mod, time ~ N())
    mod <- fit(mod)
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    expect_equal(ans$income[ans$.replicate == "Original"], data$income)
    expect_equal(mean(ans$income[ans$.replicate == "Original"]),
                 mean(ans$income[ans$.replicate == "Replicate 1"]),
                 tolerance = 0.1)
    expect_equal(mad(ans$income[ans$.replicate == "Original"]),
                 mad(ans$income[ans$.replicate == "Replicate 1"]),
                 tolerance = 0.2)
    tab <- tapply(ans$income, ans$.replicate, mean)
    expect_false(any(duplicated(tab)))
    expect_warning(replicate_data(mod, condition_on = "expected"),
                   "Ignoring value for `condition_on`.")
})

test_that("'replicate_data' works with mod_norm, noise data model", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$income <- rnorm(n = nrow(data), mean = 100, sd = 2)
    data$w <- runif(n = nrow(data), min = 500, max = 1000)
    formula <- income ~ age + sex + time
    mod <- mod_norm(formula = formula,
                    data = data,
                    weights = w)
    mod <- set_prior(mod, age ~ N())
    mod <- set_prior(mod, time ~ N())
    mod <- set_datamod_noise(mod, sd = 2)
    mod <- fit(mod)
    ans <- replicate_data(mod)
    expect_identical(names(ans), c(".replicate", names(data)))
    expect_identical(nrow(ans), nrow(data) * 20L)
    expect_equal(ans$income[ans$.replicate == "Original"], data$income)
    expect_equal(mean(ans$income[ans$.replicate == "Original"]),
                 mean(ans$income[ans$.replicate == "Replicate 1"]),
                 tolerance = 0.1)
    expect_equal(mad(ans$income[ans$.replicate == "Original"]),
                 mad(ans$income[ans$.replicate == "Replicate 1"]),
                 tolerance = 0.2)
    tab <- tapply(ans$income, ans$.replicate, mean)
    expect_false(any(duplicated(tab)))
    expect_warning(replicate_data(mod, condition_on = "expected"),
                   "Ignoring value for `condition_on`.")
})


## 'tidy' ---------------------------------------------------------------------

test_that("'tidy' works with valid inputs", {
    set.seed(0)
    data <- expand.grid(age = 0:9, time = 2000:2005, sex = c("F", "M"))
    data$popn <- rpois(n = nrow(data), lambda = 100)
    data$deaths <- rbinom(n = nrow(data), size = data$popn, prob = 0.3)
    formula <- deaths ~ age * sex
    mod <- mod_binom(formula = formula,
                    data = data,
                    size = popn)
    ans_unfit <- tidy(mod)
    expect_true(is.data.frame(ans_unfit))
    expect_identical(names(ans_unfit), c("term", "prior", "along",
                                         "n_par", "n_par_free"))
    mod_fitted <- fit(mod)
    ans_fitted <- tidy(mod_fitted)
    expect_true(is.data.frame(ans_fitted))
    expect_identical(names(ans_fitted), c("term", "prior", "along",
                                          "n_par", "n_par_free",
                                          "std_dev"))
})