# These tests reference results generated by tune 1.3.0. The code to generate
# them (and the results) are found in the `inst` directory.

test_that("verifying .loop_over_all_stages, no submodels, post estimation without tuning", {
  skip_if_not_installed("modeldata")
  skip_if_not_installed("kknn")
  skip_if_not_installed("probably")
  skip_if_not_installed("mgcv")

  load(system.file(
    "regression_tests",
    "simple_example.RData",
    package = "tune"
  ))

  # ------------------------------------------------------------------------------

  set.seed(1)
  dat <- modeldata::sim_regression(1000)
  rs <- vfold_cv(dat)

  rs_split <- rs$splits[[1]]
  rs_args <- rsample::.get_split_args(rs)

  rs_iter <- tune:::vec_list_rowwise(rs) |>
    purrr::pluck(1) |>
    mutate(
      .seeds = tune:::get_parallel_seeds(1)
    )

  # ------------------------------------------------------------------------------

  mod <- nearest_neighbor(neighbors = 11, weight_func = tune()) |>
    set_mode("regression")
  wflow <- workflow(outcome ~ ., mod, reg_cal)
  grd <- wflow |> extract_parameter_set_dials() |> grid_regular()

  static_1 <- tune:::.make_static(
    wflow,
    param_info = wflow |> extract_parameter_set_dials(),
    grid = grd,
    metrics = metric_set(rmse, rsq),
    eval_time = NULL,
    split_args = rs_args,
    control = control_grid(save_pred = TRUE)
  )

  data_1 <- tune:::.get_data_subsets(wflow, rs_split, rs_args)
  static_1 <- tune:::update_static(static_1, data_1)
  static_1$y_name <- "outcome"

  simple_res <- tune:::.loop_over_all_stages(rs_iter, grd, static_1)
  expect_true(!is.null(simple_res$.metrics[[1]]))
  expect_named(
    simple_res,
    c(".metrics", ".notes", "outcome_names", "id", ".predictions")
  )
  expect_true(nrow(simple_res) == 1)
  expect_equal(
    nrow(simple_res$.predictions[[1]]),
    nrow(data_1$pred$data) * nrow(grd)
  )

  # TODO more tests can be added when calibration method = "none" is implemented
})

test_that("verifying .loop_over_all_stages, submodels, post estimation without tuning", {
  skip_if_not_installed("modeldata")
  skip_if_not_installed("kknn")
  skip_if_not_installed("probably")
  skip_if_not_installed("mgcv")

  load(system.file(
    "regression_tests",
    "submodel_example.RData",
    package = "tune"
  ))

  ctrl <- tune::control_grid(save_pred = TRUE)

  # ------------------------------------------------------------------------------

  set.seed(1)
  dat <- modeldata::sim_regression(1000)
  rs <- vfold_cv(dat)

  rs_split <- rs$splits[[1]]
  rs_args <- rsample::.get_split_args(rs)

  rs_iter <- tune:::vec_list_rowwise(rs) |>
    purrr::pluck(1) |>
    mutate(
      .seeds = tune:::get_parallel_seeds(1)
    )

  # ------------------------------------------------------------------------------

  rec <- recipe(outcome ~ ., data = dat) |>
    step_pca(all_numeric_predictors(), num_comp = tune())

  mod <- nearest_neighbor(neighbors = tune("k"), weight_func = tune()) |>
    set_mode("regression")

  submodel_wflow <- workflow(rec, mod, reg_cal)

  # fmt: skip
  submodel_grid <-
    tibble::tribble(
      ~k,   ~weight_func, ~num_comp,
      9L,  "rectangular",        2L,
      14L,  "rectangular",        2L,
      20L,  "rectangular",        2L,
      4L,   "triangular",        2L,
      9L,   "triangular",        2L,
      14L,   "triangular",        2L,
      20L,   "triangular",        2L,
      4L, "epanechnikov",        2L,
      9L, "epanechnikov",        2L,
      14L, "epanechnikov",        2L,
      20L, "epanechnikov",        2L,
      4L,  "rectangular",       10L,
      9L,  "rectangular",       10L,
      14L,  "rectangular",       10L,
      20L,  "rectangular",       10L,
      4L,   "triangular",       10L,
      9L,   "triangular",       10L,
      14L,   "triangular",       10L,
      20L,   "triangular",       10L,
      4L, "epanechnikov",       10L,
      9L, "epanechnikov",       10L,
      14L, "epanechnikov",       10L,
      20L, "epanechnikov",       10L
    )

  # ------------------------------------------------------------------------------

  static_1 <- tune:::.make_static(
    submodel_wflow,
    param_info = submodel_wflow |> extract_parameter_set_dials(),
    grid = submodel_grid,
    metrics = metric_set(rmse, rsq),
    eval_time = NULL,
    split_args = rs_args,
    control = ctrl
  )

  data_1 <- tune:::.get_data_subsets(submodel_wflow, rs_split, rs_args)
  static_1 <- tune:::update_static(static_1, data_1)
  static_1$y_name <- "outcome"

  submodel_res <- tune:::.loop_over_all_stages(rs_iter, submodel_grid, static_1)
  expect_true(!is.null(submodel_res$.metrics[[1]]))
  expect_named(
    submodel_res,
    c(".metrics", ".notes", "outcome_names", "id", ".predictions")
  )
  expect_true(nrow(submodel_res) == 1)
  expect_equal(
    nrow(submodel_res$.predictions[[1]]),
    nrow(data_1$pred$data) * nrow(submodel_grid)
  )
})

