context("confound_set and confound_strategy")
library(testthat)
library(bidser)

# --- confound_set tests ---

test_that("confound_set returns correct motion6 variables", {
  m6 <- confound_set("motion6")
  expect_equal(length(m6), 6)
  expect_true(all(c("trans_x", "trans_y", "trans_z",
                     "rot_x", "rot_y", "rot_z") %in% m6))
})

test_that("confound_set returns 24 variables for motion24", {
  m24 <- confound_set("motion24")
  expect_equal(length(m24), 24)
  expect_true("trans_x_derivative1" %in% m24)
  expect_true("rot_z_power2" %in% m24)
  expect_true("trans_x_derivative1_power2" %in% m24)
})

test_that("confound_set 9p has 9 variables", {
  p9 <- confound_set("9p")
  expect_equal(length(p9), 9)
  expect_true(all(c("trans_x", "csf", "white_matter", "global_signal") %in% p9))
})

test_that("confound_set 36p has 36 variables", {
  p36 <- confound_set("36p")
  expect_equal(length(p36), 36)
})

test_that("confound_set is case-insensitive", {
  expect_equal(confound_set("Motion6"), confound_set("motion6"))
  expect_equal(confound_set("MOTION24"), confound_set("motion24"))
})

test_that("confound_set acompcor respects n", {
  acc <- confound_set("acompcor", n = 6)
  expect_equal(acc, "a_comp_cor_*[6]")
  acc_all <- confound_set("acompcor")
  expect_equal(acc_all, "a_comp_cor_*")
})

test_that("confound_set dvars defaults to one standardized DVARS column", {
  expect_equal(confound_set("dvars"), "std_dvars")
  expect_equal(confound_set("std_dvars"), "std_dvars")
  expect_equal(confound_set("raw_dvars"), "dvars")
  expect_equal(confound_set("non_std_dvars"), "non_std_dvars")
  expect_equal(confound_set("vx_wisestd_dvars"), "vx_wisestd_dvars")
})

test_that("confound_set errors on unknown set", {
  expect_error(confound_set("nonexistent"), "Unknown confound set")
})

test_that("list_confound_sets returns a data.frame with expected columns", {
  df <- list_confound_sets()
  expect_s3_class(df, "data.frame")
  expect_true(all(c("set", "description") %in% names(df)))
  expect_true("9p" %in% df$set)
  expect_true("legacy_default" %in% df$set)
})

test_that("confound_set('legacy_default') is the public handle for DEFAULT_CVARS2", {
  legacy <- confound_set("legacy_default")
  expect_equal(length(legacy), 26)
  # Stable public handle must stay byte-identical to the unexported constant
  # that external code reaches for via bidser:::DEFAULT_CVARS2.
  expect_identical(legacy, bidser:::DEFAULT_CVARS2)
  expect_identical(legacy, names(bidser:::CVARS_ALIASES))
  expect_true(all(c("csf", "white_matter", "global_signal",
                    "framewise_displacement",
                    "a_comp_cor_00", "t_comp_cor_05",
                    "trans_x", "rot_z") %in% legacy))
  expect_true(confound_set("LEGACY_DEFAULT") |> identical(legacy))
})


test_that("check_confounds flags zero-variance and rank-deficient columns", {
  df <- tibble::tibble(
    participant_id = "01",
    task = "rest",
    run = "01",
    x = c(1, 2, 3, 4),
    y = c(2, 4, 6, 8),
    z = c(1, 1, 1, 1)
  )

  diag <- check_confounds(df, checks = c("zero_variance", "rank"))
  expect_true("z" %in% diag$column[diag$reason == "zero_variance"])
  expect_equal(sum(diag$reason == "rank_deficient"), 1)
  expect_equal(unique(diag$action), "flag")
})

test_that("clean_confounds drops flagged columns and preserves diagnostics", {
  df <- tibble::tibble(
    participant_id = "01",
    task = "rest",
    run = "01",
    x = c(1, 2, 3, 4),
    y = c(2, 4, 6, 8),
    z = c(1, 1, 1, 1)
  )

  expect_message(
    cleaned <- clean_confounds(df, clean = c("zero_variance", "rank")),
    "Dropped"
  )
  data_cols <- setdiff(names(cleaned), c("participant_id", "task", "run"))
  expect_equal(length(data_cols), 1)
  diag <- attr(cleaned, "confound_diagnostics")
  expect_true(all(c("zero_variance", "rank_deficient") %in% diag$reason))
  expect_equal(unique(diag$action), "drop")
})

# --- confound_strategy tests ---

