#' ploting oncoplot with aplot
#' @param maf MAF object. 
#' @param genes the gene names or the number, default is 20.
#' @return \code{oncoplot} object, which is also a \code{aplot} object
#' @importFrom aplot insert_top insert_right insert_bottom
#' @export
#' @examples
#' \donttest{
#' laml.maf <- system.file("extdata", "tcga_laml.maf.gz", package = "maftools")
#' laml.clin <- system.file('extdata', 'tcga_laml_annot.tsv', package = 'maftools')
#' laml <- maftools::read.maf(maf = laml.maf, clinicalData = laml.clin)
#' oncoplot(maf = laml, genes = 20)
#' }
oncoplot <- function(maf, genes = 20) {
    yulab.utils::check_pkg("maftools")

    comp <- oncoplot_components(maf, genes)
    p_spacer <- ggplot() + ggfun::theme_transparent()

    pp <- comp$main |>
        insert_top(comp$sample, height=.2) |>
        insert_right(p_spacer, width=.01) |>
        insert_right(comp$gene, width=.2) |>
        insert_bottom(p_spacer, height=.1)

    class(pp) <- c("oncoplot", class(pp))

    return(pp)
}

#' Get reusable oncoplot components
#'
#' Build the reusable panels and tidy data used by [oncoplot()].
#' This function is intended for downstream packages that want to assemble
#' an oncoplot with additional tracks while reusing the aplotExtra mutation
#' matrix, sample burden, gene burden, and mutation color mapping.
#'
#' @param maf A MAF object.
#' @param genes The gene names or number of top mutated genes.
#'
#' @return A list with components:
#' \describe{
#'   \item{data}{Tidy oncoplot matrix.}
#'   \item{main}{Main mutation tile plot.}
#'   \item{sample}{Sample-level top annotation plot.}
#'   \item{gene}{Gene-level side annotation plot.}
#'   \item{colors}{Named mutation-class color vector.}
#' }
#' @export
oncoplot_components <- function(maf, genes = 20) {
    d <- oncoplot_tidy_onco_matrix(maf, genes)

    structure(
        list(
            data = d,
            main = oncoplot_main(maf, genes),
            sample = oncoplot_sample(maf, genes),
            gene = oncoplot_gene(maf, genes, ylab = "percentage"),
            colors = get_vcColors(websafe = FALSE)
        ),
        class = "aplotextra_oncoplot_components"
    )
}

#' @importFrom ggplot2 geom_tile
#' @importFrom ggplot2 annotation_custom
#' @importFrom ggplot2 coord_cartesian
oncoplot_main <- function(maf, genes = 20) {
    d <- oncoplot_tidy_onco_matrix(maf, genes)

    p <- ggplot(d, aes(x=.data$Sample, y=.data$Gene, fill=.data$Type)) +
        geom_tile(colour="white", linewidth=.01) + 
        oncoplot_setting(continuous = FALSE) +
        theme(legend.position = "bottom", 
            axis.text.y.left=element_text(face='italic')) 

    leg <- ggfun::get_legend(p)
    p + theme(legend.position = "none", 
            plot.margin = margin(b = 60, r = 5)) +
        annotation_custom(leg, ymin=-1.5, ymax=-.5) + 
        coord_cartesian(clip='off')
    
}

##' @importFrom aplot ylab2
oncoplot_sample <- function(maf, genes = 20, sort = FALSE) {    
    numMat <- get_oncoplot_numericMatrix(maf, genes)
    samp_sum <- obtain.sample.summary.MAF(x = maf) |>
        as.data.frame() |>
        dplyr::select(!.data$total) 

    i <- match(colnames(numMat), samp_sum$Tumor_Sample_Barcode)
    samp_sum <- samp_sum[i, , drop=FALSE]
    d <- tidyr::pivot_longer(samp_sum, -1) |>
        dplyr::rename(Sample=.data$Tumor_Sample_Barcode, Type=.data$name, Freq=.data$value)
    d$Sample <- factor(d$Sample, levels = unique(colnames(numMat)))
    if (sort) {
        td <- dplyr::group_by(.data$d, .data$Type) |> 
            dplyr::summarize(total = sum(.data$Freq)) |> 
            dplyr::arrange(.data$total) |> 
            dplyr::pull(.data$Type)
        d$Type <- factor(d$Type, levels = td)
    }

    p <- ggplot(d, aes(x=.data$Sample,y=.data$Freq,fill=.data$Type)) +
        geom_col(position="stack") +
        oncoplot_setting() +
        ylab2("TMB") 
    
    return(p)
}




#' @importFrom rlang .data
#' @importFrom ggplot2 ggplot
#' @importFrom ggplot2 aes
#' @importFrom ggplot2 coord_flip 
#' @importFrom ggplot2 geom_bar
#' @importFrom ggplot2 xlab
#' @importFrom ggplot2 scale_x_discrete
#' @importFrom ggplot2 scale_y_discrete
#' @importFrom ggplot2 margin
#' @importFrom grid textGrob
#' @importFrom grid gpar
#' @importFrom aplot xlab2
oncoplot_gene <- function(maf, genes = 20, ylab = 'gene') {
    ylab <- match.arg(ylab, c("gene", "percentage"))

    d <- oncoplot_tidy_onco_matrix(maf, genes)
    d <- d[!is.na(d$Type), ]

    p <- ggplot(d, aes(y = .data$Gene, fill = .data$Type)) + 
        geom_bar(position='stack', orientation = 'y') + 
        oncoplot_setting(noxaxis = FALSE, scale='none') #+
        # xlab("No. of samples") #+ 
        # guides(y= guide_axis_label_trans(~str_pad(rev(percent_alt), 5)))

    if (ylab == 'percentage') {
        numMat <- get_oncoplot_numericMatrix(maf, genes)
        totSamps = as.numeric(maf@summary[3, 'summary'])
        percent <- apply(numMat, 1, function(x) sum(x >0))/totSamps
        percent_alt <- paste0(round(percent * 100), '%')
        p <- p + scale_y_discrete(breaks = rownames(numMat),
                                labels = percent_alt,
                                expand = c(0, 0))
    }

    p <- p + xlab2("No. of samples")

    return(p)
}

