# Tests for integration/transfer related fxns
set.seed(42)
is_not_cran_submission <- isTRUE(as.logical(Sys.getenv("NOT_CRAN")))

pbmc_small <- suppressWarnings(UpdateSeuratObject(pbmc_small))

# Setup test objects
ref <- pbmc_small
query <- CreateSeuratObject(
  counts = as.sparse(
    GetAssayData(
      object = pbmc_small[['RNA']],
      layer = "counts") + rpois(n = ncol(pbmc_small),
      lambda = 1
    )
  )
)
query <- NormalizeData(object = query, verbose = FALSE)
query <- FindVariableFeatures(object = query, verbose = FALSE, nfeatures = 100)
ref <- FindVariableFeatures(object = ref, verbose = FALSE, nfeatures = 100)

anchors <- FindTransferAnchors(reference = ref, query = query, k.filter = 50)

# Tests for TransferData
# ------------------------------------------------------------------------------
context("TransferData")

preds.standard <- TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, verbose = FALSE)
test_that("TransferData default work", {
  # categorical metadata
  expect_equal(dim(preds.standard), c(80, 5))
  expect_equal(colnames(preds.standard)[c(1, 5)], c("predicted.id", "prediction.score.max"))
  expect_equal(rownames(preds.standard), Cells(query))
  expect_equal(preds.standard[1, 1], "1")
  expect_equal(preds.standard[1, 5], 0.4280746, tolerance = 1e-6)
  expect_equal(as.vector(rowSums(as.matrix(preds.standard[, 2:4]))), rep(1, times = ncol(query)))
  expect_true(inherits(preds.standard, "data.frame"))
  # continuous assay data
  pred.assay <- TransferData(anchorset = anchors, refdata = GetAssayData(ref[["RNA"]]), verbose = FALSE)
  expect_equal(dim(pred.assay), c(230, 80))
  expect_equal(GetAssayData(pred.assay, layer = "counts"), new("matrix"))
  expect_equal(GetAssayData(pred.assay, layer = "scale.data"), new("matrix"))
  expect_equal(colnames(pred.assay), Cells(query))
  expect_equal(rownames(pred.assay), rownames(ref[["RNA"]]))
  expect_equal(sum(GetAssayData(pred.assay)[1, ]), 64.46388, tolerance = 1e-6)
  expect_equal(sum(GetAssayData(pred.assay)[, 1]), 281.0306, tolerance = 1e-6)
  expect_true(inherits(pred.assay, "Assay"))
  expect_equal(pred.assay@var.features, logical(0))
  expect_equal(ncol(pred.assay@meta.features), 0)
})

test_that("TransferData can return predictions assay, ", {
  pred.assay <- TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, prediction.assay = TRUE, verbose = FALSE)
  expect_true(inherits(pred.assay, "Assay"))
  expect_equal(dim(pred.assay), c(4, 80))
  expect_equal(GetAssayData(pred.assay, layer = "counts"), new("matrix"))
  expect_equal(GetAssayData(pred.assay, layer = "scale.data"), new("matrix"))
  expect_equal(colnames(pred.assay), Cells(query))
  expect_equal(pred.assay@var.features, logical(0))
  expect_equal(ncol(pred.assay@meta.features), 0)
  expect_equal(sum(GetAssayData(pred.assay)[1, ]), 26.59365, tolerance = 1e-6)
  expect_equal(sum(GetAssayData(pred.assay)[, 1]), 1.428075, tolerance = 1e-6)
  expect_equal(as.vector(colSums(GetAssayData(pred.assay)[1:3, ])), rep(1, ncol(query)))
})


