
#' Plot gene tracks
#' 
#' Plot gene annotation tracks from `ensembldb` data.
#' 
#' @param locus Object of class 'locus' generated by [locus()].
#' @param filter_gene_name Vector of gene names to display.
#' @param filter_gene_biotype Vector of gene biotypes to be filtered. Use
#' [ensembldb::listGenebiotypes()] to display possible biotypes. For example, 
#' `ensembldb::listGenebiotypes(EnsDb.Hsapiens.v75)`
#' @param cex.axis Specifies font size for axis numbering.
#' @param cex.lab Specifies font size for axis titles.
#' @param cex.text Font size for gene text.
#' @param showExons Logical whether to show exons or simply show whole gene as a
#'   rectangle. If `showExons = FALSE` colours are specified by `exon_border`
#'   for rectangle border and `gene_col` for the fill colour.
#' @param maxrows Specifies maximum number of rows to display in gene 
#' annotation panel.
#' @param xticks Logical whether x axis ticks and numbers are plotted.
#' @param xlab Title for x axis. Defaults to chromosome `seqname` specified 
#' in `locus`.
#' @param border Logical whether a bounding box is plotted.
#' @param gene_col Colour for gene lines.
#' @param exon_col Fill colour for exons.
#' @param exon_border Border line colour outlining exons (or genes if
#'   `showExons` is `FALSE`). Set to `NA` for no border.
#' @param text_pos Character value of either 'top' or 'left' specifying
#'   placement of gene name labels.
#' @param italics Logical whether gene text is in italics.
#' @param highlight Vector of genes to highlight.
#' @param highlight_col Single colour or vector of colours for highlighted
#'   genes.
#' @param prioritise Vector of genes to be placed first in the gene tracks.
#' @param blanks Controls handling of genes with blank names: `"fill"` replaces
#'   blank gene symbols with ensembl gene ids. `"hide"` hides genes which are
#'   missing gene symbols.
#' @param showRecomb Logical controls alignment of right margin if
#'   recombination data present.
#' @param align Logical whether to set [par()] to align the plot.
#' @details
#' This function is called by [locus_plot()]. It can be used to plot the gene
#' annotation tracks on their own. It uses base graphics, so
#' [graphics::layout()] can be used to position adjacent plots above or below.
#' 
#' `gene_col`, `exon_col` and `exon_border` set colours for all genes, while
#' `highlight` and `highlight_col` can optionally be used together to highlight
#' specific genes of interest. For full control over every single gene, users
#' can add columns `gene_col`, `exon_col` and `exon_border` to the `TX` object
#' within the 'locus' object. Columns added to `TX` override their equivalent 
#' arguments.
#'
#' @return No return value.
#' @examples
#' if(require(EnsDb.Hsapiens.v75)) {
#' data(SLE_gwas_sub)
#' loc <- locus(SLE_gwas_sub, gene = 'UBE2L3', flank = 1e5,
#'              ens_db = "EnsDb.Hsapiens.v75")
#' genetracks(loc)
#' 
#' ## Limit the number of tracks
#' genetracks(loc, maxrows = 4)
#' 
#' ## Filter by gene biotype
#' genetracks(loc, filter_gene_biotype = 'protein_coding')
#' 
#' ## Customise colours
#' genetracks(loc, gene_col = 'grey', exon_col = 'orange',
#'            exon_border = 'darkgrey')
#' }
#' @importFrom BiocGenerics start end
#' @importFrom graphics axTicks axis lines rect text plot.new strwidth
#' @export

