# Tests for KnockoffSampler

test_that("KnockoffSampler initialization works", {
  skip_if_not_installed("knockoff")

  n = 100
  task = tgen("friedman1")$generate(n = n)
  sampler = KnockoffSampler$new(task)

  expect_s3_class(sampler, "KnockoffSampler")
  expect_equal(sampler$label, "Knockoff sampler")
  expect_s3_class(sampler$param_set, "ParamSet")
  expect_true("knockoff_fun" %in% sampler$param_set$ids())

  # Check that knockoff matrix was created
  expect_true(data.table::is.data.table(sampler$x_tilde))
  expect_equal(nrow(sampler$x_tilde), n)
  # x_tilde has feature columns + internal ..row_id column
  expect_equal(ncol(sampler$x_tilde), length(task$feature_names) + 1L)
  expect_equal(setdiff(names(sampler$x_tilde), "..row_id"), task$feature_names)
})

test_that("KnockoffSampler sampling works", {
  skip_if_not_installed("knockoff")

  task = tgen("friedman1")$generate(n = 100)
  sampler = KnockoffSampler$new(task)
  data = task$data()

  # Test single feature sampling
  sampled_data = sampler$sample("important1")

  expect_sampler_output_structure(sampled_data, task, nrows = 100)
  expect_feature_type_consistency(sampled_data, task)

  # Sampled feature values are from knockoff matrix
  expect_identical(sampled_data$important1, sampler$x_tilde$important1)
})

test_that("KnockoffSampler handles multiple features", {
  skip_if_not_installed("knockoff")

  task = tgen("friedman1")$generate(n = 100)
  sampler = KnockoffSampler$new(task)
  data = task$data()

  features = c("important1", "important2")
  sampled_data = sampler$sample(features)

  expect_sampler_output_structure(sampled_data, task, nrows = 100)
  expect_feature_type_consistency(sampled_data, task)

  # Check that sampled feature values are from knockoff matrix
  for (feat in features) {
    expect_identical(sampled_data[[feat]], sampler$x_tilde[[feat]])
  }
})

test_that("KnockoffSampler works with different numeric tasks", {
  skip_if_not_installed("knockoff")

  # Friedman1 regression task (all numeric)
  task_friedman = tgen("friedman1")$generate(n = 100)
  sampler_friedman = KnockoffSampler$new(task_friedman)
  sampled_friedman = sampler_friedman$sample("important1")
  expect_sampler_output_structure(sampled_friedman, task_friedman, nrows = 100)

  # Circle task with specified dimensions (all numeric)
  task_circle = tgen("circle", d = 4)$generate(n = 80)
  sampler_circle = KnockoffSampler$new(task_circle)
  sampled_circle = sampler_circle$sample("x1")
  expect_sampler_output_structure(sampled_circle, task_circle, nrows = 80)

  # Custom numeric task
  set.seed(123)
  custom_data = data.table::data.table(
    x1 = rnorm(50),
    x2 = runif(50),
    x3 = rexp(50),
    y = rnorm(50)
  )
  task_custom = as_task_regr(custom_data, target = "y")
  sampler_custom = KnockoffSampler$new(task_custom)
  sampled_custom = sampler_custom$sample("x1")
  expect_sampler_output_structure(sampled_custom, task_custom, nrows = 50)
})

test_that("KnockoffSampler knockoff matrix properties", {
  skip_if_not_installed("knockoff")

  task = tgen("friedman1")$generate(n = 100)
  sampler = KnockoffSampler$new(task)

  # Knockoff matrix should have same dimensions as feature matrix (plus ..row_id column)
  feature_data = task$data(cols = task$feature_names)
  expect_equal(nrow(sampler$x_tilde), nrow(feature_data))
  expect_equal(ncol(sampler$x_tilde), ncol(feature_data) + 1L)
  expect_equal(setdiff(names(sampler$x_tilde), "..row_id"), names(feature_data))

  # Knockoff values should be numeric (since task is all numeric)
  expect_true(all(sapply(sampler$x_tilde, is.numeric)))

  # Knockoff values should be finite
  expect_true(all(sapply(sampler$x_tilde, function(x) all(is.finite(x)))))
})

