#' @include helpers.R
#' @include hidden_aliases.R
#' @include ChromBackend.R
NULL

#' @title Improved in-memory Chromatographic data backend
#'
#' @name ChromBackendMemory
#'
#'
#' @description
#' `ChromBackendMemory`: This backend stores chromatographic data directly
#' in memory, making it ideal for small datasets or testing. It can be
#' initialized with a `data.frame` of chromatographic data via the `chromData`
#' parameter and a `list` of `data.frame` entries for peaks data using the
#' `peaksData` parameter. These data can be accessed with the `chromData()` and
#' `peaksData()` functions.
#'
#' @param chromData For `backendInitialize()` of a `ChromBackendMemory`
#'        backend, a `data.frame` with the chromatographic data. If not
#'        provided (or if empty), a default `data.frame` with the core
#'        chromatographicvariables will be created.
#'
#' @param object  A `ChromBackendMemory` object.
#'
#' @param peaksData For `backendInitialize()` of a `ChromBackendMemory`
#'        backend, a `list` of `data.frame` with the peaks data. If not
#'        provided (or if empty), a default `list` of empty `data.frame` with
#'        the core peaks variables will be created. The length of the list
#'        should match the number of chromatograms in the `chromData`
#'        parameter.
#'
#' @param ... Additional parameters to be passed.
#'
#' @author Philippine Louail
#'
#' @exportClass ChromBackendMemory
#'
#' @return Refer to the individual function description for information on the
#'         return value.
#'
#' @examples
#'
#' ## Method 1: Initialize backend directly
#' cdata <- data.frame(
#'     msLevel = c(1L, 1L, 1L),
#'     mz = c(112.2, 123.3, 134.4),
#'     dataOrigin = c("mem1", "mem2", "mem3")
#' )
#'
#' pdata <- list(
#'     data.frame(
#'         rtime = c(12.4, 12.8, 13.2, 14.6),
#'         intensity = c(123.3, 153.6, 2354.3, 243.4)
#'     ),
#'     data.frame(
#'         rtime = c(45.1, 46.2),
#'         intensity = c(100, 80.1)
#'     ),
#'     data.frame(
#'         rtime = c(12.4, 12.8, 13.2, 14.6),
#'         intensity = c(123.3, 153.6, 2354.3, 243.4)
#'     )
#' )
#'
#' cbm <- ChromBackendMemory()
#' cbm <- backendInitialize(cbm, chromData = cdata, peaksData = pdata)
#' cbm
#'
#' ## Method 2: Use Chromatograms constructor (recommended)
#' chr <- Chromatograms(ChromBackendMemory(), chromData = cdata, peaksData = pdata)
#' chr
#'
NULL

#' @noRd
setClass("ChromBackendMemory",
    contains = "ChromBackend",
    slots = c(
        chromData = "data.frame",
        peaksData = "list"
    ),
    prototype = prototype(
        chromData = fillCoreChromVariables(data.frame()),
        peaksData = list(.EMPTY_PEAKS_DATA),
        readonly = FALSE,
        version = "0.1"
    )
)

#' @rdname ChromBackendMemory
#'
#' @importFrom methods new
#' @export ChromBackendMemory
ChromBackendMemory <- function() {
    new("ChromBackendMemory")
}

#' @rdname ChromBackendMemory
setMethod(
    "backendInitialize", "ChromBackendMemory",
    function(object,
    chromData = fillCoreChromVariables(data.frame()),
    peaksData = list(.EMPTY_PEAKS_DATA), ...) {
        if (!is(chromData, "data.frame")) {
            stop(
                "'chromData' needs to be a 'data.frame' with the general",
                " chromatogram variables"
            )
        }
        n_cd <- nrow(chromData)
        if (n_cd) {
            if (is.null(chromData$dataOrigin)) {
                chromData$dataOrigin <- NA_character_
            }
            validChromData(chromData)
            chromData <- chromData[, !vapply(chromData, function(x) {
                all(is.na(x))
            }, logical(1)), drop = FALSE]
        } else {
            chromData <- fillCoreChromVariables(data.frame())
        }
        object@chromData <- chromData
        if (length(peaksData) > 1 || length(peaksData) == n_cd) {
            validPeaksData(peaksData)
        } else {
            peaksData <- rep(list(.EMPTY_PEAKS_DATA), n_cd)
        }
        object@peaksData <- peaksData
        object
    }
)