test_that("verifying .loop_over_all_stages, submodels only, post estimation without tuning", {
  skip_if_not_installed("modeldata")
  skip_if_not_installed("kknn")
  skip_if_not_installed("probably")
  skip_if_not_installed("mgcv")

  load(system.file(
    "regression_tests",
    "submodel_only_example.RData",
    package = "tune"
  ))

  ctrl <- tune::control_grid(save_pred = TRUE)

  # ------------------------------------------------------------------------------

  set.seed(1)
  dat <- modeldata::sim_classification(1000)
  rs <- vfold_cv(dat)

  rs_split <- rs$splits[[1]]
  rs_args <- rsample::.get_split_args(rs)

  rs_iter <- tune:::vec_list_rowwise(rs) |>
    purrr::pluck(1) |>
    mutate(
      .seeds = tune:::get_parallel_seeds(1)
    )

  # ------------------------------------------------------------------------------

  mod <- nearest_neighbor(neighbors = tune(), weight_func = "triangular") |>
    set_mode("classification")

  submodel_only_wflow <- workflow(class ~ ., mod, cls_est_post)

  submodel_only_grid <- tibble(neighbors = 3:10)

  # ------------------------------------------------------------------------------

  static_1 <- tune:::.make_static(
    submodel_only_wflow,
    param_info = submodel_only_wflow |> extract_parameter_set_dials(),
    grid = submodel_only_grid,
    metrics = metric_set(accuracy, roc_auc, brier_class),
    eval_time = NULL,
    split_args = rs_args,
    control = ctrl
  )

  data_1 <- tune:::.get_data_subsets(submodel_only_wflow, rs_split, rs_args)
  static_1 <- tune:::update_static(static_1, data_1)
  static_1$y_name <- "class"

  submodel_only_res <- tune:::.loop_over_all_stages(
    rs_iter,
    submodel_only_grid,
    static_1
  )
  expect_true(!is.null(submodel_only_res$.metrics[[1]]))
  expect_named(
    submodel_only_res,
    c(".metrics", ".notes", "outcome_names", "id", ".predictions")
  )
  expect_true(nrow(submodel_only_res) == 1)
  expect_equal(
    nrow(submodel_only_res$.predictions[[1]]),
    nrow(data_1$pred$data) * nrow(submodel_only_grid)
  )

  # TODO more tests can be added when calibration method = "none" is implemented
})

test_that("submodel calibration varies per submodel value (#1144)", {
  skip_if_not_installed("modeldata")
  skip_if_not_installed("glmnet")
  skip_if_not_installed("probably")
  skip_if_not_installed("mgcv")

  set.seed(1)
  dat <- modeldata::sim_regression(500)
  rs <- vfold_cv(dat, v = 3)

  rs_split <- rs$splits[[1]]
  rs_args <- rsample::.get_split_args(rs)

  rs_iter <- vec_list_rowwise(rs) |>
    purrr::pluck(1) |>
    mutate(.seeds = get_parallel_seeds(1))

  wflow <- workflow(outcome ~ ., glmn_spec, reg_cal)
  grd <- tibble(penalty = c(0.001, 0.1), mixture = c(1, 1))

  extract_cal_coef <- function(wf) {
    postprocessor <- extract_postprocessor(wf, estimated = TRUE)
    postprocessor$adjustments[[1]]$results$fit$estimates[[1]]$estimate |>
      coefficients()
  }

  static <- .make_static(
    wflow,
    param_info = extract_parameter_set_dials(wflow),
    grid = grd,
    metrics = metric_set(rmse),
    eval_time = NULL,
    split_args = rs_args,
    control = control_grid(
      save_pred = TRUE,
      extract = extract_cal_coef
    )
  )

  data_splits <- .get_data_subsets(wflow, rs_split, rs_args)
  static <- update_static(static, data_splits)
  static$y_name <- "outcome"

  res <- .loop_over_all_stages(rs_iter, grd, static)
  extracts <- res$.extracts[[1]]

  loop_cal_p1 <- extracts$.extracts[[1]]
  loop_cal_p2 <- extracts$.extracts[[2]]

  # The two calibration models should differ because different penalty values
  # produce different predictions on the calibration set.
  expect_false(identical(loop_cal_p1, loop_cal_p2))

  withr::local_preserve_seed()
  assign(".Random.seed", rs_iter$.seeds[[1]], envir = .GlobalEnv)
  inner_split <- rsample::internal_calibration_split(
    rs_split,
    split_args = rs_args
  )

  ref_wf <- workflow(
    outcome ~ .,
    linear_reg(penalty = 0.001, mixture = 1, engine = "glmnet"),
    reg_cal
  ) |>
    fit(
      data = analysis(inner_split),
      data_calibration = calibration(inner_split)
    )

  ref_cal <- extract_cal_coef(ref_wf)

  expect_equal(loop_cal_p1, ref_cal)
})