test_that("KnockoffSampler custom knockoff function", {
  skip_if_not_installed("knockoff")

  task = tgen("friedman1")$generate(n = 50)

  # Custom knockoff function that returns scaled versions
  custom_knockoff_fun = function(x) {
    x_matrix = as.matrix(x)
    noise = matrix(rnorm(prod(dim(x_matrix)), sd = 0.1), nrow = nrow(x_matrix))
    return(x_matrix + noise)
  }

  sampler = KnockoffSampler$new(task, knockoff_fun = custom_knockoff_fun)

  expect_true(is.function(sampler$param_set$values$knockoff_fun))
  expect_true(data.table::is.data.table(sampler$x_tilde))
  expect_equal(nrow(sampler$x_tilde), 50)
  expect_equal(ncol(sampler$x_tilde), length(task$feature_names) + 1L)

  data = task$data()
  sampled_data = sampler$sample("important1")

  expect_true(data.table::is.data.table(sampled_data))
  expect_equal(nrow(sampled_data), 50)
  expect_false(identical(sampled_data$important1, data$important1))
  expect_identical(sampled_data$important1, sampler$x_tilde$important1)
})

test_that("KnockoffSampler reproducibility", {
  skip_if_not_installed("knockoff")

  task = tgen("friedman1")$generate(n = 80)

  set.seed(123)
  sampler1 = KnockoffSampler$new(task)

  set.seed(123)
  sampler2 = KnockoffSampler$new(task)

  # Dimensions should be identical
  expect_equal(dim(sampler1$x_tilde), dim(sampler2$x_tilde))
  expect_equal(names(sampler1$x_tilde), names(sampler2$x_tilde))
  expect_true("..row_id" %in% names(sampler1$x_tilde))

  # Sampling should be consistent within each sampler
  sampled1a = sampler1$sample("important1")
  sampled1b = sampler1$sample("important1")
  expect_identical(sampled1a$important1, sampled1b$important1)
})

test_that("KnockoffSampler edge cases", {
  skip_if_not_installed("knockoff")

  # Test with minimum viable task (single feature)
  # Seeded: knockoff::create.second_order warns only for some random samples,
  # so pin the RNG to a state that reliably triggers the expected warning.
  set.seed(1L)
  single_feature_data = data.table::data.table(
    x1 = rnorm(30),
    y = rnorm(30)
  )
  task_single = as_task_regr(single_feature_data, target = "y")

  expect_warning({
    sampler_single = KnockoffSampler$new(task_single)
  })

  expect_equal(ncol(sampler_single$x_tilde), 2L) # x1 + ..row_id
  expect_equal(setdiff(names(sampler_single$x_tilde), "..row_id"), "x1")

  sampled_single = sampler_single$sample("x1")
  expect_true(data.table::is.data.table(sampled_single))
  expect_equal(nrow(sampled_single), 30)
})

test_that("KnockoffSampler fails with non-numeric features", {
  skip_if_not_installed("knockoff")

  # Factor features
  data_with_factor = data.table::data.table(
    x1 = rnorm(50),
    x2 = factor(sample(c("A", "B", "C"), 50, replace = TRUE)),
    x3 = rnorm(50),
    y = rnorm(50)
  )
  task_factor = as_task_regr(data_with_factor, target = "y")
  expect_error(KnockoffSampler$new(task_factor), class = "error")

  # Character features
  data_with_char = data.table::data.table(
    x1 = rnorm(50),
    x2 = sample(letters[1:5], 50, replace = TRUE),
    x3 = rnorm(50),
    y = rnorm(50)
  )
  task_char = as_task_regr(data_with_char, target = "y")
  expect_error(KnockoffSampler$new(task_char), class = "error")
})

