#' Annotate a SCE with BLASE Mappings
#'
#' @description
#' Annotates an SCE with the names of bulk samples that best match each
#' pseudotime bin. For each pseudotime bin, we find the highest correlation
#' with a bulk sample that was mapped against it. Because of this approach,
#' a bulk which mapped best to another pseudotime bin may be the best
#' correlation with the current pseudotime bin of interest.
#'
#' @concept annotation
#'
#' @param sce The [SingleCellExperiment::SingleCellExperiment] to annotate.
#' @param blase_results A list of [MappingResult] to use for the
#' annotation.
#' @param annotation_col String. The name of the metadata column in which
#' to store the new annotations.
#' @param include_stats Boolean. Whether or not to include metadata columns
#' containing the correlation of the best matching bin, and whether
#' that mapping was confident.
#'
#' @return A [SingleCellExperiment::SingleCellExperiment] with annotations
#' added to metadata (in a column defined by `annotation_col`), and the
#' correlations in `BLASE_Annotation_Correlation` if
#' `include_stats` is enabled.
#' @export
#'
#' @importFrom SummarizedExperiment colData
#'
#' @examples
#' counts_matrix <- matrix(
#'     c(seq_len(120) / 10, seq_len(120) / 5),
#'     ncol = 48, nrow = 5
#' )
#' sce <- SingleCellExperiment::SingleCellExperiment(assays = list(
#'     normcounts = counts_matrix, logcounts = log(counts_matrix)
#' ))
#' colnames(sce) <- seq_len(48)
#' rownames(sce) <- as.character(seq_len(5))
#' sce$cell_type <- c(rep("celltype_1", 24), rep("celltype_2", 24))
#'
#' sce$pseudotime <- seq_len(48) - 1
#' blase_data <- as.BlaseData(sce, pseudotime_slot = "pseudotime", n_bins = 4)
#' genes(blase_data) <- as.character(seq_len(5))
#'
#' bulk_counts <- matrix(seq_len(15) * 10, ncol = 3, nrow = 5)
#' colnames(bulk_counts) <- c("A", "B", "C")
#' rownames(bulk_counts) <- as.character(seq_len(5))
#'
#' # Map all bulks to bin
#' results <- map_all_best_bins(blase_data, bulk_counts)
#'
#' sce <- assign_pseudotime_bins(
#'     sce,
#'     pseudotime_slot = "pseudotime", n_bins = 4
#' )
#'
#' # Annotate SC from existing bulk
#' sce <- annotate_sce(sce, results)
#' table(sce$BLASE_Annotation)
annotate_sce <- function(
    sce,
    blase_results,
    annotation_col = "BLASE_Annotation",
    include_stats = FALSE) {
    if (!typeof(sce) == "S4" || !inherits(sce, what = "SingleCellExperiment")) {
        stop("SCE must be a Single Cell Experiment object.\n")
    }

    if (is.null(sce$pseudotime_bin)) {
        stop("SCE must have pseudotime_bin slot populated.\n")
    }

    if (
        !typeof(blase_results) == "list" || !all(
            vapply(blase_results, class, FUN.VALUE = "") == "MappingResult"
        )) {
        stop(
            "BLASE Results must be a list of ",
            "MappingResults generated by BLASE.\n"
        )
    }

    SummarizedExperiment::colData(sce)[, annotation_col] <- "Unknown"
    if (include_stats) {
        SummarizedExperiment::colData(sce)[
            , "BLASE_Annotation_Correlation"
        ] <- NA
    }

    bins <- unique(SummarizedExperiment::colData(sce)[["pseudotime_bin"]])
    for (bin in bins) {
        best_match <- PRIVATE_get_best_match(bin, blase_results)
        mask <- SummarizedExperiment::colData(sce)[, "pseudotime_bin"] == bin
        if (!is.null(best_match$mapping_result)) {
            SummarizedExperiment::colData(sce)[mask, annotation_col] <-
                bulk_name(best_match$mapping_result)

            if (include_stats) {
                SummarizedExperiment::colData(sce)[
                    mask, "BLASE_Annotation_Correlation"
                ] <- best_match$correlation
            }
        }
    }

    return(sce)
}

#' @keywords internal
PRIVATE_get_best_match <- function(pdt_bin, blase_results) {
    best_match <- NULL
    best_match_corr <- -1

    for (mapping in blase_results) {
        correlation <- mapping_history(mapping)[[pdt_bin, "correlation"]]
        if (correlation > best_match_corr) {
            best_match <- mapping
            best_match_corr <- correlation
        }
    }

    return(list(mapping_result = best_match, correlation = best_match_corr))
}