#' @rdname hidden_aliases
setMethod("backendMerge", "ChromBackendMemory", function(object, ...) {
    object <- unname(c(object, ...))
    not_empty <- lengths(object) > 0
    if (any(not_empty)) {
        res <- .df_combine(object[not_empty])
    } else {
        res <- object[[1L]]
    }
    validChromData(.chromData(res))
    validPeaksData(.peaksData(res))
    res
})

#' @rdname hidden_aliases
#' @description This method returns the chromatographic data stored in the
#' backend. If not specified otherwise it will return all defined columns in
#' the chromData slot as well as adding the `coreChromVariables` missing with
#' NA values.
setMethod(
    "chromData", "ChromBackendMemory",
    function(object, columns = chromVariables(object), drop = FALSE) {
        if (!all(columns %in% chromVariables(object))) {
            stop(
                "Some of the requested Chromatogram variables are not ",
                "available"
            )
        }
        res <- fillCoreChromVariables(.chromData(object))
        res <- res[, columns, drop = drop]
        res
    }
)

#' @rdname hidden_aliases
setReplaceMethod("chromData", "ChromBackendMemory", function(object, value) {
    if (!inherits(value, "data.frame")) {
        stop("'value' is expected to be a 'data.frame'")
    }
    if (length(object) && length(object) != nrow(value)) {
        stop("'value' has to be a 'data.frame' with ", length(object), " rows")
    }
    validChromData(value)
    object@chromData <- value
    object
})

#' @rdname hidden_aliases
setMethod("chromVariables", "ChromBackendMemory", function(object) {
    union(names(.chromData(object)), names(coreChromVariables()))
})

#' @rdname hidden_aliases
setMethod(
    "peaksData", "ChromBackendMemory",
    function(object, columns = peaksVariables(object),
    drop = FALSE, ...) {
        pv <- peaksVariables(object)
        if (!all(columns %in% pv)) {
            stop(
                "Some of the requested peaks variables are not",
                " available"
            )
        }
        if (identical(as.vector(pv), as.vector(columns))) {
            return(.peaksData(object))
        }
        if (drop && length(columns) == 1L) {
            return(lapply(.peaksData(object), function(z) z[[columns]]))
        }
        lapply(.peaksData(object), function(x) x[, columns, drop = drop])
    }
)

#' @rdname hidden_aliases
setReplaceMethod("peaksData", "ChromBackendMemory", function(object, value) {
    if (!is.list(value)) {
        stop("'value' is expected to be a list")
    }
    if (length(object) && length(object) != length(value)) {
        stop("'value' has to be a list with ", length(object), " elements")
    }
    validPeaksData(value)
    object@peaksData <- value
    object
})

#' @rdname hidden_aliases
setMethod("peaksVariables", "ChromBackendMemory", function(object) {
    union(names(.peaksData(object)[[1]]), names(corePeaksVariables()))
})

#' @rdname hidden_aliases
#' @importMethodsFrom ProtGenerics backendParallelFactor
setMethod("backendParallelFactor", "ChromBackendMemory",
          function(object, ...) {
    do <- dataOrigin(object)
    if (length(unique(do[!is.na(do)])) > 1L)
        factor(do, levels = unique(do))
    else factor()
})

#' @rdname hidden_aliases
#' @export
setMethod("isReadOnly", "ChromBackendMemory", function(object) FALSE)

#' @rdname hidden_aliases
setMethod("lengths", "ChromBackendMemory", function(x) {
    vapply(.peaksData(x), nrow, integer(1L))
})

#' Optimized intensity accessor for ChromBackendMemory.
#' Bypasses peaksData() dispatch for direct slot access.
#' @rdname hidden_aliases
setMethod("intensity", "ChromBackendMemory", function(object) {
    if (!length(object)) return(list())
    lapply(.peaksData(object), function(z) z$intensity)
})

#' Optimized rtime accessor for ChromBackendMemory.
#' Bypasses peaksData() dispatch for direct slot access.
#' @rdname hidden_aliases
setMethod("rtime", "ChromBackendMemory", function(object) {
    if (!length(object)) return(list())
    lapply(.peaksData(object), function(z) z$rtime)
})