test_that("KnockoffGaussianSampler works with all-numeric features", {
  skip_if_not_installed("knockoff")

  task = tgen("friedman1")$generate(n = 100)
  sampler = KnockoffGaussianSampler$new(task)

  expect_s3_class(sampler, "KnockoffGaussianSampler")
  expect_s3_class(sampler, "KnockoffSampler")
  expect_equal(sampler$label, "Gaussian Knockoff sampler")

  expect_true(data.table::is.data.table(sampler$x_tilde))
  expect_equal(nrow(sampler$x_tilde), 100)
  expect_equal(ncol(sampler$x_tilde), length(task$feature_names) + 1L)

  data = task$data()
  sampled_data = sampler$sample("important1")

  expect_sampler_output_structure(sampled_data, task, nrows = 100)
  expect_identical(sampled_data$important1, sampler$x_tilde$important1)

  # Test row_ids
  expect_equal(nrow(sampler$sample("important1", row_ids = 1:10)), 10)
  expect_equal(nrow(sampler$sample("important1", row_ids = 32)), 1)

  # Non-sampled features unchanged
  non_sampled = c(
    "important2",
    "important3",
    "important4",
    "important5",
    "unimportant1",
    "unimportant2",
    "unimportant3",
    "unimportant4",
    "unimportant5"
  )
  expect_equal(
    sampler$sample("important1", row_ids = 32)[, ..non_sampled],
    data[32, ..non_sampled]
  )

  # Multiple features
  features = c("important1", "important2", "important3")
  sampled_multi = sampler$sample(features)

  for (feat in features) {
    expect_false(identical(sampled_multi[[feat]], data[[feat]]))
    expect_identical(sampled_multi[[feat]], sampler$x_tilde[[feat]])
  }
})

test_that("KnockoffGaussianSampler fails with factor features", {
  skip_if_not_installed("knockoff")

  data_with_factor = data.table::data.table(
    x1 = rnorm(50),
    x2 = factor(sample(c("A", "B", "C"), 50, replace = TRUE)),
    x3 = rnorm(50),
    y = rnorm(50)
  )
  task_factor = as_task_regr(data_with_factor, target = "y")

  expect_error(
    KnockoffGaussianSampler$new(task_factor),
    regexp = "following unsupported feature types"
  )
})

test_that("KnockoffSampler works with multiple iterations", {
  skip_if_not_installed("knockoff")

  task = sim_dgp_ewald(n = 100)

  withr::with_seed(42, {
    sampler_default = KnockoffSampler$new(task = task)
  })
  withr::with_seed(42, {
    sampler_1 = KnockoffSampler$new(task = task, iters = 1)
  })
  sampler_5 = KnockoffSampler$new(task = task, iters = 5)

  # x_tilde generated with same seed, same length
  expect_equal(
    sampler_default$sample("x1"),
    sampler_1$sample("x1")
  )

  x_orig = task$data(rows = 2)
  x_sampled = sampler_1$sample("x1", row_ids = 2)

  # Other values are the same (excluding sampled feature x1)
  non_sampled_cols = setdiff(names(x_sampled), "x1")
  expect_equal(x_orig[, ..non_sampled_cols], x_sampled[, ..non_sampled_cols])

  # Multiple samples via samples_per_row
  x_orig = task$data(rows = c(1, 1, 1))
  x_sampled = sampler_5$sample("x1", row_ids = 1, samples_per_row = 3L)

  expect_equal(x_sampled$x2, x_orig$x2)

  # Sampling more draws than knockoff iters triggers replacement warning
  expect_warning(
    sampled_1_repeats <- sampler_1$sample("x1", row_ids = 1, samples_per_row = 3L),
    regexp = "requested more often than they are present"
  )
  expect_equal(nrow(sampled_1_repeats), 3L)
})

test_that("KnockoffGaussianSampler preserves feature types", {
  skip_if_not_installed("knockoff")
  test_sampler_feature_types(KnockoffGaussianSampler)
})

test_that("KnockoffSampler obeys draw-major order under samples_per_row > 1", {
  skip_if_not_installed("knockoff")
  set.seed(123L)
  n = 20L
  dt = data.table::data.table(
    y = rnorm(n),
    x1 = rnorm(n),
    x2 = rnorm(n),
    tag = seq_len(n) + 0.5
  )
  task = mlr3::as_task_regr(dt, target = "y")
  sampler = KnockoffSampler$new(task, iters = 3L)

  expect_draw_major_row_order(
    sampler,
    task,
    feature = "x1",
    tag_column = "tag",
    samples_per_row = 3L
  )
})

test_that("KnockoffSampler warns and samples with replacement when samples_per_row > iters", {
  skip_if_not_installed("knockoff")
  set.seed(123L)
  task = mlr3::tgen("friedman1")$generate(n = 20)
  sampler = KnockoffSampler$new(task, iters = 2L)

  expect_warning(
    out <- sampler$sample("important1", row_ids = task$row_ids, samples_per_row = 5L),
    regexp = "requested more often than they are present"
  )
  expect_equal(nrow(out), 20L * 5L)
})
