

# Testing collapse is time consuming, skip if ran
skip_collapse <- Sys.getenv("FILEARRAY_SKIP_COLLAPSE", unset = "") == "TRUE"
testthat::skip_if(skip_collapse)

collapse_real <- function(y,
                          keep,
                          transform = c("asis", "10log10", "square", "sqrt", "normalize")) {
    re <- switch(
        transform,
        "asis" = {
            apply(y, keep, function(x) {
                mean(x)
            })
        },
        "10log10" = {
            apply(y, keep, function(x) {
                mean(10 * log10(x))
            })
        },
        "square" = {
            apply(y, keep, function(x) {
                mean(x^2)
            })
        },
        "sqrt" = {
            apply(y, keep, function(x) {
                mean(sqrt(x))
            })
        },
        {
            stop("wrong transform")
        }
    )
    # if(storage.mode(re) != "double") {
    #     storage.mode(re) <- "double"
    # }
    re
}

collapse_cplx <- function(y,
                          keep,
                          transform = c("asis", "10log10", "square", "sqrt", "normalize")) {
    re <- switch(
        transform,
        "asis" = {
            apply(y, keep, mean)
        },
        "10log10" = {
            apply(y, keep, function(x) {
                mean(20 * log10(Mod(x)))
            })
        },
        "square" = {
            apply(y, keep, function(x) {
                mean(Mod(x)^2)
            })
        },
        "sqrt" = {
            apply(y, keep, function(x) {
                mean(Mod(x))
            })
        },
        "normalize" = {
            apply(y, keep, function(x) {
                mean(x / Mod(x))
            })
        },
        {
            stop("wrong transform")
        }
    )
    # if(storage.mode(re) != "double") {
    #     storage.mode(re) <- "double"
    # }
    re
}

expect_equivalent_cplx <- function(x, y, eps = 1e-6) {
    expect_equal(is.na(x), is.na(y))
    if (is.complex(x)) {
        expect_lte(max(abs(Re(x - y)), na.rm = TRUE), eps)
        expect_lte(max(abs(Im(x - y)), na.rm = TRUE), eps)
    } else {
        expect_lte(max(abs(x - y), na.rm = TRUE), eps)
    }
}

test_that("R/C++ - Collapse", {
    testthat::skip_on_cran()
    bsz <- get_buffer_size()
    on.exit({
        set_buffer_size(bsz)
        max_buffer_size(2097152)
    })
    set_buffer_size(16L)
    max_buffer_size(64L)
    
    # dim <- c(287, 100, 301, 7)
    dim <- c(33:36)
    set.seed(5)
    file <- tempfile()
    unlink(file, recursive = TRUE)
    x <- filearray_create(
        file,
        dim,
        type = "integer",
        partition_size = 2,
        initialize = FALSE
    )
    y <- array(1:(prod(dim)), dim)
    y[[20, 3, 3, 3]] <- NA
    storage.mode(y) <- "integer"
    x[] <- y
    
    # make sure x[] == y
    expect_equal(x[], y)
    
    
    
    # collapse
    keep <- c(1, 2, 3, 4)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    
    keep <- c(1, 4, 3, 2)
    expect_equal(
        x$collapse(
            keep = keep,
            transform = "asis",
            method = "mean"
        ),
        collapse_real(y, keep, transform = "asis")
    )
    
    keep <- c(4, 2, 3, 1)
    expect_equal(
        x$collapse(
            keep = keep,
            transform = "asis",
            method = "mean"
        ),
        collapse_real(y, keep, transform = "asis")
    )
    
    keep <- c(3, 1)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    
    keep <- c(4, 1)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    
    keep <- c(4, 2)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    keep <- c(4, 3)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    keep <- c(4)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        k <- x$collapse(keep = keep,
                        transform = transform,
                        method = "mean")
        s <- collapse_real(y, keep, transform = transform)
        diff <- max(abs(1 - s / k), na.rm = TRUE)
        # cat(transform, diff, "\n")
        expect_lt(diff, 1e-6)
    }
    keep <- c(4, 1, 3)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    keep <- c(3)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        k <- x$collapse(keep = keep,
                        transform = transform,
                        method = "mean")
        s <- collapse_real(y, keep, transform = transform)
        diff <- max(abs(1 - s / k), na.rm = TRUE)
        # cat(transform, diff, "\n")
        expect_lt(diff, 1e-6)
    }
    keep <- c(1)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        k <- x$collapse(keep = keep,
                        transform = transform,
                        method = "mean")
        s <- collapse_real(y, keep, transform = transform)
        diff <- max(abs(1 - s / k), na.rm = TRUE)
        # cat(transform, diff, "\n")
        expect_lt(diff, 1e-6)
    }
    
})