if (is_not_cran_submission) {
  test_that("TransferData handles weight.reduction properly, ", {
    # test for custom dimreduc
    custom.dr <- anchors@object.list[[1]][["pcaproject"]]
    custom.dr <- subset(x = custom.dr, cells = anchors@query.cells)
    custom.dr <- RenameCells(object = custom.dr, new.names = sapply(X = Cells(custom.dr), FUN = function(x){
      x <- gsub(pattern = "_query", replacement = "", x = x)
    }))
    expect_error(TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, weight.reduction = custom.dr, dims = 1:100))
    preds <-TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, verbose = FALSE)
    cdr.preds <- TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, weight.reduction = custom.dr, verbose = FALSE, dims = 1:30)
    expect_equal(preds, cdr.preds)
    # weight.reduction = "pca
    pca.preds <- suppressWarnings(TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, query = query, weight.reduction = "pca", verbose = FALSE))
    expect_true(inherits(pca.preds, "Seurat"))
    expect_equal(sum(GetAssayData(pca.preds[['prediction.score.id']])[1, ]), 27.83330252, tolerance = 1e-6)
    # weight.reduction = "cca"
    anchors.cca <- FindTransferAnchors(reference = ref, query = query, k.filter = 50, reduction = "cca")
    cca.preds <- TransferData(anchorset = anchors.cca, refdata = ref$RNA_snn_res.1, weight.reduction = "cca", verbose = FALSE)
    expect_true(inherits(cca.preds, "data.frame"))
    expect_equal(sum(cca.preds[, 2]), 43.61738383, tolerance = 1e-6)
  })

  test_that("TransferData with l2.norm works", {
    preds <- TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, l2.norm = TRUE, verbose = FALSE)
    expect_equal(dim(preds), c(80, 5))
    expect_equal(colnames(preds)[c(1, 5)], c("predicted.id", "prediction.score.max"))
    expect_equal(rownames(preds), Cells(query))
    expect_equal(preds[1, 1], "0")
    expect_equal(preds[1, 5], 0.3973124793, tolerance = 1e-6)
    expect_equal(as.vector(rowSums(as.matrix(preds[, 2:4]))), rep(1, times = ncol(query)))
    expect_true(inherits(preds, "data.frame"))
  })

  test_that("TransferData with other k.weight works", {
    preds <- TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, k.weight = 10, verbose = FALSE)
    expect_equal(dim(preds), c(80, 5))
    expect_equal(colnames(preds)[c(1, 5)], c("predicted.id", "prediction.score.max"))
    expect_equal(rownames(preds), Cells(query))
    expect_equal(preds[1, 1], "2")
    expect_equal(preds[1, 5], 0.6145459065, tolerance = 1e-6)
    expect_equal(as.vector(rowSums(as.matrix(preds[, 2:4]))), rep(1, times = ncol(query)))
    expect_true(inherits(preds, "data.frame"))
  })

  test_that("TransferData with reference specified works", {
    pred2 <- TransferData(anchorset = anchors, refdata = "RNA_snn_res.1", reference = ref, verbose = FALSE)
    expect_equal(preds.standard, pred2)
  })

  test_that("TransferData throws expected errors ", {
    expect_error(TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, weight.reduction = "BAD")) # better message
    expect_error(TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, weight.reduction = "cca")) # better message
    expect_error(TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, dims = 1:100))
    expect_error(ransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, k.weight = 1000))
    expect_error(suppressWarnings(TransferData(anchorset = anchors, refdata = "RNA_snn_res.1")))
    expect_error(suppressWarnings(TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1[1:10])))
    expect_error(TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, query = subset(x = query, cells = Cells(query)[1:10])))
  })

  test_that("TransferData with multiple items to transfer works ", {
    preds <- TransferData(anchorset = anchors, refdata = list(
      ids = ref$RNA_snn_res.1, groups = ref$groups, dat = GetAssayData(ref[["RNA"]])),
      verbose = FALSE)
    expect_equal(length(preds), 3)
    expect_equal(preds[[1]], preds.standard)
  })

  test_that("TransferData can return a modified query object ", {
    query <- suppressWarnings(TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, query = query, verbose = FALSE))
    expect_true("prediction.score.id" %in% Assays(query))
    expect_true("predicted.id" %in% colnames(query[[]]))
    expect_true("predicted.id.score" %in% colnames(query[[]]))
    query <- suppressWarnings(TransferData(anchorset = anchors, refdata = ref$RNA_snn_res.1, query = query, store.weights = TRUE, verbose = FALSE))
    expect_equal(dim(Tool(query, slot = "TransferData")$weights.matrix), c(128, 80))
  })
}