#' Plot PCA
#'
#' Plot multi-dimensional scaling plot using algorithm of BiocSingular::runPCA(). 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 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.
#' @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_pca(lmr)
#'
#' @importFrom BiocSingular runPCA
#'
#' @export
plot_pca <- function(x, 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))
    }

    pca_res <- BiocSingular::runPCA(t(x), rank = max(plot_dims))

    plot_data <- data.frame(
        dim1 = pca_res$x[, plot_dims[1]],
        dim2 = pca_res$x[, plot_dims[2]]
    )

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

    xlabel <- glue::glue("PCA Dim {plot_dims[1]}")
    ylabel <- glue::glue("PCA Dim {plot_dims[2]}")

    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))
}