#!/usr/bin/env Rscript

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

setDTthreads(8)

# =====================================================
# INPUT / OUTPUT
# =====================================================

infile <- "/BLUES/eric/ONT_WGBS/Heatmap/DNA_methylation_matrix_FULL_no_rRNA_tRNA.tsv"

outdir <- "/BLUES/eric/ONT_WGBS/Figure_2/Violin_plot"
dir.create(outdir, recursive = TRUE, showWarnings = FALSE)

out_pdf <- file.path(
  outdir,
  "TE_class_DNA_methylation_violin.pdf"
)

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

dt <- fread(infile)

dt <- dt[!grepl("^5S$|^7SK$|^7SLRNA$", subfamily)]

# =====================================================
# DEFINE TE CLASS
# =====================================================

dt[, class := fifelse(
  grepl("^L1|^L2|^CR1|^RTE|LINE", subfamily, ignore.case = TRUE),
  "LINE",

  fifelse(
    grepl("^Alu|^MIR|^SINE", subfamily, ignore.case = TRUE),
    "SINE",

    fifelse(
      grepl("^ERV|LTR|^MER", subfamily, ignore.case = TRUE),
      "LTR",

      fifelse(
        grepl("^hAT|^TcMar|DNA", subfamily, ignore.case = TRUE),
        "DNA",
        "Other"
      )
    )
  )
)]

# =====================================================
# LONG FORMAT
# =====================================================

df <- melt(
  dt,
  id.vars = c("subfamily", "class"),
  measure.vars = c("Primed", "Naive", "TSC"),
  variable.name = "celltype",
  value.name = "beta"
)

df[, celltype := factor(
  celltype,
  levels = c("Primed", "Naive", "TSC")
)]

df[, class := factor(
  class,
  levels = c("LINE", "SINE", "LTR", "DNA", "Other")
)]

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

p <- ggplot(
  df,
  aes(
    x = celltype,
    y = beta,
    fill = celltype
  )
) +

  geom_violin(
    scale = "width",
    trim = TRUE,
    color = "black",
    linewidth = 0.2
  ) +

  geom_boxplot(
    width = 0.12,
    outlier.size = 0.2,
    fill = "white"
  ) +

  facet_wrap(
    ~ class,
    nrow = 1,
    scales = "fixed"
  ) +

  scale_fill_manual(
    values = c(
      "Primed" = "#d7301f",
      "Naive"  = "#fee08b",
      "TSC"    = "#4575b4"
    )
  ) +

  labs(
    title = "DNA methylation distribution across TE classes",
    subtitle = "Violin plot of TE subfamily methylation beta values",
    x = NULL,
    y = "DNA methylation (beta)"
  ) +

  coord_cartesian(
    ylim = c(0, 1)
  ) +

  theme_bw(base_size = 12) +

  theme(
    strip.background = element_rect(
      fill = "grey95",
      color = "black"
    ),

    strip.text = element_text(
      face = "bold",
      size = 12
    ),

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

    axis.text.y = element_text(
      size = 10
    ),

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

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

    legend.position = "none"
  )

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

ggsave(
  out_pdf,
  p,
  width = 13,
  height = 4.8,
  useDingbats = FALSE
)

cat("Saved:\n")
cat(out_pdf, "\n")