# =============================================================================
# Tests for pretrained learner support
# =============================================================================

# -----------------------------------------------------------------------------
# assert_pretrained() unit tests
# -----------------------------------------------------------------------------

test_that("assert_pretrained returns FALSE for untrained learner", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  resampling = rsmp("holdout")$instantiate(task)

  expect_false(assert_pretrained(learner, task, resampling))
})

test_that("assert_pretrained returns TRUE for valid pretrained setup", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  learner$train(task)
  resampling = rsmp("holdout")$instantiate(task)

  expect_true(assert_pretrained(learner, task, resampling))
})

test_that("assert_pretrained errors when trained learner has multi-fold resampling", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  learner$train(task)
  resampling = rsmp("cv", folds = 3)$instantiate(task)

  expect_error(assert_pretrained(learner, task, resampling), "not compatible")
})

# -----------------------------------------------------------------------------
# PFI with pretrained learner
# -----------------------------------------------------------------------------

test_that("PFI works with pretrained learner", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  resampling = rsmp("holdout")$instantiate(task)
  measure = msr("regr.mse")

  learner$train(task)

  pfi = PFI$new(
    task = task,
    learner = learner,
    measure = measure,
    resampling = resampling,
    n_repeats = 3L
  )

  pfi$compute()
  expect_importance_dt(pfi$importance(), features = pfi$features)
  checkmate::expect_r6(pfi$resample_result, "ResampleResult")

  # Original learner must still have its model after compute
  expect_false(is.null(learner$model))
})

test_that("PFI pretrained and non-pretrained produce comparable results", {
  set.seed(123)
  task = tsk("mtcars")
  resampling = rsmp("holdout")$instantiate(task)
  measure = msr("regr.mse")

  # Pretrained path
  learner_pre = lrn("regr.rpart")
  learner_pre$train(task)

  set.seed(1)
  pfi_pre = PFI$new(
    task = task,
    learner = learner_pre,
    measure = measure,
    resampling = resampling,
    n_repeats = 5L
  )
  pfi_pre$compute()

  # Non-pretrained path (fresh learner, same resampling)
  set.seed(1)
  pfi_fresh = PFI$new(
    task = task,
    learner = lrn("regr.rpart"),
    measure = measure,
    resampling = resampling,
    n_repeats = 5L
  )
  pfi_fresh$compute()

  # Both should produce valid importance tables with same features
  expect_importance_dt(pfi_pre$importance(), features = pfi_pre$features)
  expect_importance_dt(pfi_fresh$importance(), features = pfi_fresh$features)
  expect_setequal(pfi_pre$importance()$feature, pfi_fresh$importance()$feature)
})

test_that("PFI with pretrained learner errors on non-instantiated resampling", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  learner$train(task)
  measure = msr("regr.mse")

  expect_error(
    PFI$new(
      task = task,
      learner = learner,
      measure = measure,
      resampling = rsmp("holdout")
    ),
    "instantiated"
  )
})

test_that("PFI with pretrained learner errors on multi-fold resampling at construction", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  learner$train(task)
  measure = msr("regr.mse")

  expect_error(
    PFI$new(
      task = task,
      learner = learner,
      measure = measure,
      resampling = rsmp("cv", folds = 3)$instantiate(task),
      n_repeats = 2L
    ),
    "not compatible"
  )
})

test_that("PFI with pretrained learner works for classification", {
  task = tsk("penguins")
  learner = lrn("classif.rpart", predict_type = "prob")
  resampling = rsmp("holdout")$instantiate(task)
  measure = msr("classif.ce")

  learner$train(task)

  pfi = PFI$new(
    task = task,
    learner = learner,
    measure = measure,
    resampling = resampling,
    n_repeats = 3L
  )

  pfi$compute()
  expect_importance_dt(pfi$importance(), features = pfi$features)
})

# -----------------------------------------------------------------------------
# MarginalSAGE with pretrained learner
# -----------------------------------------------------------------------------

test_that("MarginalSAGE works with pretrained learner", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  resampling = rsmp("holdout")$instantiate(task)
  measure = msr("regr.mse")

  learner$train(task)

  sage = MarginalSAGE$new(
    task = task,
    learner = learner,
    measure = measure,
    resampling = resampling,
    n_permutations = 2L,
    n_samples = 20L
  )

  sage$compute()
  expect_importance_dt(sage$importance(), features = sage$features)
  checkmate::expect_r6(sage$resample_result, "ResampleResult")
})

test_that("MarginalSAGE pretrained and non-pretrained produce comparable results", {
  set.seed(123)
  task = tsk("mtcars")
  resampling = rsmp("holdout")$instantiate(task)
  measure = msr("regr.mse")

  # Pretrained path
  learner_pre = lrn("regr.rpart")
  learner_pre$train(task)

  set.seed(1)
  sage_pre = MarginalSAGE$new(
    task = task,
    learner = learner_pre,
    measure = measure,
    resampling = resampling,
    n_permutations = 3L,
    n_samples = 20L
  )
  sage_pre$compute()

  # Non-pretrained path
  set.seed(1)
  sage_fresh = MarginalSAGE$new(
    task = task,
    learner = lrn("regr.rpart"),
    measure = measure,
    resampling = resampling,
    n_permutations = 3L,
    n_samples = 20L
  )
  sage_fresh$compute()

  expect_importance_dt(sage_pre$importance(), features = sage_pre$features)
  expect_importance_dt(sage_fresh$importance(), features = sage_fresh$features)
  expect_setequal(sage_pre$importance()$feature, sage_fresh$importance()$feature)
})

test_that("MarginalSAGE with pretrained learner errors on non-instantiated resampling", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  learner$train(task)
  measure = msr("regr.mse")

  expect_error(
    MarginalSAGE$new(
      task = task,
      learner = learner,
      measure = measure,
      resampling = rsmp("holdout"),
      n_permutations = 2L,
      n_samples = 20L
    ),
    "instantiated"
  )
})

test_that("MarginalSAGE with pretrained learner errors on multi-fold resampling at construction", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  learner$train(task)
  measure = msr("regr.mse")

  expect_error(
    MarginalSAGE$new(
      task = task,
      learner = learner,
      measure = measure,
      resampling = rsmp("cv", folds = 3)$instantiate(task),
      n_permutations = 2L,
      n_samples = 20L
    ),
    "not compatible"
  )
})

test_that("MarginalSAGE with pretrained learner works for classification", {
  task = tsk("penguins")
  learner = lrn("classif.rpart", predict_type = "prob")
  resampling = rsmp("holdout")$instantiate(task)
  measure = msr("classif.ce")

  learner$train(task)

  sage = MarginalSAGE$new(
    task = task,
    learner = learner,
    measure = measure,
    resampling = resampling,
    n_permutations = 2L,
    n_samples = 20L
  )

  sage$compute()
  expect_importance_dt(sage$importance(), features = sage$features)
})

# -----------------------------------------------------------------------------
# LOCO warns for pretrained learner (refit-based methods)
# -----------------------------------------------------------------------------

test_that("LOCO warns when given a pretrained learner", {
  task = tsk("mtcars")
  learner = lrn("regr.rpart")
  learner$train(task)

  expect_warning(
    LOCO$new(
      task = task,
      learner = learner,
      measure = msr("regr.mse")
    ),
    "already trained"
  )
})