#' @importFrom ggplot2 theme_minimal
oncoplot_scale <- function(continuous = TRUE, scale = 'y') {
    scale <- match.arg(scale, c('x', 'y', 'none'))
    if (scale == 'none') {
        return(NULL)
    } 
    
    if (scale == 'x') {
        if (continuous) {
            scale_fun <- ggplot2::scale_x_continuous
        } else {
            scale_fun <- ggplot2::scale_x_discrete
        }
    } else {
        if (continuous) {
            scale_fun <- ggplot2::scale_y_continuous
        } else{
            scale_fun <- ggplot2::scale_y_discrete
        }
    }
    
    scale_fun(expand = c(0, 0))

}

#' @importFrom ggplot2 theme_minimal
#' @importFrom ggplot2 theme
#' @importFrom ggplot2 ylab
oncoplot_setting <- function(noxaxis = TRUE, continuous = TRUE, scale='y') {  
    list(
        theme_minimal(),
        if (noxaxis) ggfun::theme_noxaxis(),
        theme(legend.position = "none", panel.grid.major = element_blank()),
        oncoplot_scale(continuous = continuous, scale = scale),
        oncoplot_fill(),
        xlab(NULL),
        ylab(NULL)
    )
}

#' @importFrom ggplot2 scale_fill_manual
oncoplot_fill <- function(breaks=NULL, values=NULL, name = NULL, na.value = "#bdbdbd") {
    vc_col <- get_vcColors(websafe = FALSE)  # maftools color setting
    if (is.null(values)) {
        values <- vc_col
    } 
    if (is.null(breaks)) {
        vc_lev <- names(vc_col)
        breaks <- factor(vc_lev, levels = rev(vc_lev))
    }
    scale_fill_manual(name = NULL, 
            breaks = breaks, 
            values = values, 
            na.value = na.value)
}


get_oncoplot_genes <- function(maf, genes = 20) {
    if (is.numeric(genes)) {
        if (length(genes) == 1) {
            # Top N genes
            genes <- obtain.gene.summary.MAF(x = maf)[1:genes, "Hugo_Symbol"]  
        } else {
            # as index
            genes <- obtain.gene.summary.MAF(x = maf)[genes, "Hugo_Symbol"]
        }
    }
    # if character, return as it is.

    return(genes)
}

get_oncoplot_numericMatrix <- function(maf, genes = 20) {
    genes <- get_oncoplot_genes(maf, genes = genes)
    om <- createOncoMatrix(m = maf, g = genes)
    numMat = om$numericMatrix  # gene/sample ~ frequency
    return(numMat)
}


oncoplot_tidy_onco_matrix <- function(maf, genes = 20) {
    genes <- get_oncoplot_genes(maf, genes)
    om <- createOncoMatrix(m = maf, g = genes)
    mat_origin = om$oncoMatrix # gene/sample ~ variant type

    d = mat_origin[rev(rownames(mat_origin)),] |>
        as.data.frame() |>
        tibble::rownames_to_column('Gene') |>
        tidyr::pivot_longer(-1) |>
        dplyr::rename(Sample = "name", Type = "value")

    d$Type[d$Type == ""] <- NA #"Non_Mut"
    d$Gene <- factor(d$Gene, levels = rev(rownames(mat_origin)))
    d$Sample <- factor(d$Sample, levels = colnames(mat_origin))
    return(d)
}

obtain.sample.summary.MAF <- function(x){
    x@variant.classification.summary
}

obtain.gene.summary.MAF <- function(x){
    x@gene.summary
}

# To add annotation on right y axis: https://github.com/tidyverse/ggplot2/issues/3171
# guide_axis_label_trans <- function(label_trans = identity, ...) {
#   axis_guide <- guide_axis(...)
#   axis_guide$label_trans <- rlang::as_function(label_trans)
#   class(axis_guide) <- c("guide_axis_trans", class(axis_guide))
#   axis_guide
# }

# guide_train.guide_axis_trans <- function(x, ...) {
#   trained <- NextMethod()
#   trained$key$.label <- x$label_trans(trained$key$.label)
#   trained
# }

#' @import maftools
#' @importFrom yulab.utils get_fun_from_pkg
get_vcColors <- yulab.utils::get_fun_from_pkg('maftools', 'get_vcColors')
createOncoMatrix <- yulab.utils::get_fun_from_pkg('maftools', 'createOncoMatrix')

if(getRversion() >= "2.15.1")  {
    utils::globalVariables(c(".", "Hugo_Symbol", 
       "Tumor_Sample_Barcode", "Variant_Classification","Variant_Classification_temp","Variant_Type")
    )
}