test_that("R/C++ - Float", {
    testthat::skip_on_cran()
    bsz <- get_buffer_size()
    on.exit({
        set_buffer_size(bsz)
        max_buffer_size(2097152)
    })
    set_buffer_size(16L)
    max_buffer_size(64L)
    
    # dim <- c(287, 100, 301, 7)
    dim <- c(33:36)
    set.seed(5)
    file <- tempfile()
    unlink(file, recursive = TRUE)
    x <- filearray_create(
        file,
        dim,
        type = "float",
        partition_size = 2,
        initialize = FALSE
    )
    y <- array(rnorm(length(x))^2, dim)
    y[[20, 3, 3, 3]] <- NA
    x[] <- y
    
    # make sure x[] == y
    eps <- 10^(ceiling(log10(max(abs(
        y
    ), na.rm = TRUE))) - 7)
    expect_equal(x[], y, tolerance = eps)
    y <- x[]
    
    
    
    # collapse
    keep <- c(1, 2, 3, 4)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    
    keep <- c(1, 4, 3, 2)
    expect_equal(
        x$collapse(
            keep = keep,
            transform = "asis",
            method = "mean"
        ),
        collapse_real(y, keep, transform = "asis")
    )
    
    keep <- c(4, 2, 3, 1)
    expect_equal(
        x$collapse(
            keep = keep,
            transform = "asis",
            method = "mean"
        ),
        collapse_real(y, keep, transform = "asis")
    )
    
    keep <- c(3, 1)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    
    keep <- c(4, 1)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    
    keep <- c(4, 2)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    keep <- c(4, 3)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    keep <- c(4)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        k <- x$collapse(keep = keep,
                        transform = transform,
                        method = "mean")
        s <- collapse_real(y, keep, transform = transform)
        diff <- max(abs(1 - s / k), na.rm = TRUE)
        # cat(transform, diff, "\n")
        expect_lt(diff, 1e-6)
    }
    keep <- c(4, 1, 3)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_real(y, keep, transform = transform)
        )
    }
    keep <- c(3)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        k <- x$collapse(keep = keep,
                        transform = transform,
                        method = "mean")
        s <- collapse_real(y, keep, transform = transform)
        diff <- max(abs(1 - s / k), na.rm = TRUE)
        # cat(transform, diff, "\n")
        expect_lt(diff, 1e-6)
    }
    keep <- c(1)
    for (transform in c("asis", "10log10", "square", "sqrt")) {
        k <- x$collapse(keep = keep,
                        transform = transform,
                        method = "mean")
        s <- collapse_real(y, keep, transform = transform)
        diff <- max(abs(1 - s / k), na.rm = TRUE)
        # cat(transform, diff, "\n")
        expect_lt(diff, 1e-6)
    }
    
})

test_that("R/C++ - Collapse (complex)", {
    testthat::skip_on_cran()
    bsz <- get_buffer_size()
    on.exit({
        set_buffer_size(bsz)
        max_buffer_size(2097152)
    })
    set_buffer_size(16L)
    max_buffer_size(64L)
    
    # dim <- c(287, 100, 301, 7)
    dim <- c(33:36)
    set.seed(5)
    file <- tempfile()
    unlink(file, recursive = TRUE)
    x <- filearray_create(
        file,
        dim,
        type = "complex",
        partition_size = 2,
        initialize = FALSE
    )
    y <- array(rnorm(length(x)) + rnorm(length(x)) * 1i, dim)
    y[[20, 3, 3, 3]] <- NA
    x[] <- y
    
    # make sure x[] == y
    expect_equivalent_cplx(x[], y)
    
    y <- x[]
    
    # collapse
    keep <- c(1, 2, 3, 4)
    for (transform in c("asis", "10log10", "square", "sqrt", "normalize")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_cplx(y, keep, transform = transform)
        )
    }
    
    keep <- c(1, 4, 3, 2)
    expect_equal(
        x$collapse(
            keep = keep,
            transform = "asis",
            method = "mean"
        ),
        collapse_cplx(y, keep, transform = "asis")
    )
    
    keep <- c(4, 2, 3, 1)
    expect_equal(
        x$collapse(
            keep = keep,
            transform = "asis",
            method = "mean"
        ),
        collapse_cplx(y, keep, transform = "asis")
    )
    
    keep <- c(3, 1)
    for (transform in c("asis", "10log10", "square", "sqrt", "normalize")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_cplx(y, keep, transform = transform)
        )
    }
    
    keep <- c(4, 1)
    for (transform in c("asis", "10log10", "square", "sqrt", "normalize")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_cplx(y, keep, transform = transform)
        )
    }
    
    keep <- c(4, 2)
    for (transform in c("asis", "10log10", "square", "sqrt", "normalize")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_cplx(y, keep, transform = transform)
        )
    }
    keep <- c(4, 3)
    for (transform in c("asis", "10log10", "square", "sqrt", "normalize")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_cplx(y, keep, transform = transform)
        )
    }
    keep <- c(4)
    for (transform in c("asis", "10log10", "square", "sqrt", "normalize")) {
        expect_equal(max(abs(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ) -
                collapse_cplx(y, keep, transform = transform)
        ), na.rm = TRUE), 0)
    }
    keep <- c(4, 1, 3)
    for (transform in c("asis", "10log10", "square", "sqrt", "normalize")) {
        expect_equal(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ),
            collapse_cplx(y, keep, transform = transform)
        )
    }
    keep <- c(3)
    for (transform in c("asis", "10log10", "square", "sqrt", "normalize")) {
        expect_equal(max(abs(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ) -
                collapse_cplx(y, keep, transform = transform)
        ), na.rm = TRUE), 0)
    }
    keep <- c(1)
    for (transform in c("asis", "10log10", "square", "sqrt", "normalize")) {
        expect_equal(max(abs(
            x$collapse(
                keep = keep,
                transform = transform,
                method = "mean"
            ) -
                collapse_cplx(y, keep, transform = transform)
        ), na.rm = TRUE), 0)
    }
    
})


Sys.setenv("FILEARRAY_SKIP_COLLAPSE" = "TRUE")