test_that("confound_strategy creates pcabasic80 object", {
  strat <- confound_strategy("pcabasic80")
  expect_s3_class(strat, "confound_strategy")
  expect_equal(strat$name, "pcabasic80")
  expect_equal(strat$perc_var, 80)
  expect_true(length(strat$pca_vars) > 0)
  expect_true(length(strat$raw_vars) > 0)
  # Should include motion vars and compcor wildcards
  expect_true("trans_x" %in% strat$pca_vars)
  expect_true("csf" %in% strat$pca_vars)
  # Cosine should be in raw_vars
  expect_true(any(grepl("cosine", strat$raw_vars)))
})

test_that("confound_strategy custom works", {
  strat <- confound_strategy(
    pca_vars = c("trans_x", "trans_y"),
    raw_vars = c("csf"),
    npcs = 2
  )
  expect_s3_class(strat, "confound_strategy")
  expect_equal(strat$name, "custom")
  expect_equal(strat$npcs, 2)
  expect_equal(strat$pca_vars, c("trans_x", "trans_y"))
  expect_equal(strat$raw_vars, "csf")
})

test_that("confound_strategy errors on unknown name", {
  expect_error(confound_strategy("nonexistent"), "Unknown confound strategy")
})

test_that("confound_strategy errors when pca_vars missing", {
  expect_error(confound_strategy(raw_vars = "csf"), "pca_vars must be specified")
})

test_that("list_confound_strategies returns a data.frame", {
  df <- list_confound_strategies()
  expect_s3_class(df, "data.frame")
  expect_true("pcabasic80" %in% df$strategy)
})


# --- read_confounds with strategy ---

# Helper to create a richer confounds file for strategy testing
create_rich_confounds_proj <- function() {
  temp_dir <- tempfile("bids_strat_")
  dir.create(temp_dir)
  readr::write_tsv(tibble::tibble(participant_id = "01"),
                   file.path(temp_dir, "participants.tsv"))
  jsonlite::write_json(list(Name = "TestStrategy", BIDSVersion = "1.7.0"),
                       file.path(temp_dir, "dataset_description.json"),
                       auto_unbox = TRUE)
  dir.create(file.path(temp_dir, "sub-01"))
  deriv_root <- file.path(temp_dir, "derivatives", "fmriprep")
  conf_dir <- file.path(deriv_root, "sub-01", "func")
  dir.create(conf_dir, recursive = TRUE)
  jsonlite::write_json(
    list(
      Name = "fMRIPrep",
      BIDSVersion = "1.7.0",
      DatasetType = "derivative",
      GeneratedBy = list(list(Name = "fMRIPrep"))
    ),
    file.path(deriv_root, "dataset_description.json"),
    auto_unbox = TRUE
  )

  n <- 20
  set.seed(42)
  conf_data <- tibble::tibble(
    CSF = rnorm(n),
    WhiteMatter = rnorm(n),
    GlobalSignal = rnorm(n),
    FramewiseDisplacement = abs(rnorm(n, 0, 0.1)),
    X = rnorm(n), Y = rnorm(n), Z = rnorm(n),
    RotX = rnorm(n), RotY = rnorm(n), RotZ = rnorm(n),
    X_derivative1 = rnorm(n), Y_derivative1 = rnorm(n), Z_derivative1 = rnorm(n),
    RotX_derivative1 = rnorm(n), RotY_derivative1 = rnorm(n), RotZ_derivative1 = rnorm(n),
    X_power2 = rnorm(n), Y_power2 = rnorm(n), Z_power2 = rnorm(n),
    RotX_power2 = rnorm(n), RotY_power2 = rnorm(n), RotZ_power2 = rnorm(n),
    X_derivative1_power2 = rnorm(n), Y_derivative1_power2 = rnorm(n),
    Z_derivative1_power2 = rnorm(n),
    RotX_derivative1_power2 = rnorm(n), RotY_derivative1_power2 = rnorm(n),
    RotZ_derivative1_power2 = rnorm(n),
    aCompCor00 = rnorm(n), aCompCor01 = rnorm(n), aCompCor02 = rnorm(n),
    tCompCor00 = rnorm(n), tCompCor01 = rnorm(n), tCompCor02 = rnorm(n),
    cosine_00 = rnorm(n), cosine_01 = rnorm(n), cosine_02 = rnorm(n)
  )

  conf_file <- file.path(conf_dir,
                         "sub-01_task-test_run-01_desc-confounds_timeseries.tsv")
  readr::write_tsv(conf_data, conf_file)
  list(path = temp_dir, proj = bids_project(temp_dir, fmriprep = TRUE), n = n)
}


