# MESSAGE -----------------------------------------------------------------
#
# author: Yulin Lyu
# email: lvyulin@pku.edu.cn
#
# require: R whatever
#
# ---

# * 1. Load packages ------------------------------------------------------

setwd("exampleData/ATAC")

# grammar
library(tidyverse)
library(magrittr)
library(glue)
library(data.table)

# graphics
library(rtracklayer)
library(ggplot2)
library(ggsci)
library(patchwork)

# * 2. Load data ----------------------------------------------------------

transGene <- readRDS("../../data/hg19canonicalTrans.rds") %>% as.data.table()
transMeta <- readRDS("../../data/hg19trans.rds") %>% as.data.table()
exonMeta <- readRDS("../../data/hg19exon.rds") %>% as.data.table()

bwFile <- list.files("/mnt/f/exampleData/ATAC/bw", full.names = T)[c(2, 1, 3)]
sampleName <- str_remove_all(bwFile, ".*/|.bw")

# * 3. Plot ---------------------------------------------------------------

plotTrack <- function(gene = "SP5") {
  stopifnot(gene %in% transGene$gene_name)
  
  trans <- transGene[gene_name == gene, trans_id]
  chr <- transMeta[tx_name == trans, seqnames] %>% as.vector()
  startPos <- transMeta[tx_name == trans, start]
  endPos <- transMeta[tx_name == trans, end]
  strand <- transMeta[tx_name == trans, strand] %>% as.vector()
  
  extend <- max(2000, (endPos - startPos) * 0.05)
  startEx <- startPos - extend
  endEx <- endPos + extend
  
  exonData <- exonMeta[name == gene]
  exonData <- exonData[map_lgl(tx_name, ~ {trans %in% .x}), .(start, end)]
  
  usedData <- map(bwFile, ~ import.bw(
    .x, format = "bw", selection = BigWigSelection(GRanges(chr, IRanges(startEx, endEx)))))
  names(usedData) <- sampleName
  
  usedData %<>% map(~ as.data.table(.x))
  usedData %<>% imap(~ {data.table(sample = .y, pos = .x$start[1]:.x$end[nrow(.x)], value = rep(.x$score, .x$width))})
  
  plotData <- purrr::reduce(usedData, rbind)
  plotData$sample %<>% factor(levels = sampleName)
  
  usedColor <- ArchR::ArchRPalettes$stallion %>% unname()
  
  g1 <- ggplot(plotData, aes(pos, value)) +
    geom_area(aes(fill = sample), show.legend = F) +
    scale_fill_manual(values = usedColor) +
    labs(y = "Signal intensity") +
    scale_y_continuous(expand = expansion(c(0, 0.05))) +
    scale_x_continuous(expand = c(0, 0)) +
    facet_wrap(~ sample, ncol = 1, strip.position = "right") +
    theme(
      aspect.ratio = 1/10,
      panel.background = element_blank(),
      panel.grid = element_blank(),
      panel.border = element_rect(fill = NA),
      panel.spacing = unit(0, "lines"),
      strip.background = element_blank(),
      plot.margin = margin(),
      axis.ticks.x = element_blank(),
      axis.text.x = element_blank(),
      axis.title.x = element_blank()
    )
  
  segmentData <- data.table(
    start = seq(startPos, endPos, length.out = 10)[2:9]
  )
  
  if(strand == "-") {
    segmentData[, end := start - 1][]
  }else{
    segmentData[, end := start + 1][]
  }
  
  sizeBar <- log10(endPos - startPos) %>% floor() %>% 10^.
  
  g2 <- ggplot(exonData) +
    geom_linerange(x = 0, ymin = startPos, ymax = endPos, size = 0.5) +
    geom_linerange(aes(x = 0, ymin = start, ymax = end), size = 3) +
    geom_segment(
      data = segmentData, aes(y = start, yend = end), x = 0, xend = 0,
      arrow = arrow(length = unit(.4, "lines")), size = 0.3) +
    annotate("segment", x = 1, xend = 1, y = endEx - sizeBar, yend = endEx, size = 1.5) +
    annotate("text", x = 1.2, y = endEx, label = humanFormat(sizeBar), hjust = 1, size = 3) +
    scale_y_continuous(expand = c(0, 0), limits = c(startEx, endEx)) +
    scale_x_continuous(expand = expansion(c(.1, .1))) +
    labs(x = "", y = gene) +
    coord_flip() +
    theme(
      aspect.ratio = 1/10,
      panel.background = element_blank(),
      panel.grid = element_blank(),
      axis.ticks = element_blank(),
      axis.text = element_blank(),
      plot.margin = margin()
    )
  
  g1 + g2 + plot_layout(ncol = 1, heights = c(length(sampleName), 1))
  
  ggsave(str_c("graphics/", gene, "_track.png"), width = 10, height = length(sampleName) + 1)
}

humanFormat <- function(n) {
  stopifnot(is.numeric(n), n > 0)
  
  size <- log10(n)
  
  if(size < 3) {
    as.character(n)
  } else if(size < 6) {
    paste0(round(n / 1e3), "K")
  } else if(size < 9) {
    paste0(round(n / 1e6), "M")
  } else {
    paste0(round(n / 1e9), "G")
  }
}

plotTrack("SP5")
plotTrack("CAT")