#!/usr/bin/env r
#******************************************
## Plot boxplots of expression distribution
## across samples in the ENCODE dataset
## and
## Barplot with the ratio of genes
## expressed in nucleus / cytosol
#******************************************

#-------------------------
## Load required libraries
#-------------------------
library(tidyverse)
library(vroom)
library(ggpubr)

#------------------------
## Plotting Variables
#------------------------
theme_publication <- function(
  size_small = 8, size_medium = 9, size_big = 10,
  frame_size = 1, font = "NimbusSan"
) {
  (theme_bw() +
    theme(
      panel.border = element_rect(linewidth = frame_size),
      text = element_text(size = size_big, family = font),
      axis.text = element_text(size = size_small),
      legend.text = element_text(size = size_medium),
      strip.text = element_text(size = size_big, face = "bold"),
      strip.background = element_blank(),
      plot.tag = element_text(face = "bold")
    )
  )
}

# Colours to use for plotting
cols <- c(
  "Protein-Coding" = "#373790",
  "lncRNAs" = "#01a101",
  "CLS loci" = "#bf62a6",
  "Random decoy" = "#45536e"
)


#-------------------
## Working directory
#-------------------
setwd("~/encode_cell_lines_v47/")

#-----------------------
## Load and process data
#-----------------------
meta <- vroom(
  "data/polyAplus/meta.tsv"
) %>%
  select(-c("File accession", "Paired end")) %>%
  distinct() %>%
  unite(Sample, `Experiment accession`, `Biological replicate(s)`) %>%
  select(Sample, everything()) %>%
  rename(Biosample = `Biosample term name`) %>%
  mutate(Biosample = ifelse(Biosample == "endothelial cell of umbilical vein", "HUVEC", Biosample))

CLS_genes <-
  vroom(
    "data/disjoint.ids/cls.loci.disjoint.ids",
    col_names = "gene_id",
    delim = ".",
    col_select = 1
  )

decoy_genes <-
  vroom(
    "data/disjoint.ids/decoys.loci.disjoint.ids",
    col_names = "gene_id",
    delim = "."
  )

lncRNA_genes <-
  vroom(
    "data/disjoint.ids/lncRNA.loci.v27.disjoint.ids",
    col_names = "gene_id",
    delim = ".",
    col_select = 1
  )

pc_genes <-
  vroom(
    "data/disjoint.ids/protein_coding.loci.v47.disjoint.ids",
    col_names = "gene_id",
    delim = ".",
    col_select = 1
  )

annotation <-
  vroom(
    "data/annotation_transcripts.tsv.gz",
    col_names = c("chr", "start", "end", "strand", "gene_id", "gene_type", "transcript_id", "transcript_type", "gene_name"),
    col_select = c("gene_id", "gene_name", "gene_type")
  ) %>%
  distinct()

expression <-
  vroom(
    "data/expression_polyAplus.tsv.gz",
    col_select = c(sample_name, gene_id, TPM, expected_count, effective_length)
  ) %>%
  rename(Sample = sample_name) %>%
  left_join(
    meta %>%
      select(Sample, Biosample, Fraction)
  )

# Filter lowly expressed genes
cutoff <- 1

filtered_av_expression <-
  expression %>%
  group_by(gene_id, Biosample, Fraction) %>%
  summarise(av_exp = mean(TPM, na.rm = TRUE)) %>%
  ungroup() %>%
  filter(av_exp > cutoff) %>%
  select(-av_exp) %>%
  left_join(expression)

filtered_av_expression <-
  filtered_av_expression %>%
  select(Sample, gene_id, TPM, Biosample, Fraction)

#---------------
## Generate plot
#---------------

# Adjust gene_type column to focus to the types of interest
filtered_av_expression <-
  filtered_av_expression %>%
  left_join(annotation) %>%
  separate_wider_delim(
    gene_id, delim = ".",
    names = c("gene_id", "version"),
    too_few = "align_start"
  ) %>%
  select(-version) %>%
  mutate(
    gene_type = case_when(
      gene_id %in% pc_genes$gene_id ~ "protein_coding",
      gene_id %in% CLS_genes$gene_id ~ "CLS",
      gene_id %in% lncRNA_genes$gene_id ~ "lncRNA",
      gene_id %in% decoy_genes$gene_id ~ "decoy",
      .default = "rest"
    )
  )