test_that("read_confounds works with confound_strategy object", {
  setup <- create_rich_confounds_proj()
  on.exit(unlink(setup$path, recursive = TRUE, force = TRUE), add = TRUE)

  strat <- confound_strategy(
    pca_vars = c("csf", "white_matter", "trans_x", "trans_y", "trans_z",
                 "rot_x", "rot_y", "rot_z"),
    raw_vars = c("cosine_*"),
    perc_var = 80
  )
  conf <- read_confounds(setup$proj, cvars = strat, nest = FALSE)
  expect_s3_class(conf, "tbl_df")
  expect_equal(nrow(conf), setup$n)
  # Should have PC columns from PCA
  pc_cols <- grep("^PC", names(conf), value = TRUE)
  expect_true(length(pc_cols) > 0)
  # Should have cosine columns (raw)
  cos_cols <- grep("^cosine", names(conf), value = TRUE)
  expect_true(length(cos_cols) > 0)
  # Should have metadata columns
  expect_true(all(c("participant_id", "run", "session") %in% names(conf)))
})


test_that("read_confounds with strategy: PCA reduces columns", {
  setup <- create_rich_confounds_proj()
  on.exit(unlink(setup$path, recursive = TRUE, force = TRUE), add = TRUE)

  strat <- confound_strategy(
    pca_vars = c("csf", "white_matter", "trans_x", "trans_y", "trans_z",
                 "rot_x", "rot_y", "rot_z"),
    npcs = 3
  )
  conf <- read_confounds(setup$proj, cvars = strat, nest = FALSE)
  pc_cols <- grep("^PC", names(conf), value = TRUE)
  expect_equal(length(pc_cols), 3)
})


test_that("read_confounds with confound_set still works", {
  setup <- create_rich_confounds_proj()
  on.exit(unlink(setup$path, recursive = TRUE, force = TRUE), add = TRUE)

  conf <- read_confounds(setup$proj, cvars = confound_set("motion6"), nest = FALSE)
  expect_s3_class(conf, "tbl_df")
  expect_equal(nrow(conf), setup$n)
  # Should have 6 motion columns plus metadata
  data_cols <- setdiff(names(conf), c("participant_id", "task", "run", "session"))
  expect_equal(length(data_cols), 6)
})

test_that("read_confounds with confound_set('dvars') selects one DVARS variant", {
  temp_dir <- tempfile("bids_dvars_")
  dir.create(temp_dir)
  on.exit(unlink(temp_dir, recursive = TRUE, force = TRUE), add = TRUE)

  readr::write_tsv(tibble::tibble(participant_id = "01"),
                   file.path(temp_dir, "participants.tsv"))
  jsonlite::write_json(list(Name = "DvarsTest", BIDSVersion = "1.8.0"),
                       file.path(temp_dir, "dataset_description.json"),
                       auto_unbox = TRUE)
  deriv_root <- file.path(temp_dir, "derivatives", "fmriprep")
  conf_dir <- file.path(deriv_root, "sub-01", "func")
  dir.create(conf_dir, recursive = TRUE)
  jsonlite::write_json(
    list(
      Name = "fMRIPrep",
      BIDSVersion = "1.8.0",
      DatasetType = "derivative",
      GeneratedBy = list(list(Name = "fMRIPrep"))
    ),
    file.path(deriv_root, "dataset_description.json"),
    auto_unbox = TRUE
  )
  readr::write_tsv(
    tibble::tibble(
      dvars = c(1, 2, 3),
      std_dvars = c(0.1, 0.2, 0.3),
      non_std_dvars = c(10, 20, 30)
    ),
    file.path(conf_dir, "sub-01_task-rest_run-01_desc-confounds_timeseries.tsv")
  )

  proj <- bids_project(temp_dir, fmriprep = TRUE)
  conf <- read_confounds(proj, cvars = confound_set("dvars"), nest = FALSE)
  data_cols <- setdiff(names(conf), c("participant_id", "task", "run", "session"))

  expect_equal(data_cols, "std_dvars")
})


test_that("read_confounds default cvars resolves canonical names to old-style columns", {
  setup <- create_rich_confounds_proj()
  on.exit(unlink(setup$path, recursive = TRUE, force = TRUE), add = TRUE)

  # No explicit cvars: exercises the new default confound_set("legacy_default").
  # The mock dataset uses old fMRIPrep column names, so this also verifies the
  # canonical default resolves back through the alias table.
  conf_default <- read_confounds(setup$proj, nest = FALSE)
  expect_s3_class(conf_default, "tbl_df")
  expect_equal(nrow(conf_default), setup$n)

  data_cols <- setdiff(names(conf_default),
                       c("participant_id", "task", "run", "session"))
  # Motion, global, FD and CompCor present in the mock should all resolve.
  expect_true(all(c("CSF", "WhiteMatter", "GlobalSignal",
                    "FramewiseDisplacement",
                    "X", "Y", "Z", "RotX", "RotY", "RotZ",
                    "aCompCor00", "tCompCor00") %in% data_cols))

  # Passing the explicit public handle must match the implicit default exactly.
  conf_explicit <- read_confounds(setup$proj,
                                  cvars = confound_set("legacy_default"),
                                  nest = FALSE)
  expect_identical(names(conf_default), names(conf_explicit))
})