test_estimate <- function() {
  test_data <- data.frame(
    biomass = c(0, 11, 8, 10, 0, 12, 14, 9, 0, 13, 8, 10, 0, 11, 16, 19),
    ppt = c(2, 1, 2, 3, 2, 2, 3, 1, 1, 3, 1, 2, 0, 1, 2, 3),
    county = factor(rep(c("A", "B", "C", "D"), each = 4))
  )

  spec1 <- basal::specify(
    biomass ~ ppt + (1 | county),
    level = "unit"
  )

  spec2 <- basal::specify(
    biomass ~ ppt + (1 | county),
    level = "unit",
    model_stage = "zi"
  )

  for (engine in list(
    basal::engine_rstanarm(),
    basal::engine_brms()
  )) {
    args <- list(
      object = spec1,
      data = test_data,
      chains = 1,
      iter = 20,
      burn_in = 5,
      thin = 1,
      seed = 1,
      ncores = 1,
      nthreads = 1,
      engine = engine,
      refresh = 0
    )
    fit1 <- do.call(
      basal::fit,
      args
    )

    args$object <- spec2

    fit2 <- do.call(
      basal::fit,
      args
    )

    est_data <- test_data
    est_data$dummy <- 1

    for (fit in list(
      fit1, fit2
    )) {
      testthat::expect_no_error(
        basal::estimate(
          fit,
          ndraws = 10,
          max_preds = 10,
          seed = 1
        )
      )

      testthat::expect_no_error(
        basal::estimate(
          fit,
          newdata = est_data,
          ndraws = 10,
          max_preds = 10,
          seed = 1
        )
      )

      testthat::expect_no_error(
        basal::estimate(
          fit,
          newdata = est_data,
          domain = "dummy",
          ndraws = 10,
          max_preds = 10,
          seed = 1
        )
      )

      testthat::expect_no_error(
        basal::estimate(
          fit,
          newdata = est_data,
          domain = NULL,
          ndraws = 10,
          max_preds = 10,
          seed = 1
        )
      )

      testthat::expect_no_error(
        basal::estimate(
          fit,
          stat = c(mean = mean, median = median),
          aggregation_statistic = c(median = median),
          ndraws = 10,
          max_preds = 10,
          seed = 1
        )
      )

      testthat::expect_error(
        basal::estimate(
          fit,
          domain = "state",
          ndraws = 10,
          max_preds = 10
        )
        # not present in newdata
      )
    }
  }
}

testthat::test_that("estimate handles basic arguments and domains", {
  testthat::skip_on_cran()
  test_estimate()
})

