library(testthat)
library(DOtools)
library(Seurat)
library(SingleCellExperiment)

# Load SCE object
sce_data <- readRDS(system.file("extdata", "sce_data.rds", package = "DOtools"))

# Convert to Seurat
sce_seurat <- as.Seurat(sce_data)

# Ensure RNA assay has a data slot
if (!"data" %in% slotNames(sce_seurat[["RNA"]])) {
  sce_seurat[["RNA"]]@data <- as.matrix(GetAssayData(sce_seurat, slot = "counts"))
}

# Ensure metadata has necessary columns
if (!("condition" %in% colnames(sce_seurat@meta.data))) {
  sce_seurat@meta.data$condition <- sample(c("healthy", "disease"), ncol(sce_seurat), replace = TRUE)
}
if (!("cell_type" %in% colnames(sce_seurat@meta.data))) {
  sce_seurat@meta.data$cell_type <- sample(c("A","B"), ncol(sce_seurat), replace = TRUE)
}

# Create ListTest for comparisons
ListTest <- list(c("healthy", "disease"))

# ---- EXISTING TESTS ----

test_that("DO.VlnPlot returns a ggplot object for a single feature", {
  p <- DO.VlnPlot(
    sce_object = sce_seurat,
    SeuV5 = TRUE,
    Feature = "NKG7",
    ListTest = ListTest,
    ctrl.condition = "healthy",
    group.by = "condition"
  )
  expect_s3_class(p, "ggplot")
})

test_that("DO.VlnPlot works with group.by.2 for multi-group testing", {
  # Suppress the ggplot2 warning about vln.df$group
  suppressWarnings({
    p <- DO.VlnPlot(
      sce_object = sce_seurat,
      SeuV5 = TRUE,
      Feature = "NKG7",
      ListTest = ListTest,
      ctrl.condition = "healthy",
      group.by = "condition",
      group.by.2 = "cell_type"
    )
  })
  expect_s3_class(p, "ggplot")
})

test_that("DO.VlnPlot stops when insufficient colors are provided", {
  expect_error(
    DO.VlnPlot(
      sce_object = sce_seurat,
      SeuV5 = TRUE,
      Feature = "NKG7",
      ListTest = ListTest,
      ctrl.condition = "healthy",
      group.by = "condition",
      vector_colors = c("#FF0000")
    ),
    "Only 1 colors provided"
  )
})

# ---- TESTS FOR RETURNVALUES ----

test_that("DO.VlnPlot returnValues diagnosis", {
  # Test what the function actually returns
  result <- DO.VlnPlot(
    sce_object = sce_seurat,
    SeuV5 = TRUE,
    Feature = "NKG7",
    ListTest = ListTest,
    ctrl.condition = "healthy",
    group.by = "condition",
    returnValues = TRUE
  )

  cat("Class of result:", class(result), "\n")
  cat("Type of result:", typeof(result), "\n")
  if (is.list(result)) {
    cat("List names:", names(result), "\n")
    cat("List length:", length(result), "\n")
  }

  expect_true(TRUE)
})

# tests
test_that("DO.VlnPlot basic functionality with various parameters", {
  # Test 1: Basic plot without returnValues
  p1 <- DO.VlnPlot(
    sce_object = sce_seurat,
    SeuV5 = TRUE,
    Feature = "NKG7",
    ListTest = ListTest,
    ctrl.condition = "healthy",
    group.by = "condition"
  )
  expect_s3_class(p1, "ggplot")

  # Test 2: With wilcox_test = FALSE
  p2 <- DO.VlnPlot(
    sce_object = sce_seurat,
    SeuV5 = TRUE,
    Feature = "NKG7",
    ListTest = ListTest,
    ctrl.condition = "healthy",
    group.by = "condition",
    wilcox_test = FALSE
  )
  expect_s3_class(p2, "ggplot")

  # Test 3: With custom jitter args
  p3 <- DO.VlnPlot(
    sce_object = sce_seurat,
    SeuV5 = TRUE,
    Feature = "NKG7",
    ListTest = ListTest,
    ctrl.condition = "healthy",
    group.by = "condition",
    geom_jitter_args = c(0.1, 0.2, 0.3)
  )
  expect_s3_class(p3, "ggplot")
})

test_that("DO.VlnPlot works with SeuV5 = FALSE", {
  p <- DO.VlnPlot(
    sce_object = sce_seurat,
    SeuV5 = FALSE,
    Feature = "NKG7",
    ListTest = ListTest,
    ctrl.condition = "healthy",
    group.by = "condition"
  )
  expect_s3_class(p, "ggplot")
})

