#' MultiAssayExperiment wrapper for IntegratedLearner
#' 
#' This function offers a simple interface to the \code{IntegratedLearner()}
#' from a \code{\link{MultiAssayExperiment}} object. Instead of wrangling the
#' data, the function does it for the user and feeds it to
#' \code{IntegratedLearner()}.
#'
#' @param mae \code{\link{MultiAssayExperiment}}
#' 
#' @param experiment \code{Character vector} or \code{Integer vector}. Specifies
#' experiments from \code{mae}.
#' 
#' @param assay.type \code{Character vector}. Specified the name of assay for
#' each experiment. The length must be equal to \code{experiment}.
#' 
#' @param outcome.col \code{Character scalar}. Specifies the name of column
#' from \code{colData} that includes the outcome variable.
#' 
#' @param valid.col \code{Character scalar}. Specifies the name of column
#' from \code{colData} that includes the information on which sample belongs to
#' validation set. The variable must be \code{logical}. (Default: \code{NULL})
#' 
#' @param ... additional arguments passed to \code{IntegratedLearner()}
#' \itemize{
#'   \item \code{na.rm}: \code{Logical scalar}. Should features with missing
#'   values be dropped? (Default: \code{FALSE})
#' }
#' 
#' @return A \code{SuperLearner} object containing the trained model fits.
#' 
#' @export

IntegratedLearnerFromMAE <- function(
        mae, experiment, assay.type, outcome.col, valid.col = NULL, ...){
    ############################### INPUT CHECK ################################
    # MultiAssayExperiment and mia are required to run this function. However,
    # they are not direct dependency of the package; they are only needed when
    # this function is run. tidyr is needed for tidyr::pivot_wider
    .require_package("MultiAssayExperiment")
    .require_package("mia")
    .require_package("tidyr")
    # The object must be MAE
    if( !is(mae, "MultiAssayExperiment") ){
        stop("'mae' must be MultiAssayExperiment.", call. = FALSE)
    }
    # Experiments inside MAR must be SE objects
    if( !all(vapply(MultiAssayExperiment::experiments(mae), function(x)
        is(x, "SummarizedExperiment"), FUN.VALUE = logical(1))) ){
        stop("Experiments in 'mae' must be SummarizedExperiment objects.",
             call. = FALSE)
    }
    # experiment must be either integers specifying the index of experiments or
    # characters specifying the names of experiments
    is_index <- mia:::.is_integer(experiment) &&
        all(experiment <= length(experiments(mae)))
    is_character <- is.character(experiment) && all(experiment %in% names(mae))
    if( !(is_index || is_character) ){
        stop("'experiment' must be a vector specifying either names or ",
             "indices of experiments.", call. = FALSE)
    }
    # assay.type must be a vector of character values. At this point, we do not
    # check if the assays can be found. Moreover, there must be as many
    # assay.types as there are experiments.
    if( !(is.character(assay.type) &&
          length(assay.type) == length(experiment)) ){
        stop("'assay.type' must be a character vector specifying assay for ",
             "each experiment.", call. = FALSE)
    }
    # outcome.col and valid.col must be character values but we dot check
    # here whether they can be found since they can be inside experiments.
    if( !mia:::.is_a_string(outcome.col) ){
        stop("'outcome.col' must be a single character value.")
    }
    if( !(is.null(valid.col) || mia:::.is_a_string(valid.col)) ){
        stop("'valid.col' must be a single character value.")
    }
    ############################# INPUT CHECK END ##############################
    # Subset the MAE
    mae <- mae[ , , experiment]
    # Get data from experiments as a single large table
    data <- .get_data_from_MAE(
        mae, experiment, assay.type, outcome.col, valid.col)
    # Divide data into feature table, and metadata
    data <- .wrangle_data(data, outcome.col, valid.col, ...)
    # Fit the model
    fit <- do.call(IntegratedLearner, data)
    return(fit)
}

################################ HELP FUNCTIONS ################################

