#' Plot MDS
#'
#' Plot multi-dimensional scaling plot using algorithm of limma::plotMDS(). It is recommended this be done with the
#' log-methylation-ratio matrix generated by bsseq_to_log_methy_ratio().
#'
#' @param x the log-methylation-ratio matrix.
#' @param top the number of top genes used to calculate pairwise distances.
#' @param plot_dims the numeric vector of the two dimensions to be plotted.
#' @param labels the character vector of labels for data points. By default uses
#'   column names of x, set to NULL to plot points.
#' @param groups the character vector of groups the data points will be coloured
#'   by. Colour palette can be adjusted using scale_colour_*() functions from
#'   ggplot2. If groups is numeric, the points will be coloured by a continuous
#'   colour palette. By default, groups is NULL and the points will not be
#'   coloured.
#' @param legend_name the name for the legend.
#'
#' @return ggplot object of the MDS plot.
#'
#' @examples
#' nmr <- load_example_nanomethresult()
#' bss <- methy_to_bsseq(nmr)
#' lmr <- bsseq_to_log_methy_ratio(bss)
#' plot_mds(lmr)
#'
#' @importFrom limma plotMDS
#' @importFrom ggplot2 draw_key_point
#' @export
plot_mds <- function(x, top = 500, plot_dims = c(1, 2), labels = colnames(x), groups = NULL, legend_name = "group") {
    if (!is.null(labels)) {
        assertthat::assert_that(ncol(x) == length(labels))
    }

    if (!is.null(groups)) {
        assertthat::assert_that(ncol(x) == length(groups))
    }

    mds_res <- limma::plotMDS(
        x,
        top = top,
        plot = FALSE
    )

    var_exp1 <- round(100 * mds_res$var.explained[plot_dims[1]])
    var_exp2 <- round(100 * mds_res$var.explained[plot_dims[2]])

    if ("eigen.vectors" %in% names(mds_res)) {
        plot_data <- data.frame(
            dim1 = mds_res$eigen.vectors[, plot_dims[1]],
            dim2 = mds_res$eigen.vectors[, plot_dims[2]]
        )
        xlabel <- glue::glue("Leading logFC Dim {plot_dims[1]} ({var_exp1}%)")
        ylabel <- glue::glue("Leading logFC Dim {plot_dims[2]} ({var_exp2}%)")
    } else {
        plot_data <- data.frame(
            dim1 = mds_res$cmdscale.out[, plot_dims[1]],
            dim2 = mds_res$cmdscale.out[, plot_dims[2]]
        )
        xlabel <- glue::glue("Leading logFC Dim {plot_dims[1]}")
        ylabel <- glue::glue("Leading logFC Dim {plot_dims[2]}")
    }

    plot_data$labels <- labels
    plot_data$group <- groups

    p <- ggplot2::ggplot(
        plot_data,
        ggplot2::aes(x = .data$dim1, y = .data$dim2)
    )

    if (is.null(groups)) {
        if (is.null(labels)) {
            # no labels or groups
            p <- p + ggplot2::geom_point()
        } else {
            # no groups, but labels
            p <- p + ggplot2::geom_label(aes(label = labels))
        }
    } else {
        if (is.numeric(groups)) {
            # continuous colour palette ignores labels
            if (!is.null(labels)) {
                message("Ignoring labels as groups is numeric. Set `labels=NULL` to suppress this message.")
            }
            p <- p +
                ggplot2::geom_point(aes(colour = .data$group)) +
                ggplot2::scale_color_continuous(name = legend_name)
        } else {
            # discrete colour palette
            if (is.null(labels)) {
                # no labels, but groups
                p <- p + ggplot2::geom_point(aes(colour = .data$group)) +
                    ggplot2::guides(color = ggplot2::guide_legend(title = legend_name))
            } else {
                # labels and groups
                # key_glyph causes the legend to display points
                p <- p + ggplot2::geom_label(
                    aes(label = labels, colour = .data$group),
                    key_glyph = ggplot2::draw_key_point
                )
            }
        }
    }

    p +
        ggplot2::theme_bw() +
        ggplot2::xlab(xlabel) +
        ggplot2::ylab(ylabel) +
        ggplot2::scale_x_continuous(expand = c(.1, .1))
}