library(testthat)
library(mockery)
library(DOtools)
library(SingleCellExperiment)
library(Seurat)
library(Matrix)
library(dplyr)

# Define missing helper functions first
.suppressDeprecationWarnings <- function(expr) {
  withCallingHandlers(
    expr,
    warning = function(w) {
      if (grepl("deprecated", w$message, ignore.case = TRUE)) {
        invokeRestart("muffleWarning")
      }
    }
  )
}

grepv <- function(pattern, x, ignore.case = FALSE, invert = FALSE, value = FALSE, ...) {
  result <- grep(pattern, x, ignore.case = ignore.case, invert = invert, value = TRUE, ...)
  if (length(result) == 0) return(character(0))
  return(result)
}

.logger <- function(message) {
  message(paste0("[INFO] ", message))
}

# Enhanced test objects with proper metadata
make_sce <- function() {
  set.seed(1)
  n_genes <- 20
  n_cells <- 30
  mat_counts <- matrix(rpois(n_genes * n_cells, 5), nrow = n_genes)
  rownames(mat_counts) <- paste0("gene", seq_len(n_genes))
  colnames(mat_counts) <- paste0("cell", seq_len(n_cells))
  mat_counts <- Matrix(mat_counts, sparse = TRUE)

  sce <- SingleCellExperiment::SingleCellExperiment(
    assays = list(
      counts = mat_counts,
      logcounts = log1p(mat_counts)
    )
  )

  # Add metadata with expected column names
  sce[["condition"]] <- rep(c("healthy", "disease"), length.out = n_cells)
  sce[["cluster"]] <- rep(c("A", "B", "C"), each = 10, length.out = n_cells)
  sce[["seurat_clusters"]] <- sce[["cluster"]]  # Add expected default column
  sce[["cell_type"]] <- rep(c("Tcell", "Bcell", "Monocyte"), length.out = n_cells)

  return(sce)
}

make_seurat <- function() {
  sce <- make_sce()
  seu <- Seurat::as.Seurat(sce, data = "logcounts")

  # Ensure we have RNA assay and proper metadata
  if ("originalexp" %in% names(seu@assays)) {
    seu@assays$RNA <- seu@assays$originalexp
    seu@assays$originalexp <- NULL
  }
  DefaultAssay(seu) <- "RNA"

  # Add required metadata
  seu$seurat_clusters <- seu$cluster
  seu$condition <- sce$condition

  return(seu)
}

# Improved mock that handles parameter passing correctly
fake_basilisk <- function(env, fun, args) {
  # Basic validation
  if (is.null(args$features) || length(args$features) == 0) {
    stop("Features cannot be empty")
  }

  # Ensure all expected parameters are in args
  expected_params <- c(
    "sce_object", "group_by", "features", "value_plot", "z_score",
    "swap_axes", "cmap", "clustering_method", "clustering_metric",
    "cluster_x_axis", "cluster_y_axis", "figsize", "linewidth",
    "vmin", "vcenter", "vmax", "legend_title", "add_stats",
    "pval_cutoff", "square", "showP", "logcounts", "test"  # Added test here
  )

  # Add missing parameters with NULL values
  for (param in expected_params) {
    if (!param %in% names(args)) {
      args[[param]] <- NULL
    }
  }

  # Return structured result
  list(
    fake_heatmap = TRUE,
    args_passed = args,
    success = TRUE,
    type = if (!is.null(args$showP) && args$showP) "plot" else "dict"
  )
}

# Mock the statistical functions to avoid real DE analysis
mock_find_all_markers <- function(object, features, group.by, test.use, ...) {
  # Return a mock dataframe that matches the expected structure
  groups <- unique(object@meta.data[[group.by]])
  mock_results <- data.frame()

  for (group in groups) {
    for (feature in features) {
      mock_results <- rbind(mock_results, data.frame(
        gene = feature,
        cluster = group,
        avg_log2FC = rnorm(1, 0, 0.5),
        p_val = runif(1, 0, 0.1),
        p_val_adj = runif(1, 0, 0.1),
        pct.1 = runif(1, 0.1, 0.9),
        pct.2 = runif(1, 0.1, 0.9)
      ))
    }
  }

  return(mock_results)
}

mock_find_markers <- function(object, features, group.by, ident.1, ident.2, test.use, ...) {
  # Return a mock dataframe for pairwise comparisons
  mock_results <- data.frame()

  for (feature in features) {
    mock_results <- rbind(mock_results, data.frame(
      avg_log2FC = rnorm(1, 0, 1),
      p_val = runif(1, 0, 0.05),
      p_val_adj = runif(1, 0, 0.05),
      pct.1 = runif(1, 0.3, 0.8),
      pct.2 = runif(1, 0.3, 0.8)
    ))
  }

  rownames(mock_results) <- features
  return(mock_results)
}

