postpsassR: postpsass data visualization

#' @title Drawing circos plot
#' @description Drawing circos plot
#' @param data Generated by load_data_files function
#' @param tracks Choose tracks to plot from c("common", "unique_1", "unique_2", "divergence_major_allele", "divergence_minor_allele")
#' @param tracks.names tracks' name to change
#' @param base.color color for background
#' @param bg.col background color(default: "white")
#' @param top.track to check whether is the top track (defaul: FALSE)
#' @param point.size point size (default: 0.1)
#' @param color.palette point color
#' @export


draw_circos_plot <- function(data,
                             output.file = NULL, width = 2400, height = 2400, res = 120,
                             tracks = c("common", "unique_1", "unique_2", "divergence_major_allele", "divergence_minor_allele"),
                             tracks.names = c(NULL, NULL, NULL, NULL, NULL),
                             base.color = c("lightgreen", "lightpink"),
                             point.size = 0.1,
                             color.unplaced = FALSE,
                             color.palette = c("0"="dodgerblue3", "1"="purple3", "2"="grey20", "3"="firebrick2", "4"="lightblue"),
                             sector.titles.expand = NULL) {
    if (! is.null(output.file)){
        png(output.file, width=width, height=height,  res=res)
    }
    # get sector information
    n_sectors <- length(data$lengths$plot)
    sectors <- names(data$lengths$plot)
    sector_width <- data$lengths$plot
    if (!is.null(data$names)) {
        sector_names <- c(data$names, "Unplaced"="unplaced")
    } else {
        sector_names <- sectors
        names(sector_names) <- sectors
    }
    # Setup sector colors
    bgs <- rep(base.color, n_sectors %/% 2)


    # Setup point colors
    # if (!color.unplaced) {
    #     color.palette["0"] = color.palette["2"]
    #     color.palette["1"] = color.palette["2"]
    # }


    # Create sector lengths matrix
    sector_lengths <- matrix(c(rep(0, n_sectors), data$lengths$plot), ncol=2)


    # Setup gaps between sectors
    gaps <- rep(1, n_sectors)
    gaps[length(gaps)] <- 12  # Bigger gap after the last sector to leave some space for track titles


    # Calculate width of track and sector.title.expand factor based on number of tracks
    n_tracks <- length(tracks)
    track.height <- 0
    if (is.null(sector.titles.expand)) sector.titles.expand <- 1.1 + n_tracks / 10
    if (n_tracks == 1) {
        track.height <- 0.3
    } else if (n_tracks >= 2) {
        track.height <- 0.6 / n_tracks
    }

    a = 7.5
    # Reset circos parameters
    circlize::circos.clear()


    # Setup circos parameters
    circlize::circos.par("track.height" = track.height,
                         "start.degree" = 90 - a,
                         "gap.degree" = gaps,
                         "cell.padding" = c(0.01, 0.1, 0.01, 0.1),
                         "points.overflow.warning" = FALSE,
                         "track.margin" = c(0, 0.015),
                         "unit.circle.segments" = 100)


    # Initialize circos plot
    circlize::circos.initialize(factors = sectors, xlim = sector_lengths, sector.width = sector_width)

    # Draw specified tracks
    for (i in c(1:length(tracks))) {

        top.track <- FALSE
        track.label = NULL
        if (i == 1) top.track <- TRUE
        # print(paste0(' - Warning: unknown sex \"', tracks[i], '\" for window SNP track.'))

        if (tracks[i] == "common") {
            if (! is.null(tracks.names[0])){track.label = tracks.names[0]}
            track_circos(data$common, track_label=track.label, bg.col = bgs,
                         top.track = top.track, color.point = color.palette[i],
                         sector.names = sector_names, sector.titles.expand = sector.titles.expand, sectors = sectors)

        } else if (tracks[i] == "unique_1") {

            if (! is.null(tracks.names[1])){track.label = tracks.names[1]}
            track_circos(data$unique_1, track_label=track.label, bg.col = bgs,
                         top.track = top.track, color.point = color.palette[i],
                         sector.names = sector_names, sector.titles.expand = sector.titles.expand, sectors = sectors)

        } else if (tracks[i] == "unique_2") {

            # TO BE IMPLEMENTED

            if (! is.null(tracks.names[2])){track.label = tracks.names[2]}
            track_circos(data$unique_2, track_label=track.label, bg.col = bgs,
                         top.track = top.track, color.point = color.palette[i],
                         sector.names = sector_names, sector.titles.expand = sector.titles.expand, sectors = sectors)

        } else if (tracks[i] == "divergence_major_allele") {

            if (! is.null(tracks.names[3])){track.label = tracks.names[3]}
            track_circos(data$divergence_major_allele, track_label=track.label, bg.col = bgs,
                         top.track = top.track, color.point = color.palette[i],
                         sector.names = sector_names, sector.titles.expand = sector.titles.expand, sectors = sectors)

        } else if (tracks[i] == "divergence_minor_allele") {

            if (! is.null(tracks.names[4])){track.label = tracks.names[4]}
            track_circos(data$divergence_minor_allele, track_label=track.label, bg.col = bgs,
                         top.track = top.track, color.point = color.palette[i],
                         sector.names = sector_names, sector.titles.expand = sector.titles.expand, sectors = sectors)

        } else {

            print(paste0("Warning: unknown track type \"", tracks[i], "\" ..."))

        }
    }

    # Close output file
    if (!is.null(output.file)) {
        dev.off()
    }
}