# Generate summary table
filtered_av_expression %>%
  filter(gene_type != "rest") %>%
  group_by(Biosample, Fraction, gene_type, gene_id) %>%
  summarise(av_exp = mean(TPM, na.rm = TRUE)) %>%
  summarise(
    N_genes = n(),
    median_TPM = median(av_exp, na.rm = TRUE),
    sum_TPM = sum(av_exp, na.rm = TRUE)
  ) %>%
  mutate(perc_TPM = (sum_TPM / sum(sum_TPM)) * 100) %>%
  ungroup() %>%
  mutate(
    across(where(is.double), ~ round(.x, 2)),
    Fraction = factor(Fraction,
      levels = c("whole_cell", "cytosolic", "nuclear"),
      labels = c("Whole-cell", "Cytosolic", "Nuclear")
    ),
    gene_type = factor(gene_type,
      levels = c(
        "protein_coding", "lncRNA", "CLS", "decoy"
      )
    )
  ) %>%
  arrange(Biosample, Fraction, gene_type) %>%
  pivot_wider(
    names_from = "gene_type",
    values_from = c(
      "N_genes", "median_TPM", "sum_TPM", "perc_TPM"
    ),
    names_sep = "."
  ) %>%
  mutate(
    across(where(~ any(is.na(.x))), ~ replace_na(.x, 0))
  ) %>%
  vroom_write("data/polyAplus/expression_summary_filtered.tsv.gz")


# Boxplots of expression distribution
p <-
  filtered_av_expression %>%
  filter(
    gene_type %in%
      c("protein_coding", "lncRNA", "CLS", "decoy")
  ) %>%
  group_by(Biosample, Fraction, gene_type, gene_id) %>%
  summarise(TPM = mean(TPM, na.rm = TRUE)) %>%
  ungroup() %>%
  mutate(
    gene_type = factor(gene_type,
      levels = c("protein_coding", "lncRNA", "CLS", "decoy"),
      labels = c("Protein-Coding", "lncRNAs", "CLS loci", "Random decoy"),
    ),
    Fraction = factor(Fraction,
      levels = c("nuclear", "cytosolic", "whole_cell"),
      labels = c("Nuclear", "Cytosolic", "Whole-cell")
    ),
    TPM = log2(TPM + 1)
  ) %>%
  ggplot(aes(x = gene_type, y = TPM, colour = gene_type)) +
  geom_violin(trim = FALSE, fill = "grey", lwd = 0.3) +
  facet_grid(Fraction~Biosample, scales = "free_y") +
  theme_publication(
    font = "Helvetica"
  ) +
  theme(
    axis.text.x = element_blank(),
    axis.ticks.x = element_blank(),
    legend.position = "top"
  ) +
  labs(
    x = "Gene Type",
    y = "Expression (log2(TPM + 1))"
  ) +
  scale_colour_manual(
    values = cols,
    guide = guide_legend(character(0), nrow = 1)
  )

pdf(
  "plots/expression_distribution_polyAplus.pdf",
  height = 4,
  width = 12
)

  add_summary(p, "median_iqr", size = 0.01)


dev.off()

# Barplot with the ratio of genes
pdf(
    "plots/ratio_expressed_genes_polyAplus.pdf",
    height = 4,
    width = 4
)
filtered_av_expression %>%
  filter(
    gene_type %in% c("protein_coding", "lncRNA", "CLS") &
      Fraction %in% c("nuclear", "cytosolic")
  ) %>%
  select(gene_id, Biosample, Fraction, gene_type) %>%
  distinct() %>%
  count(Biosample, Fraction, gene_type) %>%
  mutate(
    gene_type = factor(gene_type,
      levels = c("protein_coding", "lncRNA", "CLS"),
      labels = c("Protein-Coding", "lncRNAs", "CLS loci"),
    )
  ) %>%
  pivot_wider(names_from = "Fraction", values_from = "n") %>%
  group_by(Biosample, gene_type) %>%
  summarise(ratio = nuclear / cytosolic) %>%
  ggplot(aes(x = ratio, y = Biosample, fill = gene_type)) +
  geom_bar(stat = "identity", position = "dodge") +
  theme_publication(font = "Helvetica") +
  theme(
    axis.line.y.left = element_line(),
    axis.ticks.x = element_blank(),
    panel.border = element_blank(),
    panel.grid.major.y = element_blank(),
    legend.position = "top"
  ) +
  labs(
    x = "Number of genes in nucleus over cytosol",
    y = ""
  ) +
  scale_fill_manual(
    values = cols,
    guide = guide_legend(character(0), nrow = 1)
  )
dev.off()