test_that("DO.VlnPlot works with metadata feature", {
  # Add a metadata feature for testing
  sce_seurat$test_metadata <- rnorm(ncol(sce_seurat))

  p <- DO.VlnPlot(
    sce_object = sce_seurat,
    SeuV5 = TRUE,
    Feature = "test_metadata",
    ListTest = ListTest,
    ctrl.condition = "healthy",
    group.by = "condition"
  )
  expect_s3_class(p, "ggplot")
})

test_that("DO.VlnPlot error conditions", {
  # Feature not found
  expect_error(
    DO.VlnPlot(
      sce_object = sce_seurat,
      SeuV5 = TRUE,
      Feature = "nonexistent_gene",
      ListTest = ListTest,
      ctrl.condition = "healthy",
      group.by = "condition"
    ),
    "Feature not found in Seurat Object"
  )

  # ctrl.condition is NULL
  expect_error(
    DO.VlnPlot(
      sce_object = sce_seurat,
      SeuV5 = TRUE,
      Feature = "NKG7",
      ListTest = ListTest,
      ctrl.condition = NULL,
      group.by = "condition"
    ),
    "Please specify the ctrl condition"
  )
})

test_that("DO.VlnPlot works with NULL ListTest", {
  p <- DO.VlnPlot(
    sce_object = sce_seurat,
    SeuV5 = TRUE,
    Feature = "NKG7",
    ListTest = NULL,
    ctrl.condition = "healthy",
    group.by = "condition"
  )
  expect_s3_class(p, "ggplot")
})

test_that("DO.VlnPlot warning conditions", {
  # Groups with only 0 expression
  original_data <- sce_seurat[["RNA"]]@data["NKG7", ]
  sce_seurat[["RNA"]]@data["NKG7", ] <- 0

  expect_warning(
    DO.VlnPlot(
      sce_object = sce_seurat,
      SeuV5 = TRUE,
      Feature = "NKG7",
      ListTest = ListTest,
      ctrl.condition = "healthy",
      group.by = "condition"
    ),
    "Some comparisons have no expression"
  )

  # Restore original data
  sce_seurat[["RNA"]]@data["NKG7", ] <- original_data
})

test_that("DO.VlnPlot returnValues behavior", {
  result <- DO.VlnPlot(
    sce_object = sce_seurat,
    SeuV5 = TRUE,
    Feature = "NKG7",
    ListTest = ListTest,
    ctrl.condition = "healthy",
    group.by = "condition",
    returnValues = TRUE
  )

  expect_true(is.list(result) || ggplot2::is.ggplot(result))
})

# Test edge cases that don't depend on returnValues
test_that("DO.VlnPlot edge cases", {
  # Test with group.by.2 and custom parameters
  suppressWarnings({
    p <- DO.VlnPlot(
      sce_object = sce_seurat,
      SeuV5 = TRUE,
      Feature = "NKG7",
      ListTest = ListTest,
      ctrl.condition = "healthy",
      group.by = "condition",
      group.by.2 = "cell_type",
      stat_pos_mod = 1.5,
      hjust.wilcox.2 = 0.7
    )
  })
  expect_s3_class(p, "ggplot")
})

# Test the remove_zeros functionality specifically
test_that("DO.VlnPlot removes zero comparisons", {
  modified_seurat <- sce_seurat
  modified_seurat[["RNA"]]@data["NKG7", ] <- 0

  expect_warning(
    DO.VlnPlot(
      sce_object = modified_seurat,
      SeuV5 = TRUE,
      Feature = "NKG7",
      ListTest = ListTest,
      ctrl.condition = "healthy",
      group.by = "condition"
    ),
    "Removing Test"
  )
})

# comprehensive test
test_that("DO.VlnPlot comprehensive functionality test", {
  # Test multiple parameter
  test_cases <- list(
    list(SeuV5 = TRUE, wilcox_test = TRUE),
    list(SeuV5 = TRUE, wilcox_test = FALSE),
    list(SeuV5 = FALSE, wilcox_test = TRUE),
    list(SeuV5 = FALSE, wilcox_test = FALSE)
  )

  for (params in test_cases) {
    p <- DO.VlnPlot(
      sce_object = sce_seurat,
      SeuV5 = params$SeuV5,
      Feature = "NKG7",
      ListTest = ListTest,
      ctrl.condition = "healthy",
      group.by = "condition",
      wilcox_test = params$wilcox_test
    )
    expect_s3_class(p, "ggplot")
  }
})