genetracks <- function(locus,
                       filter_gene_name = NULL,
                       filter_gene_biotype = NULL,
                       border = FALSE,
                       cex.axis = 0.9,
                       cex.lab = 1,
                       cex.text = 0.7,
                       gene_col = ifelse(showExons, 'blue4', 'skyblue'),
                       exon_col = 'blue4',
                       exon_border = 'blue4',
                       showExons = TRUE,
                       maxrows = NULL,
                       text_pos = 'top',
                       italics = FALSE,
                       xticks = TRUE,
                       xlab = NULL,
                       highlight = NULL,
                       highlight_col = "red",
                       prioritise = highlight,
                       blanks = c("fill", "hide"),
                       showRecomb = TRUE,
                       align = TRUE) {
  if (!inherits(locus, "locus")) stop("Object of class 'locus' required")
  blanks <- match.arg(blanks)
  TX <- locus$TX
  EX <- locus$EX
  xrange <- locus$xrange
  if (!is.null(filter_gene_name)) {
    TX <- TX[TX$gene_name %in% filter_gene_name, ]
  }
  if (!is.null(filter_gene_biotype)) {
    TX <- TX[TX$gene_biotype %in% filter_gene_biotype, ]
  }
  if (is.null(xlab)) xlab <- paste("Chromosome", locus$seqname, "(Mb)")
  
  recomb <- !is.null(locus$recomb) & showRecomb
  if (align) {
    op <- par(mar = c(ifelse(xticks, 3.5, 1), 3.5, 0.25,
                      ifelse(recomb, 3.5, 1.5)))
    on.exit(par(op))
  }
  
  if (nrow(TX) != 0) {
    TX <- gene_colours(TX, gene_col, exon_col, exon_border, showExons,
                       highlight, highlight_col)
    TX <- mapRow(TX, xlim = xrange, cex.text = cex.text, text_pos = text_pos,
                 blanks = blanks, prioritise = prioritise)
    maxrows <- if (is.null(maxrows)) max(TX$row) else min(c(max(TX$row), maxrows))
    if (max(TX$row) > maxrows) message(max(TX$row), " tracks needed to show all genes")
    TX <- TX[TX$row <= maxrows, ]
  } else maxrows <- 1
  
  plot(NA, xlim = xrange,
       ylim = c(-maxrows - 0.3, -0.3), 
       bty = if (border) 'o' else 'n',
       yaxt = 'n', xaxt = 'n',
       xlab = if (xticks) xlab else "",
       ylab = "",
       cex.lab = cex.lab,
       font.main = 1,
       mgp = c(1.7, 0.4, 0))
  if (xticks) {
    xd <- diff(xrange) * 0.04
    axis(1, at = xrange + c(-xd, xd), labels = FALSE, lwd.ticks = 0)  # extend line
    axis(1, at = axTicks(1), labels = axTicks(1) / 1e6, cex.axis = cex.axis,
         lwd = 0, lwd.ticks = 1,
         tcl = -0.3, mgp = c(1.7, 0.4, 0))
  }
  if (nrow(TX) == 0) {
    message("No genes to plot")
    return(invisible(NULL))
  }
  
  exheight <- switch(text_pos, "top" = 0.15, "left" = 0.3)
  if (showExons) {
    for (i in seq_len(nrow(TX))) {
      lines(TX[i, c('start', 'end')], rep(-TX[i, 'row'], 2),
            col = TX$gene_col[i], lwd = 1.5, lend = 1)
      e <- EX[EX$gene_id == TX$gene_id[i], ]
      exstart <- start(e)
      exend <- end(e)
      rect(exstart, -TX[i, 'row'] - exheight, exend, -TX[i, 'row'] + exheight,
           col = TX$exon_col[i], border = TX$exon_border[i],
           lwd = 0.5, lend = 2, ljoin = 1)
    }
  } else {
    # without exons
    rect(TX[, 'start'], -TX[, 'row'] - exheight,
         TX[, 'end'], -TX[, 'row'] + exheight,
         col = TX$gene_col, lwd = 1, lend = 2, ljoin = 1, border = exon_border)
  }
  
  font <- if (italics) 3 else NULL
  if (text_pos == "top") {
    tfilter <- which(TX$tmin >= (xrange[1] - diff(xrange) * 0.04) & 
                       (TX$tmax <= xrange[2] + diff(xrange) * 0.04))
    for (i in tfilter) {
      text(TX$mean[i], -TX[i, 'row'] + 0.45,
           labels = bquote_gene(TX$gene_name[i], TX$strand[i], italics),
           cex = cex.text, xpd = NA)
    }
  } else if (text_pos == "left") {
    tfilter <- if (border) {
      which(TX$tmin >= xrange[1])
    } else seq_len(nrow(TX))
    for (i in tfilter) {
      text(max(c(TX$start[i], xrange[1] - diff(xrange) * 0.04)), -TX[i, 'row'],
           labels = bquote_gene(TX$gene_name[i], TX$strand[i], italics),
           cex = cex.text, pos = 2, xpd = NA)
    }
  }
  
}