# Test basic functionality with improved setup
test_that("DO.Heatmap basic run with add_stats = FALSE works", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  expect_silent({
    res <- DO.Heatmap(
      sce_object = sce,
      features = rownames(sce)[1:3],
      group_by = "cluster",
      add_stats = FALSE,
      showP = FALSE,
      value_plot = "expr"
    )
  })

  expect_true(res$fake_heatmap)
  expect_true(res$success)
})

# Fixed error condition tests
test_that("DO.Heatmap validates inputs thoroughly", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Test missing logcounts assay
  sce_no_logcounts <- sce
  assays(sce_no_logcounts) <- assays(sce_no_logcounts)[names(assays(sce_no_logcounts)) != "logcounts"]

  expect_error(
    DO.Heatmap(
      sce_object = sce_no_logcounts,
      features = rownames(sce)[1:2],
      group_by = "cluster",
      add_stats = FALSE
    ),
    "logcounts not found in assays"
  )
})

# Fixed statistical analysis tests
test_that("DO.Heatmap statistical analysis works correctly", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)
  stub(DO.Heatmap, "Seurat::FindAllMarkers", mock_find_all_markers)
  stub(DO.Heatmap, "Seurat::FindMarkers", mock_find_markers)

  # Test with pre-computed p-values
  df_pvals <- data.frame(
    A = c(0.01, 0.05, 0.1),
    B = c(0.001, 0.01, 0.5),
    C = c(0.1, 0.2, 0.3)
  )
  rownames(df_pvals) <- rownames(sce)[1:3]

  res <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    add_stats = TRUE,
    df_pvals = df_pvals
  )

  expect_true(res$fake_heatmap)
})

# Fixed fold change functionality
test_that("DO.Heatmap fold change calculations work", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)
  stub(DO.Heatmap, "Seurat::FindMarkers", mock_find_markers)

  # Test with explicit ident groups
  res <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    value_plot = "fc",
    group_fc = "condition",
    group_fc_ident_1 = "disease",
    group_fc_ident_2 = "healthy",
    add_stats = TRUE
  )

  expect_true(res$fake_heatmap)
  expect_equal(res$args_passed$value_plot, "fc")
})

# Fixed statistical test parameters test
test_that("DO.Heatmap statistical test parameters work", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)
  stub(DO.Heatmap, "Seurat::FindAllMarkers", mock_find_all_markers)

  # Test only with basic tests that work reliably
  simple_tests <- c("wilcox", "t")

  for (test_type in simple_tests) {
    res <- DO.Heatmap(
      sce_object = sce,
      features = rownames(sce)[1:3],
      group_by = "cluster",
      test = test_type,
      add_stats = TRUE
    )

    expect_true(res$fake_heatmap)
  }

  # Test only_pos parameter
  res_pos <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    only_pos = TRUE,
    add_stats = TRUE
  )

  res_neg <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    only_pos = FALSE,
    add_stats = TRUE
  )

  expect_true(res_pos$fake_heatmap)
  expect_true(res_neg$fake_heatmap)
})

# Test Seurat object handling
test_that("DO.Heatmap works correctly with Seurat objects", {
  seu <- make_seurat()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)
  stub(DO.Heatmap, "Seurat::FindAllMarkers", mock_find_all_markers)

  # Test basic Seurat functionality
  res <- DO.Heatmap(
    sce_object = seu,
    features = rownames(seu)[1:3],
    group_by = "condition",
    assay_normalized = "RNA",
    add_stats = TRUE
  )

  expect_true(res$fake_heatmap)
})

# Test parameter validation
test_that("DO.Heatmap validates parameters correctly", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Test various parameter combinations
  expect_silent(
    DO.Heatmap(
      sce_object = sce,
      features = rownames(sce)[1:3],
      group_by = "cluster",
      z_score = NULL,
      add_stats = FALSE
    )
  )

  expect_silent(
    DO.Heatmap(
      sce_object = sce,
      features = rownames(sce)[1:3],
      group_by = "cluster",
      z_score = "group",
      add_stats = FALSE
    )
  )
})

# Test edge cases and boundary conditions
test_that("DO.Heatmap handles edge cases", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Test with single feature
  res <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1],
    group_by = "cluster",
    add_stats = FALSE
  )

  expect_true(res$fake_heatmap)
  expect_equal(length(res$args_passed$features), 1)
})

# Test visualization parameters
test_that("DO.Heatmap visualization parameters work", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Test different color maps
  color_maps <- c("Reds", "Blues", "viridis")

  for (cmap in color_maps) {
    res <- DO.Heatmap(
      sce_object = sce,
      features = rownames(sce)[1:3],
      group_by = "cluster",
      cmap = cmap,
      add_stats = FALSE
    )

    expect_true(res$fake_heatmap)
    expect_equal(res$args_passed$cmap, cmap)
  }
})