big_test_estimate <- function() {
  test_data <- data.frame(
    biomass = c(0, 11, 8, 10, 0, 12, 14, 9, 0, 13, 8, 10, 0, 11, 16, 19),
    ppt = c(2, 1, 2, 3, 2, 2, 3, 1, 1, 3, 1, 2, 0, 1, 2, 3),
    county = factor(rep(c("A", "B", "C", "D"), each = 4)),
    d = factor(rep(c("A", "B", "C", "D", "E", "F", "G", "H"), each = 2))
  )

  agg_data <- basal::aggregate(
    test_data, "biomass", "county", test_data, 1e6
  )

  custom_args <- list(
    formula = biomass ~ ppt,
    level = NULL
  )

  manual_args <- list(
    response_name = "biomass",
    auxiliary_variables = "ppt"
  )

  for (base_args in list(custom_args, manual_args)) {
    for (model_type in c("custom", "BHF", "FH")) {
      args <- base_args
      if (model_type != "custom") {
        domain_name <- "county"
      } else {
        domain_name <- NULL
        args$formula <- biomass ~ ppt + (1 | county)
      }
      for (transform in list(
        NULL,
        basal::transform(function(x) {x^2}, sqrt)
      )) {
        for (model_stage in c("single", "zi")) {
          if (model_type == "custom" || model_type == "BHF") {
            level <- "unit" # level doesn't have to be unit for custom models
          } else {
            level <- "area"
          }
          if (level == "area" || model_type == "FH") {
            thedata <- agg_data
          } else {
            thedata <- test_data
          }
          domain_name <- if (model_type != "custom" || level == "area") {
            "county"
          } else {
            NULL
          }
          expecting_error <- FALSE
          if ((level == "area" || model_type == "FH") &&
                model_stage == "zi") {
            expecting_error <- TRUE
          } else if (model_type == "custom" && is.null(args$formula)) {
            expecting_error <- TRUE
          } else if (model_type == "custom" && !is.null(domain_name)) {
            expecting_error <- TRUE
          } else if ((model_type == "FH" || level == "area") &&
                       !is.null(transform)) {
            expecting_error <- TRUE
          } else if (!is.null(args$response_name) && !is.null(args$formula)) {
            expecting_error <- TRUE
          }
          if (expecting_error) {
            next
          } else {
            spec <- basal::specify(
              formula = args$formula,
              response_name = args$response_name,
              auxiliary_variables = args$auxiliary_variables,
              model = model_type,
              domain_name = domain_name,
              variable_transform = transform,
              model_stage = model_stage,
              level = level
            )
            for (engine in list(
              basal::engine_rstanarm(),
              basal::engine_brms()
            )) {
              if (!(level %in% engine$level &&
                      model_type %in% engine$model &&
                      model_stage %in% engine$model_stage)) {
                next
              }
              model <- (
                suppressWarnings(basal::fit(
                  spec,
                  thedata,
                  chains = 1,
                  burn_in = 2,
                  thin = 1,
                  iter = 4,
                  ncores = 1,
                  nthreads = 1,
                  engine = engine,
                  refresh = 0
                ))
              ) |> try()
          
              for (stat in list(
                list(mean = mean),
                c(mean = mean, var = var, med = median)
              )) {
                for (agg_stat in list(
                  list(mean = mean), list(sd = sd), list(med = median)
                )) {
                  for (domain in list(NULL, "d", "county")) {
                    if (level == "area") {
                      thenewdata = NULL
                    } else {
                      thenewdata = thedata
                    }
                    t <- testthat::expect_no_error(suppressWarnings(
                      basal::estimate(
                        model,
                        newdata = thenewdata, ndraws = 2,
                        stat = stat, aggregation_statistic = agg_stat,
                        domain = domain
                      )
                    )) |> try()
                    if (inherits(t, "try-error")) {
                      browser()
                    }
                  }
                }
              }
            }
          }
        }
      }
    }
  }
}

testthat::test_that("more estimate tests work as normal", {
  testthat::skip_on_cran()
  big_test_estimate()
})

testthat::test_that("estimate returns expected domain information", {
  test_data <- data.frame(
    biomass = c(10, 12, 15, 9, 18, 21, 13, 17, 20, 11, 16, 19),
    ppt = c(1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3),
    county = factor(rep(c("A", "B", "C", "D"), each = 3))
  )

  spec <- basal::specify(
    model = "BHF",
    response_name = "biomass",
    auxiliary_variables = "ppt",
    domain_name = "county"
  )

  fit <- basal::fit(
    spec,
    data = test_data,
    chains = 1,
    iter = 20,
    burn_in = 5,
    thin = 1,
    seed = 1,
    ncores = 1,
    nthreads = 1,
    engine = basal::engine_rstanarm(),
    refresh = 0
  )

  result <- basal::estimate(
    fit,
    domain = "county",
    stat = c(mean = mean),
    aggregation_statistic = c(mean = mean),
    ndraws = 10,
    max_preds = 10,
    seed = 1
  )

  testthat::expect_s3_class(result, "basal_estimate")
  testthat::expect_equal(result$params$domain, "county")
  testthat::expect_equal(result$params$ndraws, 10)
  testthat::expect_equal(result$params$max_preds, 10)
  testthat::expect_true("county" %in% names(result$preds))
})

testthat::test_that("Estimate plays well with area-level models", {
  testthat::skip_on_cran()

  test_data <- data.frame(
    biomass = c(10, 12, 15, 9, 18, 21, 13, 17, 20, 11, 16, 19),
    ppt = c(1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3),
    county = factor(rep(c("A", "B", "C", "D"), each = 3))
  )

  test_data <- basal::aggregate(test_data, "biomass", "county", test_data, 1e4)

  fit <- basal::fit(
    basal::specify(biomass ~ ppt + (1 | county), level = "area"),
    test_data,
    iter = 10,
    burn_in = 5,
    chains = 1
  )

  testthat::expect_no_error(suppressWarnings(basal::estimate(fit, ndraws = 2)))
})