#' @importFrom utils capture.output head
#' @rdname hidden_aliases
setMethod("show", "ChromBackendMemory", function(object) {
    cpd <- .chromData(object)
    cat(class(object), "with", nrow(cpd), "chromatograms\n")
    if (nrow(cpd)) {
        cpd <- fillCoreChromVariables(cpd)
        cpd <- cpd[, c("chromIndex", "msLevel", "mz"), drop = FALSE]
        txt <- capture.output(print(head(cpd)))
        cat(txt, sep = "\n")
        cp_cols <- .chromData(object)[
            ,
            !(colnames(.chromData(object)) %in%
                colnames(cpd))
        ]
        cat("...", length(cp_cols), "more  chromatogram variables/columns\n")
        cat("...", ncol(.peaksData(object)[[1]]), "peaksData variables\n")
    }
})

#' @rdname hidden_aliases
#' @export
setMethod(
    "supportsSetBackend", "ChromBackendMemory",
    function(object, ...) TRUE
)

#' @importFrom MsCoreUtils i2index
#' @importMethodsFrom S4Vectors [ [<-
#' @rdname hidden_aliases
setMethod("[", "ChromBackendMemory", function(x, i, j, ..., drop = FALSE) {
    if (!length(i)) {
        return(ChromBackendMemory())
    }
    i <- i2index(i, length = length(x))
    x@chromData <- .chromData(x)[i, ]
    x@peaksData <- .peaksData(x)[i]
    x
})

#' @rdname hidden_aliases
setMethod("$", "ChromBackendMemory", function(x, name) {
    if (name %in% union(chromVariables(x), names(coreChromVariables()))) {
        res <- chromData(x, columns = name, drop = TRUE)
    } else if (name %in% peaksVariables(x)) {
        res <- peaksData(x, columns = name, drop = TRUE)
    } else {
        stop("The requested variable '", name, "' is not available")
    }
    res
})

#' @rdname hidden_aliases
setReplaceMethod("$", "ChromBackendMemory", function(x, name, value) {
    if (length(x) && length(value) != length(x)) {
        stop(
            "length of 'value' needs to match the number of chromatograms ",
            "in object."
        )
    }
    if (name %in% peaksVariables(x)) {
        if (!is.list(value)) {
            stop("The value for peaksData should be a list")
        }
        for (i in seq_along(value)) peaksData(x)[[i]][[name]] <- value[[i]]
    } else {
        chromData(x)[, name] <- value
    }
    x
})

#' @rdname hidden_aliases
setMethod("chromExtract", "ChromBackendMemory", function(object, peak.table, by) {
    required_cols <- c("rtMin", "rtMax", by)
    .validate_chromExtract_input(
        object = object,
        peak.table = peak.table,
        by = by, required_cols = required_cols
    )

    matched <- .match_chromdata_peaktable(
        object = object,
        peak.table = peak.table,
        by = by
    )
    object <- matched$object
    chrom_keys <- matched$chrom_keys
    peak_keys  <- matched$peak_keys
    obj_sp <- split(object, chrom_keys) ##  UT need to check that
    pk_split <- split(peak.table, peak_keys)

    overl_cols <- .check_overl_columns(peak.table = peak.table,
                                        object = object,
                                        required_cols = required_cols)
    new_data <- mapply(function(obj, pks) { ## could switch to bpmapply ?
        d <- .chromData(obj)
        d <- suppressWarnings(cbind(d, pks[!overl_cols]))
        d[, names(peak.table)[overl_cols]] <- pks[, overl_cols]
        p <- vector("list", nrow(d))
        for (z in seq(nrow(d))) {
            rt <- rtime(obj)[[1]]
            inrt <- rt >= d$rtMin[z] & rt <= d$rtMax[z]
            p[[z]] <- peaksData(obj)[[1]][inrt, , drop = FALSE]
        }
        list(cd = d, pd = p)
    }, obj = obj_sp,
    pks = pk_split, SIMPLIFY = FALSE)

    new_cdata <- .fast_rbind(lapply(new_data, function(z) z$cd))
    rownames(new_cdata) <- NULL
    new_pdata <- unlist(lapply(new_data, function(z) z$pd), recursive = FALSE)
    backendInitialize(new("ChromBackendMemory"),
                      chromData = new_cdata, peaksData = new_pdata)

})
