# MESSAGE -----------------------------------------------------------------

#

# author: Yulin Lyu

# email: lvyulin@pku.edu.cn

#

# require: R whatever

#

# ---

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

setwd("exampleData/RNA")

# grammar

library(tidyverse)

library(magrittr)

library(glue)

library(data.table)

# analysis

library(DESeq2)

# graphics

library(ComplexHeatmap)

library(circlize)

library(RColorBrewer)

library(ggsci)

library(scales)

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

vsd <- readRDS("mid/vsd.rds")

vsdMtx <- assay(vsd)

diffData <- fread("mid/ES_vs_Fib.DEG.csv")

colnames(diffData)[1] <- "gene"

diffData[is.na(padj), padj := 1][]

diffData[, p := -log10(padj)][]

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

diffData[, type := "ns"][]

diffData[log2FoldChange > 3 & padj < 0.05, type := "up"][log2FoldChange < -3 & padj < 0.05, type := "down"][]

upGene <- diffData[order(p, decreasing = T)][type == "up"][1:10, gene]

downGene <- diffData[order(p, decreasing = T)][type == "down"][1:10, gene]

# heatmap with many genes

diffGene <- diffData[type != "ns", gene]

geneType <- diffData[type != "ns", type]

markGene <- c(upGene, downGene)

heatData <- vsdMtx[diffGene, ]

heatData %<>% apply(1, scale) %>% t() %>% set_colnames(colnames(vsdMtx))

colorFun <- colorRamp2(seq(1.5, -1.5, len = 9), brewer.pal(9, "RdBu"))

pdf("plot/heatmapMany.pdf", width = 6, height = 8)

Heatmap(

  heatData, col = colorFun, border = T,

  cluster_rows = T, cluster_columns = T,

  show_row_names = F, show_column_names = T,

  show_row_dend = F, show_column_dend = T,

  column_names_rot = 45,

  row_split = geneType,

  column_split = rep(c("Fib", "CiPS", "ES"), each = 2),

  left_annotation = rowAnnotation(

    type = anno_block(

      width = unit(.1, "in"), gp = gpar(fill = pal_nejm()(2)[2:1]))),

  right_annotation = rowAnnotation(

    gene = anno_mark(

      match(markGene, diffGene), markGene,

      labels_gp = gpar(fontface = "italic", fontsize = 10))),

  heatmap_legend_param = list(

    title = "Scaled expression",

    title_position = "lefttop-rot",

    legend_height = unit(1.5, "in")),

  width = unit(3, "in"),

  height = unit(6, "in")

)

dev.off()

# heatmap with few genes

setkey(diffData, gene)

usedGene <- c(upGene, downGene)

geneType <- diffData[usedGene, type]

heatData <- vsdMtx[usedGene, ]

heatData %<>% apply(1, scale) %>% t() %>% set_colnames(colnames(vsdMtx))

colorFun <- colorRamp2(seq(1.5, -1.5, len = 9), brewer.pal(9, "RdBu"))

pdf("plot/heatmapFew.pdf", width = 6, height = 6)

Heatmap(

  heatData, col = colorFun, border = F,

  cluster_rows = T, cluster_columns = T,

  show_row_names = T, show_column_names = T,

  show_row_dend = F, show_column_dend = T,

  column_names_rot = 45,

  row_names_gp = gpar(fontface = "italic"),

  row_split = geneType,

  column_split = rep(c("Fib", "CiPS", "ES"), each = 2),

  cell_fun = function(j, i, x, y, width, height, fill){

    grid.rect(x = x, y = y, width = width, height = height, 

              gp = gpar(col = "white", fill = NA))},

  left_annotation = rowAnnotation(

    type = anno_block(

      width = unit(.1, "in"), gp = gpar(fill = pal_nejm()(2)[2:1]))),

  heatmap_legend_param = list(

    title = "Scaled expression",

    title_position = "lefttop-rot",

    legend_height = unit(1.5, "in")),

  width = unit(3, "in"),

  height = unit(4, "in")

)

dev.off()