#' @title Drawing circos track
#' @description Drawing circos track
#' @param data Generated by load_data_files function
#' @param track_label Label for track
#' @param bg.col background color(default: "white")
#' @param top.track to check whether is the top track (defaul: FALSE)
#' @param point.size point size (default: 0.1)
#' @param color.point point color (default: "black")
#' @param sector.names sector names (default: NULL)

track_circos <- function(data, track_label, bg.col = "white",
                         ylim = c(0, 1.025 * max(data[[3]]) + 0.01), top.track = FALSE, point.size = 0.1,
                         color.point = "black", sector.names = NULL, sector.titles.expand = 1.3, sectors = NULL)
{

    # print(track_label)
    # lb = unlist(strsplit(track_label, ' '))
    # print(lb[1])
    # y_label = bquote(atop(.(lb[1]), .(lb[2])))
    y_label = bquote(.(track_label))
    # if (track_label == "unique_1") {
    #     y_label <- expression(bold(atop("Ameiurus", "melas")))
    # } else if (track_label=="unique_2") {
    #     y_label <- expression(bold(atop("Ictalurus", "punctatus")))
    # }

    # Draw the top track of the plot, showing sex-bias
    circlize::circos.track(factors = data$Contig,
                           x = data$Position,
                           y = data[[3]],
                           ylim = ylim,
                           bg.col = bg.col,
                           panel.fun = function(x, y) {  # panel.fun is the function drawing the track

                               # Get useful sector information
                               sector.index <- circlize::get.cell.meta.data("sector.index")
                               xcenter <- circlize::get.cell.meta.data("xcenter")
                               ymin <- circlize::get.cell.meta.data("ylim")[1]
                               ymax <- circlize::get.cell.meta.data("ylim")[2]
                               xmin <- circlize::get.cell.meta.data("xlim")[1]
                               xmax <- circlize::get.cell.meta.data("xlim")[2]
                               xplot <- circlize::get.cell.meta.data("xplot")

                               # Add top axis and titles to sectors
                               if (top.track) {

                                   # Create x axis on top of sectors
                                   circlize::circos.axis(h = "top",
                                                         major.at = c(0, xmax / 3, 2 * xmax / 3, xmax),  # Label every 1/3 of the axis
                                                         labels.cex = 1.2,
                                                         labels.facing = "outside",
                                                         direction="outside",
                                                         labels = convert_to_mb(c(0, xmax / 3, 2 * xmax / 3, xmax)),  # Conversion to Mb
                                                         minor.ticks = 4,
                                                         labels.pos.adjust = TRUE)

                                   # Add sector names
                                   if (!is.null(sector.names)) {
                                       circlize::circos.text(xcenter,
                                                             sector.titles.expand * ymax,
                                                             sector.names[sector.index],
                                                             cex = 1.5,
                                                             facing = "inside",
                                                             niceFacing = TRUE)
                                   }
                               }

                               # Plot the data
                               circlize::circos.points(x, y,
                                                       cex = point.size,
                                                       bg = bg.col,
                                                       col = color.point)

                               # Add Y axis on the first sector only
                               if (sector.index == sectors[1]) {

                                   # Create y axis
                                   circlize::circos.yaxis(side = "left",
                                                          at = c(ylim[1], (ylim[2] - ylim[1]) / 2 + ylim[1], ylim[2]),  # 3 labels
                                                          sector.index = sectors[1],
                                                          labels.cex = 1.2,
                                                          labels.niceFacing = FALSE,
                                                          labels = round(c(ylim[1], (ylim[2] - ylim[1]) / 2 + ylim[1], ylim[2]), 0))

                                   #Add y axis labels
                                   label_offset <- - 7.5 * (xmax - xmin) / (xplot[1] - xplot[2])  # Axis title will be plotted 5° on the left of the axis
                                   circlize::circos.text(label_offset,
                                                         0.5 * (ymax - ymin) + ymin,
                                                         y_label,
                                                         sector.index = sectors[1],
                                                         facing = "clockwise",
                                                         cex = 1.3,
                                                         font = 2)
                               }
                           }
    )
}