#!/usr/bin/env Rscript

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

setDTthreads(8)

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

infile <- "/BLUES/eric/ONT_WGBS/Figure_5/annotation_files/annotation_summary.tsv"

out_pdf <- "/BLUES/eric/ONT_WGBS/Figure_5/DhMR_genomic_distribution.pdf"

# =====================================================
# LOAD
# =====================================================

dt <- fread(infile)

# =====================================================
# ORDER
# =====================================================

dt[
  ,
  comparison := factor(
    comparison,
    levels = c(
      "Naive_vs_Primed",
      "TSC_vs_Naive"
    ),
    labels = c(
      "Naive vs Primed",
      "TSC vs Naive"
    )
  )
]

dt[
  ,
  category := factor(
    category,
    levels = c(
      "promoter",
      "exon",
      "intron",
      "intergenic"
    )
  )
]

# =====================================================
# PERCENT
# =====================================================

dt[
  ,
  percent := n / sum(n) * 100,
  by = comparison
]

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

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

# =====================================================
# TOTAL LABELS
# =====================================================

total_dt <- dt[
  ,
  .(total = sum(n)),
  by = comparison
]

# =====================================================
# PLOT
# =====================================================

p <- ggplot(
  dt,
  aes(
    comparison,
    percent,
    fill = category
  )
) +

  geom_col(
    width = 0.75
  ) +

  geom_text(
    aes(
      label = sprintf("%.1f%%", percent)
    ),
    position = position_stack(vjust = 0.5),
    size = 5,
    color = "white",
    fontface = "bold"
  ) +

  geom_text(
    data = total_dt,
    aes(
      x = comparison,
      y = 103,
      label = paste0("n = ", total)
    ),
    inherit.aes = FALSE,
    size = 5,
    fontface = "bold"
  ) +

  scale_fill_manual(
    values = genomic_cols
  ) +

  expand_limits(y = 108) +

  labs(
    title = "Genomic distribution of DhMR windows",
    subtitle = "1 kb windows, |Delta 5hmC| > 0.05, minimum CpG coverage = 30x",
    x = NULL,
    y = "Percent",
    fill = NULL
  ) +

  theme_bw(base_size = 15) +

  theme(
    plot.title = element_text(
      hjust = 0.5,
      face = "bold",
      size = 22
    ),

    plot.subtitle = element_text(
      hjust = 0.5,
      size = 13
    ),

    axis.text.x = element_text(
      face = "bold",
      size = 16
    ),

    axis.title.y = element_text(
      face = "bold",
      size = 18
    ),

    legend.position = "right",

    legend.text = element_text(
      size = 14
    ),

    plot.margin = margin(
      t = 20,
      r = 20,
      b = 15,
      l = 20
    )
  )

# =====================================================
# SAVE
# =====================================================

ggsave(
  out_pdf,
  p,
  width = 8,
  height = 6
)

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