#!/usr/bin/env Rscript

suppressPackageStartupMessages({
  library(data.table)
  library(ggplot2)
  library(patchwork)
})

setDTthreads(8)

# =====================================================
# INPUT
# =====================================================

infile <- "annotation_summary_clean.tsv"

dt <- fread(infile)

dt[, state := factor(state, levels=c("Prime","Naive","TSC"))]
# =====================================================
# PERCENT
# =====================================================

dt[, percent := n / sum(n) * 100, by=.(state, type)]

# =====================================================
# COLORS
# =====================================================

genomic_cols <- c(
  promoter   = "#4C78A8",
  exon       = "#F58518",
  intron     = "#54A24B",
  intergenic = "#E45756"
)

te_cols <- c(
  "TE overlap" = "#B279A2",
  "Non-TE"     = "grey75"
)

chrom_cols <- c(
  chrX      = "#E45756",
  autosome  = "#4C78A8"
)

# =====================================================
# GENOMIC
# =====================================================

g_dt <- dt[type=="genomic"]

p1 <- ggplot(
  g_dt,
  aes(state, percent, fill=category)
) +

  geom_col(width=0.75) +

  scale_fill_manual(values=genomic_cols) +

  labs(
    title="Genomic annotation",
    x=NULL,
    y="Percent"
  ) +

  theme_bw(base_size=13)

# =====================================================
# TE
# =====================================================

te_dt <- dt[type=="te"]

p2 <- ggplot(
  te_dt,
  aes(state, percent, fill=category)
) +

  geom_col(width=0.75) +

  scale_fill_manual(values=te_cols) +

  labs(
    title="TE overlap",
    x=NULL,
    y="Percent"
  ) +

  theme_bw(base_size=13)

# =====================================================
# CHROM
# =====================================================

c_dt <- dt[type=="chromosome"]

p3 <- ggplot(
  c_dt,
  aes(state, percent, fill=category)
) +

  geom_col(width=0.75) +

  scale_fill_manual(values=chrom_cols) +

  labs(
    title="Chromosome distribution",
    x=NULL,
    y="Percent"
  ) +

  theme_bw(base_size=13)

# =====================================================
# COMBINE
# =====================================================

final_plot <- p1 + p2 + p3 +
  plot_layout(ncol=3)

ggsave(
  "annotation_barplots.pdf",
  final_plot,
  width=15,
  height=5
)

cat("\n[DONE]\n")