library(mariner)
library(marinerData)
library(data.table)
library(InteractionSet)
library(glue, include.only = "glue")
library(strawr)
library(dplyr, include.only = "arrange")
library(GenomicRanges)
library(GenomeInfoDb)

## Shared objects --------------------------------------------------------------

## Read .hic file paths
hicFiles <- c(
    LEUK_HEK_PJA27_inter_30.hic(),
    LEUK_HEK_PJA30_inter_30.hic()
)
names(hicFiles) <- c("FS", "WT")

## Reference BEDPE files (loops called with SIP)
bedpeFiles <- c(
    FS_5kbLoops.txt(),
    WT_5kbLoops.txt()
)
names(bedpeFiles) <- c("FS", "WT")

## Read in bedpeFiles as a list of GInteractions
giList <-
    lapply(bedpeFiles, fread) |>
    lapply(as_ginteractions)

## Create out-of-order, interchromosomal ranges
## (for testing)
giList <-
    lapply(giList, \(x) {
        ## Create a sample of regions to permute
        set.seed(123)
        s <- sample(
            x = seq_len(length(regions(x))),
            size = length(regions(x)),
            replace = FALSE
        )

        ## Scramble regions
        regions(x) <- regions(x)[s]
        x
    })

## Merge pairs
mgi <- mergePairs(
    x = giList,
    radius = 50e03
)

## Bin MergedGInteractions
bgi <- assignToBins(x = mgi, binSize = 50e03)

## Bin with snapping
sgi <- snapToBins(x = mgi, binSize = 50e03)

## Test pullHicMatrices --------------------------------------------------------

## TODO: Change .pullHicMatrices to pullHicMatrices
test_that("pullHicMatrices checks for binning", {
    ## Inform if not binned
    .handleBinning(x = mgi, binSize = 50e03) |>
        expect_message("Pairs are not binned.*")

    ## Still returns binned object
    .handleBinning(x = mgi, binSize = 50e03) |>
        suppressMessages() |>
        expect_s4_class("MergedGInteractions")

    ## Accepts binned pairs
    .handleBinning(
        x = bgi,
        binSize = 50e03
    ) |>
        expect_s4_class("MergedGInteractions")

    ## Throws message when using a non-divisible binSize
    .handleBinning(
        x = bgi,
        binSize = 250e03
    ) |>
        expect_message(".*Snapping to binSize=250000.*")

    ## No message for divisible binSizes
    .handleBinning(
        x = sgi,
        binSize = 10e03
    ) |>
        expect_message(NA)
})

test_that("straw returns data in expected order", {
    ## Define chromosome map & arrange by index
    chrMap <-
        readHicChroms(hicFiles[1]) |>
        arrange(index) |>
        head(n = 25)

    ## Check every combination of chromosomes
    for (i in seq(2, nrow(chrMap))) {
        for (j in seq(2, nrow(chrMap))) {
            chr1loc <- glue('{chrMap[i, "name"]}:0:0')
            chr2loc <- chrMap[j, "name"]
            sparseMat <-
                strawr::straw(
                    norm = "NONE",
                    fname = hicFiles[1],
                    chr1loc = chr1loc,
                    chr2loc = chr2loc,
                    unit = "BP",
                    binsize = 2500000,
                    matrix = "observed"
                )

            ## All values in column "x" should be 0
            if (nrow(sparseMat) != 0) {
                if (i <= j) {
                    expect_true(all(sparseMat$x == 0),
                        label = paste(i, j)
                    )
                } else {
                    expect_false(all(sparseMat$x == 0),
                        label = paste(i, j)
                    )
                }
            }
        }
    }
})

test_that("Check chromosomes in .hic file", {
    ## Assign to x (to avoid modifying in place)
    x <- bgi

    ## Seqnames mismatch throws error
    .checkHicChroms(x, hicFiles[1]) |>
        expect_error("There's.*craziness.*")

    ## Corrected seqnames don't throw error
    GenomeInfoDb::seqlevelsStyle(x) <- "ENSEMBL"
    .checkHicChroms(x, hicFiles[1]) |>
        expect_null()
})

test_that("Blocking works for individual ranges", {
    .blockRanges(
        start = c(0, 40, 45, 50, 110, 120, 0),
        end = c(10, 50, 55, 60, 120, 130, 25),
        blockSize = 50
    ) |>
        expect_identical(
            list(
                start = c(0, 25, 25, 50, 100, 100, 0),
                end = c(50, 75, 75, 100, 150, 150, 50)
            )
        )

    .blockRanges(start = 0, end = 100, blockSize = 50) |>
        expect_error(".*blockSize/2.*")

    .blockRanges(start = 0, end = 50, blockSize = 50) |>
        expect_error(".*blockSize/2.*")

    .blockRanges(start = 0, end = 26, blockSize = 50) |>
        expect_error(".*blockSize/2.*")
})

