#!/usr/bin/env Rscript

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

setDTthreads(8)

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

infile <- "annotation_summary_autosome_clean.tsv"

dt <- fread(infile)

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

# keep only genomic + TE panels
dt <- dt[type %in% c("genomic", "te")]

# =====================================================
# 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"
)

# =====================================================
# 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",
    subtitle="Autosome-only regions",
    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",
    subtitle="Autosome-only regions",
    x=NULL,
    y="Percent"
  ) +
  theme_bw(base_size=13)

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

final_plot <- p1 + p2 +
  plot_layout(ncol=2)

ggsave(
  "annotation_barplots_autosome_only.pdf",
  final_plot,
  width=10,
  height=5
)

ggsave(
  "annotation_barplots_autosome_only.png",
  final_plot,
  width=10,
  height=5,
  dpi=300
)

cat("\n[DONE] Autosome-only annotation plot saved.\n")