bquote_gene <- function(gene, strand, italics) {
  if (strand == "+") {
    if (!italics) {bquote(.(gene) * symbol("\256"))
    } else bquote(italic(.(gene)) * symbol("\256"))
  } else {
    if (!italics) {bquote(symbol("\254") * .(gene))
    } else bquote(symbol("\254") * italic(.(gene)))
  }
}

# map genes into rows without overlap
mapRow <- function(TX, gap = diff(xlim) * 0.02, cex.text = 0.7, 
                   xlim = range(TX[, c('start', 'end')]),
                   text_pos = 'top', blanks = "fill", prioritise = NULL) {
  blank <- TX$gene_name == ""
  if (any(blank)) {
    if (blanks == "fill") {
      TX$gene_name[blank] <- TX$gene_id[blank]
    } else if (blanks == "hide") {
      TX <- TX[!blank, ]
    }
  }
  w <- match(prioritise, TX$gene_name)
  w <- w[!is.na(w)]
  nprot <- TX$gene_biotype != "protein_coding"
  o <- order(nprot)
  o <- o[!o %in% w]
  TX <- TX[c(w, o), ]
  
  gw <- strwidth(paste0("--", TX$gene_name), units = "inch", 
                 cex = cex.text) * diff(xlim) / par("pin")[1]
  TX$mean <- rowMeans(TX[, c('start', 'end')])
  if (text_pos == 'top') {
    xd <- diff(xlim) * 0.04
    xlim2 <- xlim + c(-xd, xd)
    # left edge
    fix_left <- (TX$mean - gw / 2) < xlim2[1] & (xlim2[1] + gw) < TX$end
    if (any(fix_left)) {
      TX$mean[fix_left] <- xlim2[1] + gw[fix_left] / 2
    }
    # right edge
    fix_right <- (TX$mean + gw / 2) > xlim2[2] & (xlim2[2] - gw) > TX$start
    if (any(fix_right)) {
      TX$mean[fix_right] <- xlim2[2] - gw[fix_right] / 2
    }
    TX$tmin <- TX$mean - gw / 2
    TX$tmax <- TX$mean + gw / 2
  } else if (text_pos == 'left') {
    TX$tmin <- TX$start - gw - gap
    TX$tmax <- TX$end
  } else if (text_pos == 'none') {
    TX$tmax <- TX$tmin <- TX$mean
  }
  TX$min <- pmin(TX$start, TX$end, TX$tmin) - gap / 2
  TX$max <- pmax(TX$start, TX$end, TX$tmax) + gap / 2
  TX$row <- 0
  j <- 1
  while (any(TX$row == 0)) {
    xset <- which(TX$row == 0)
    for (i in xset) {
      # overlap detection
      if (!any(TX$min[i] < TX$max[TX$row == j] &
               TX$max[i] > TX$min[TX$row == j])) {
        TX$row[i] <- j
      }
    }
    j <- j + 1
  }
  TX
}


# highlight selected genes
gene_colours <- function(TX, gene_col, exon_col, exon_border, showExons,
                         highlight, highlight_col) {
  if (is.null(TX$gene_col)) TX$gene_col <- gene_col
  if (is.null(TX$exon_col)) TX$exon_col <- exon_col
  if (is.null(TX$exon_border)) TX$exon_border <- exon_border
  w <- match(highlight, TX$gene_name)
  w <- w[!is.na(w)]
  if (length(w) > 0) {
    TX$gene_col[w] <- highlight_col
    if (showExons) {
      TX$exon_col[w] <- highlight_col
      TX$exon_border[w] <- highlight_col
    }
  }
  TX
}