test_that("GInteractions correctly map to blocks", {
    x <-
        GInteractions(
            GRanges(
                seqnames = "chr1",
                ranges = IRanges(
                    start = c(0, 30, 10, 0, 0),
                    end = c(20, 50, 30, 20, 20)
                )
            ),
            GRanges(
                seqnames = c(rep("chr1", 4), "chr2"),
                ranges = IRanges(
                    start = c(0, 30, 40, 10, 10),
                    end = c(20, 50, 60, 30, 30)
                )
            )
        )

    exp <-
        GInteractions(
            GRanges(
                seqnames = "chr1",
                ranges = IRanges(
                    start = c(0, 25, 0, 0),
                    end = c(50, 75, 50, 50)
                )
            ),
            GRanges(
                seqnames = c(rep("chr1", 3), "chr2"),
                ranges = IRanges(
                    start = c(0, 25, 25, 0),
                    end = c(50, 75, 75, 50)
                )
            )
        )
    exp$block <- seq(1L, 4L)
    exp$xIndex <- list(c(1L, 4L), 2L, 3L, 5L)

    .mapToBlocks(x, 50) |>
        expect_identical(exp)
})

test_that("Straw args are checked correctly", {
    .checkStrawArgs(
        files = hicFiles,
        half = "both",
        norm = "KR",
        binSize = 10e03,
        matrix = "oe"
    ) |>
        expect_null()

    .checkStrawArgs(
        files = hicFiles,
        half = "neither",
        norm = "KR",
        binSize = 10e03,
        matrix = "oe"
    ) |>
        expect_error()

    .checkStrawArgs(
        files = hicFiles,
        half = "both",
        norm = "SCALE",
        binSize = 10e03,
        matrix = "oe"
    ) |>
        expect_error()

    .checkStrawArgs(
        files = hicFiles,
        half = "both",
        norm = "none",
        binSize = 10e03,
        matrix = "oe"
    ) |>
        expect_error()

    .checkStrawArgs(
        files = hicFiles,
        half = "both",
        norm = "KR",
        binSize = 101e03,
        matrix = "oe"
    ) |>
        expect_error()

    .checkStrawArgs(
        files = hicFiles,
        half = "both",
        norm = "KR",
        binSize = 10e03,
        matrix = "obs"
    ) |>
        expect_error()
})

test_that("pullHicPixels pulls correct counts", {
    ## Assign to x (to avoid modifying in place)
    x <- assignToBins(bgi[1:10], 2.5e06)
    GenomeInfoDb::seqlevelsStyle(x) <- "ENSEMBL"

    ## Give names to hicFiles
    namedHicFiles <- hicFiles
    names(namedHicFiles) <- c("Mut", "WT")

    ## Pull pixels
    imat <-
        pullHicPixels(
            x = x,
            binSize = 2.5e06,
            files = namedHicFiles,
            matrix = "observed",
            norm = "NONE",
            half = "both"
        )

    ## Get reference counts with straw directly
    chr1loc <- paste(seqnames1(x), start1(x), start1(x), sep = ":")
    chr2loc <- paste(seqnames2(x), start2(x), start2(x), sep = ":")
    pullCounts <- \(fname, i) {
        straw(
            norm = "NONE",
            fname = fname,
            chr1loc = chr1loc[i],
            chr2loc = chr2loc[i],
            unit = "BP",
            binsize = 2.5e06,
            matrix = "observed"
        )$counts
    }
    Mut <-
        lapply(seq_along(chr1loc), \(i) {
            cnts <- pullCounts(namedHicFiles["Mut"], i)
            if (length(cnts) == 0) cnts <- 0
            cnts
        }) |>
        do.call(rbind, args = _)

    WT <-
        lapply(seq_along(chr1loc), \(i) {
            cnts <- pullCounts(namedHicFiles["WT"], i)
            if (length(cnts) == 0) cnts <- 0
            cnts
        }) |>
        do.call(rbind, args = _)

    ## Test
    expect_identical(
        as.matrix(counts(imat)),
        as.matrix(data.frame(Mut, WT))
    )
})