# Test groups ordering
test_that("DO.Heatmap groups ordering works", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Test different group orders
  group_orders <- list(
    c("A", "B", "C"),
    c("C", "B", "A"),
    NULL
  )

  for (groups_order in group_orders) {
    res <- DO.Heatmap(
      sce_object = sce,
      features = rownames(sce)[1:3],
      group_by = "cluster",
      groups_order = groups_order,
      add_stats = FALSE
    )

    expect_true(res$fake_heatmap)
  }
})

# Test axis swapping and rotation
test_that("DO.Heatmap axis parameters work", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Test axis swapping
  res_swap <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    swap_axes = TRUE,
    add_stats = FALSE
  )

  expect_true(res_swap$fake_heatmap)
  expect_true(res_swap$args_passed$swap_axes)
})

# Test figure size and styling
test_that("DO.Heatmap figure styling works", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Test different figure sizes
  fig_sizes <- list(c(4, 5), c(6, 8))

  for (figsize in fig_sizes) {
    res <- DO.Heatmap(
      sce_object = sce,
      features = rownames(sce)[1:3],
      group_by = "cluster",
      figsize = figsize,
      add_stats = FALSE
    )

    expect_true(res$fake_heatmap)
    expect_equal(res$args_passed$figsize, figsize)
  }
})

# Test that all parameters are passed correctly
test_that("DO.Heatmap passes all parameters correctly", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Create a comprehensive test with many parameters
  res <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    assay_normalized = "RNA",
    group_by = "cluster",
    groups_order = c("A", "B", "C"),
    value_plot = "expr",
    z_score = "group",
    clip_value = TRUE,
    max_fc = 3,
    path = tempdir(),
    filename = "test_heatmap.png",
    swap_axes = FALSE,
    cmap = "Blues",
    title = "Test Heatmap",
    title_fontprop = list(size = 14),
    clustering_method = "average",
    clustering_metric = "euclidean",
    cluster_x_axis = TRUE,
    cluster_y_axis = FALSE,
    figsize = c(6, 4),
    linewidth = 0.2,
    ticks_fontdict = list(size = 10),
    xticks_rotation = 45,
    yticks_rotation = 0,
    vmin = 0,
    vmax = 10,
    legend_title = "Expression",
    add_stats = FALSE,  # Disable stats for simplicity
    square = TRUE,
    showP = TRUE,
    logcounts = TRUE
  )

  expect_true(res$fake_heatmap)
  # Verify key parameters are passed correctly
  expect_equal(res$args_passed$group_by, "cluster")
  expect_equal(res$args_passed$value_plot, "expr")
})

# Test return value types
test_that("DO.Heatmap return values are correct", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Test with showP = TRUE
  res_showP_true <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    showP = TRUE,
    add_stats = FALSE
  )

  expect_true(res_showP_true$fake_heatmap)

  # Test with showP = FALSE
  res_showP_false <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    showP = FALSE,
    add_stats = FALSE
  )

  expect_true(res_showP_false$fake_heatmap)
})

# Test performance with larger datasets
test_that("DO.Heatmap handles larger feature sets", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # Test with many features
  many_features <- rownames(sce)[1:min(10, nrow(sce))]

  res <- DO.Heatmap(
    sce_object = sce,
    features = many_features,
    group_by = "cluster",
    add_stats = FALSE
  )

  expect_true(res$fake_heatmap)
})

# Test parameter interactions
test_that("DO.Heatmap parameter interactions work correctly", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)
  stub(DO.Heatmap, "Seurat::FindAllMarkers", mock_find_all_markers)

  # Test combination of parameters
  res <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    cluster_x_axis = TRUE,
    cluster_y_axis = FALSE,
    add_stats = TRUE,
    pval_cutoff = 0.01
  )

  expect_true(res$fake_heatmap)
})

# Test that the function handles warnings appropriately
test_that("DO.Heatmap handles warnings gracefully", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)

  # This should not produce warnings
  expect_silent(
    DO.Heatmap(
      sce_object = sce,
      features = rownames(sce)[1:3],
      group_by = "cluster",
      add_stats = FALSE
    )
  )
})

# Test that all code paths are exercised
test_that("DO.Heatmap exercises all major code paths", {
  sce <- make_sce()
  stub(DO.Heatmap, "basilisk::basiliskRun", fake_basilisk)
  stub(DO.Heatmap, "Seurat::FindAllMarkers", mock_find_all_markers)
  stub(DO.Heatmap, "Seurat::FindMarkers", mock_find_markers)

  # Path 1: Basic functionality without stats
  res1 <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    add_stats = FALSE
  )

  # Path 2: With stats and automatic p-value calculation
  res2 <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    add_stats = TRUE
  )

  # Path 3: With FC calculation
  res3 <- DO.Heatmap(
    sce_object = sce,
    features = rownames(sce)[1:3],
    group_by = "cluster",
    value_plot = "fc",
    group_fc = "condition",
    add_stats = TRUE
  )

  expect_true(res1$fake_heatmap)
  expect_true(res2$fake_heatmap)
  expect_true(res3$fake_heatmap)
})