# This function fetches data from single experiments and create a single table.
.get_data_from_MAE <- function(
        mae, experiment, assay.type, outcome.col, valid.col){
    # Loop through experiments
    feature_table <- lapply(seq_len(length(experiment)), function(i){
        # Get experiment and assay name for this experiment
        exp <- experiment[[i]]
        assay_name <- assay.type[[i]]
        # Get SE object
        tse <- MultiAssayExperiment::getWithColData(mae, exp, verbose = FALSE)
        # Check if the assay name is correct
        mia:::.check_assay_present(assay_name, tse)
        # Outcome column must be found from colData
        if( !(mia:::.is_a_string(outcome.col) &&
              outcome.col %in% colnames(colData(tse))) ){
            stop("'outcome.col' must specify a column from colData.",
                 call. = FALSE)
        }
        if( !(is.null(valid.col) || (mia:::.is_a_string(valid.col) &&
                                     valid.col %in% colnames(colData(tse)))) ){
            stop("'valid.col' must be NULL or specify a column from colData.",
                 call. = FALSE)
        }
        # Also check that the values in validation set are correct
        if( !is.null(valid.col) && !is.logical(tse[[valid.col]]) ){
            stop("'valid.col' must specify a column that has boolean values.",
                 call. = FALSE)
        }
        # Melt data to single table
        df <- mia::meltSE(
            tse, assay.type = assay_name,
            row.name = "featureID", col.name = "subjectID",
            add.col = c(outcome.col, valid.col))
        # Replace the assay name with general name
        colnames(df)[ colnames(df) == assay_name ] <- "value"
        # Add omic type
        df[["featureType"]] <- as.character(exp)
        return(df)
    })
    # Combine data
    feature_table <- do.call(rbind, feature_table)
    # Convert outcome to numeric
    feature_table[[outcome.col]] <- as.numeric(
        as.factor(feature_table[[outcome.col]]))
    return(feature_table)
}

# This function gets single table as input and outputs an arguments ready for
# IntegratedLearner()
.wrangle_data <- function(
        feature_table, outcome.col, valid.col, na.rm = FALSE, ...){
    #
    if( !mia:::.is_a_bool(na.rm) ){
        stop("'na.rm' must be TRUE or FALSE.", call. = FALSE)
    }
    #
    # Divide data into metadata and feature table
    sample_metadata <- feature_table[
        , c("subjectID", outcome.col, valid.col), drop = FALSE]
    sample_metadata <- sample_metadata[ !duplicated(sample_metadata), ]
    feature_metadata <- feature_table[
        , c("featureID", "featureType"), drop = FALSE]
    feature_metadata <- feature_metadata[ !duplicated(feature_metadata), ]
    feature_table <- feature_table[
        , c("featureID", "subjectID", "value"), drop = FALSE]
    
    # Rename outcome column to Y which is require by IntegratedLearner
    colnames(sample_metadata)[ colnames(sample_metadata) == outcome.col ] <- "Y"
    # Convert feature table into wide format where rows represent features
    # and column samples
    feature_table <- tidyr::pivot_wider(
        feature_table, names_from = "subjectID", values_from = "value")
    
    # Convert datasets into data.frame and add rownames
    sample_metadata <- sample_metadata |> as.data.frame()
    rownames(sample_metadata) <- sample_metadata[["subjectID"]]
    feature_metadata <- feature_metadata |> as.data.frame()
    rownames(feature_metadata) <- feature_metadata[["featureID"]]
    feature_table <- feature_table |> as.data.frame()
    rownames(feature_table) <- feature_table[["featureID"]]
    # Feature table must have only numeric values
    feature_table[["featureID"]] <- NULL
    
    # Check if there are missing values. Drop them if specified. SuperLearner
    # does not support missing values.
    if( any(is.na(feature_table)) && na.rm ){
        keep <- complete.cases(feature_table)
        feature_table <- feature_table[keep, ]
        feature_metadata <- feature_metadata[keep, ]
        # Give error or warning depending on if there are still features left
        # after removing those that are not complete.
        FUN <- if( sum(keep) == 0 ) stop else warning
        FUN(sum(!keep), "/", length(keep),
            " features dropped due to missing data.", call. = FALSE)
    }
    
    # If user specified validation set, divide data into training and validation
    # datasets
    if( !is.null(valid.col) ){
        # Divide feature table
        feature_table_valid <- feature_table[
            , sample_metadata[[valid.col]], drop = FALSE]
        feature_table <- feature_table[
            , !sample_metadata[[valid.col]], drop = FALSE]
        # Divide sample metadata
        sample_metadata_valid <- sample_metadata[sample_metadata[[valid.col]], ]
        sample_metadata <- sample_metadata[!sample_metadata[[valid.col]], ]
        # Create an argument list
        args <- list(
            feature_table = feature_table,
            sample_metadata = sample_metadata,
            feature_metadata = feature_metadata,
            feature_table_valid = feature_table_valid,
            sample_metadata_valid = sample_metadata_valid
        )
    } else{
        # Create an argument list
        args <- list(
            feature_table = feature_table,
            sample_metadata = sample_metadata,
            feature_metadata = feature_metadata
        )
    }
    # Add those arguments that are passed with "..."
    args <- c(args, list(...))
    return(args)
}