test_that("pullHicMatrices for square, regular arrays", {
    ## Intrachromosomal
    ## On-diagonal square
    gi <- read.table(text = "
            1 51500000 52000000 1 51500000 52000000
            1 150000000 150500000 1 150000000 150500000") |>
        as_ginteractions()
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "upper")

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "both")
    exp <- array(
        data = c(
            67, 22, 13, 9, 8, 22, 59, 20, 10, 7,
            13, 20, 51, 15, 6, 9, 10, 15, 49, 23,
            8, 7, 6, 23, 66, 63, 25, 15, 4, 2, 25,
            68, 28, 7, 3, 15, 28, 87, 45, 11, 4, 7,
            45, 100, 26, 2, 3, 11, 26, 86
        ),
        dim = c(5, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "upper")
    exp <- array(
        data = c(
            67, NA, NA, NA, NA, 22, 59, NA, NA, NA,
            13, 20, 51, NA, NA, 9, 10, 15, 49, NA,
            8, 7, 6, 23, 66, 63, NA, NA, NA, NA, 25,
            68, NA, NA, NA, 15, 28, 87, NA, NA, 4, 7,
            45, 100, NA, 2, 3, 11, 26, 86
        ),
        dim = c(5, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "lower")
    exp <- array(
        data = c(
            67, 22, 13, 9, 8, NA, 59, 20, 10, 7, NA,
            NA, 51, 15, 6, NA, NA, NA, 49, 23, NA,
            NA, NA, NA, 66, 63, 25, 15, 4, 2, NA, 68,
            28, 7, 3, NA, NA, 87, 45, 11, NA, NA, NA,
            100, 26, NA, NA, NA, NA, 86
        ),
        dim = c(5, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)


    ## Intrachromosomal
    ## Off-diagonal
    ## Define square where start1 < start2
    ## (i.e. upper triangular)
    gi <- read.table(text = "
            1 51500000 51700000 1 51800000 52000000
            1 150000000 150200000 1 150300000 150500000") |>
        as_ginteractions()
    expect_true(all(start1(gi) < start2(gi)))

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "both")
    exp <- array(
        data = c(9, 10, 8, 7, 4, 7, 2, 3),
        dim = c(2, 2, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "upper")
    exp <- array(
        data = c(9, 10, 8, 7, 4, 7, 2, 3),
        dim = c(2, 2, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "lower")
    exp <- array(data = NA_real_, dim = c(2, 2, 2, 1))
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)


    ## Intrachromosomal
    ## Off-diagonal
    ## Define square where start1 > start2
    ## (i.e. lower triangular)
    gi <- read.table(text = "
            1 51800000 52000000 1 51500000 51700000
            1 150300000 150500000 1 150000000 150200000") |>
        as_ginteractions()
    expect_true(all(start1(gi) > start2(gi)))

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "both")
    exp <- array(
        data = c(9, 8, 10, 7, 4, 2, 7, 3),
        dim = c(2, 2, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "upper")
    exp <- array(data = NA_real_, dim = c(2, 2, 2, 1))
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "lower")
    exp <- array(
        data = c(9, 8, 10, 7, 4, 2, 7, 3),
        dim = c(2, 2, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)


    ## Interchromosomal
    ## Define square where seqnames1 < seqnames2
    gi <- read.table(text = "
            1 50000000 55000000 2 55000000 60000000
            1 150000000 155000000 2 155000000 160000000") |>
        as_ginteractions() |>
        suppressWarnings()
    expect_true(all(seqnames1(gi) < seqnames2(gi)))

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "both")
    exp <- array(
        data = c(1, 1, 1, 0, 0, 1, 0, 0),
        dim = c(2, 2, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "upper")
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "lower")
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## Interchromosomal
    ## Define square where seqnames1 > seqnames2
    gi <- read.table(text = "
            2 50000000 55000000 1 55000000 60000000
            2 150000000 155000000 1 155000000 160000000") |>
        as_ginteractions() |>
        suppressWarnings()
    expect_true(all(seqnames1(gi) > seqnames2(gi)))

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "both")
    exp <- array(
        data = c(0, 0, 1, 1, 2, 1, 2, 0),
        dim = c(2, 2, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "upper")
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "lower")
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)
})

test_that("pullHicMatrices for rectangular, regular arrays", {
    ## Intrachromosomal
    ## On-diagonal rectangle
    gi <- read.table(text = "
            1 51000000 51300000 1 51000000 51500000
            1 150000000 150300000 1 150000000 150500000") |>
        as_ginteractions()

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "both")
    exp <- array(
        data = c(
            53, 15, 5, 15, 68, 19, 5, 19, 69, 1,
            8, 12, 4, 5, 2, 63, 25, 15, 25, 68, 28,
            15, 28, 87, 4, 7, 45, 2, 3, 11
        ),
        dim = c(3, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "upper")
    exp <- array(
        data = c(
            53, NA, NA, 15, 68, NA, 5, 19, 69, 1,
            8, 12, 4, 5, 2, 63, NA, NA, 25, 68, NA,
            15, 28, 87, 4, 7, 45, 2, 3, 11
        ),
        dim = c(3, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "lower")
    exp <- array(
        data = c(
            53, 15, 5, NA, 68, 19, NA, NA, 69, NA,
            NA, NA, NA, NA, NA, 63, 25, 15, NA, 68,
            28, NA, NA, 87, NA, NA, NA, NA, NA, NA
        ),
        dim = c(3, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)


    ## Intrachromosomal
    ## Off-diagonal
    ## Define rectangle where start1 < start2
    ## (i.e. upper triangular)
    gi <- read.table(text = "
            1 51000000 51300000 1 51300000 51800000
            1 150000000 150300000 1 150300000 150800000") |>
        as_ginteractions()
    expect_true(all(start1(gi) < start2(gi)))

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "both")
    exp <- array(
        data = c(
            1, 8, 12, 4, 5, 2, 1, 5, 2, 2, 3, 8,
            1, 4, 2, 4, 7, 45, 2, 3, 11, 6, 7, 5,
            3, 3, 3, 6, 2, 3
        ),
        dim = c(3, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "upper")
    exp <- array(
        data = c(
            1, 8, 12, 4, 5, 2, 1, 5, 2, 2, 3, 8,
            1, 4, 2, 4, 7, 45, 2, 3, 11, 6, 7, 5,
            3, 3, 3, 6, 2, 3
        ),
        dim = c(3, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "lower")
    exp <- array(data = NA_real_, dim = c(3, 5, 2, 1))
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)


    ## Intrachromosomal
    ## Off-diagonal
    ## Define rectangle where start1 > start2
    ## (i.e. lower triangular)
    gi <- read.table(text = "
            1 51300000 51800000 1 51000000 51300000
            1 150300000 150800000 1 150000000 150300000") |>
        as_ginteractions()
    expect_true(all(start1(gi) > start2(gi)))

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "both")
    exp <- array(
        data = c(
            1, 4, 1, 2, 1, 8, 5, 5, 3, 4, 12, 2, 2,
            8, 2, 4, 2, 6, 3, 6, 7, 3, 7, 3, 2, 45,
            11, 5, 3, 3
        ),
        dim = c(5, 3, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "upper")
    exp <- array(data = NA_real_, dim = c(5, 3, 2, 1))
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 100e03, half = "lower")
    exp <- array(
        data = c(
            1, 4, 1, 2, 1, 8, 5, 5, 3, 4, 12, 2, 2,
            8, 2, 4, 2, 6, 3, 6, 7, 3, 7, 3, 2, 45,
            11, 5, 3, 3
        ),
        dim = c(5, 3, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)


    ## Interchromosomal
    ## Pull rectangle where seqnames1 < seqnames2
    gi <- read.table(text = "
            1 50000000 57500000 2 50000000 62500000
            1 150000000 157500000 2 150000000 162500000") |>
        as_ginteractions() |>
        suppressWarnings()
    expect_true(all(seqnames1(gi) < seqnames2(gi)))

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "both")
    exp <- array(
        data = c(
            1, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 1,
            0, 3, 2, 0, 2, 0, 2, 1, 0, 1, 0, 0, 0,
            2, 1, 1, 0
        ),
        dim = c(3, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "upper")
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "lower")
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)


    ## Interchromosomal
    ## Pull rectangle where seqnames1 > seqnames2
    gi <- read.table(text = "
            2 50000000 57500000 1 50000000 62500000
            2 150000000 157500000 1 150000000 162500000") |>
        as_ginteractions() |>
        suppressWarnings()
    expect_true(all(seqnames1(gi) > seqnames2(gi)))

    ## half="both"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "both")
    exp <- array(
        data = c(
            1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0,
            2, 0, 2, 0, 0, 0, 2, 1, 2, 1, 0, 2, 0,
            1, 1, 1, 1
        ),
        dim = c(3, 5, 2, 1)
    )
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="upper"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "upper")
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)

    ## half="lower"
    res <- pullHicMatrices(gi, hicFiles[1], 2.5e06, half = "lower")
    expect_identical(interactions(res), gi)
    expect_identical(unname(as.array(counts(res))), exp)
})


test_that("Pull irregular arrays", {
    gi <- read.table(text = "
            1 51000000 51300000 1 51000000 51500000
            2 52000000 52300000 3 52000000 52500000
            1 150000000 150500000 1 150000000 150300000
            2 52000000 52300000 2 52000000 52800000") |>
        as_ginteractions()

    iset <- pullHicMatrices(gi, hicFiles, 100e03, half = "both")

    ## Are counts pulled & accessed correctly?
    iset1 <- pullHicMatrices(gi[1], hicFiles[1], 50e03, half = "both")
    iset2 <- pullHicMatrices(gi, hicFiles[1], 50e03, half = "both")
    exp <- counts(iset1) |> as.matrix()
    expect_identical(counts(iset2)[, , 1, 1] |> as.matrix(), exp)
    expect_identical(as.list(counts(iset2))[[1]][[1]], exp)
})
