#!/usr/bin/env Rscript

suppressPackageStartupMessages({
  library(data.table)
  library(ggplot2)
  library(patchwork)
  library(ggrepel)
  library(ggtext)
  library(grid)
})

setDTthreads(8)

base_dir <- "/BLUES/eric/ONT_WGBS/Heatmap"
outdir   <- "/BLUES/eric/ONT_WGBS/Figure_3/MA_Plot"
dir.create(outdir, showWarnings = FALSE, recursive = TRUE)

dna_file <- file.path(base_dir, "DNA_methylation_matrix_FULL_no_rRNA_tRNA.tsv")
rna_fc_file <- "/BLUES/eric/ONT_WGBS/Figure_5/MA_plot/RNAseq_log2FC.tsv"

pdf_out <- file.path(outdir, "Grid_based_classification_2.pdf")
tsv_out <- file.path(outdir, "Grid_based_classification_2.tsv")

dna <- fread(dna_file)
rna_fc <- fread(rna_fc_file)

dt <- merge(
  dna[, .(
    subfamily,
    DNA_NP = Naive - Primed,
    DNA_TN = TSC - Naive
  )],
  rna_fc[, .(
    subfamily,
    RNA_NP = log2FC_NP,
    RNA_TN = log2FC_TN
  )],
  by = "subfamily"
)

dt <- dt[
  is.finite(DNA_NP) & is.finite(RNA_NP) &
    is.finite(DNA_TN) & is.finite(RNA_TN)
]

# ======================
# CUTOFFS
# ======================

rna_cut <- 2
dna_cut <- 0.1

# ======================
# REGION CLASSIFICATION
# ======================

dt[, region_PN :=
     fifelse(RNA_NP > rna_cut, "1",
     fifelse(RNA_NP > 0,       "2",
     fifelse(RNA_NP >= -rna_cut, "3", "4")))
]

dt[, region_TN := "none"]

dt[RNA_TN < -rna_cut & DNA_TN >=  dna_cut, region_TN := "A"]
dt[RNA_TN < -rna_cut & DNA_TN >  -dna_cut & DNA_TN < dna_cut, region_TN := "B"]
dt[RNA_TN >  rna_cut & DNA_TN >= -dna_cut & DNA_TN <= dna_cut, region_TN := "C"]
dt[RNA_TN >= -rna_cut & RNA_TN <= rna_cut &
     DNA_TN >= -dna_cut & DNA_TN <= dna_cut, region_TN := "D"]

dt[, group := "none"]

dt[region_PN == "1" & region_TN == "A",
   group := "DNA methylation dependent"]

dt[region_PN == "1" & region_TN == "B",
   group := "Naive specific active"]

dt[region_TN == "C" & region_PN %in% c("2", "3"),
   group := "TSC specific"]

dt[region_TN == "B" & region_PN %in% c("2", "3"),
   group := "stem-cell specific"]

dt[region_PN == "4" & region_TN == "D",
   group := "primed specific"]

dt[region_PN == "4" & region_TN == "C",
   group := "Naive repressed"]

fwrite(dt[group != "none"], tsv_out, sep = "\t")

highlight_dt <- dt[group != "none"]

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

cols <- c(
  "DNA methylation dependent" = "#1f77b4",
  "Naive specific active"     = "#ff7f0e",
  "TSC specific"              = "#d62728",
  "stem-cell specific"        = "#2ca02c",
  "primed specific"           = "#9467bd",
  "Naive repressed"           = "#8c564b"
)

# ======================
# AXIS LABEL FUNCTIONS
# ======================

x_label_fun <- function(x) {
  ifelse(x %in% c(-rna_cut, rna_cut),
         paste0("<span style='color:red'>", x, "</span>"),
         x)
}

y_label_fun <- function(y) {
  ifelse(y %in% c(-dna_cut, dna_cut),
         paste0("<span style='color:red'>", y, "</span>"),
         y)
}

# ======================
# Naive - Primed
# ======================

plot_np <- ggplot() +
  geom_point(data = dt, aes(RNA_NP, DNA_NP),
             color = "grey75", size = 0.8, alpha = 0.7) +
  geom_point(data = highlight_dt,
             aes(RNA_NP, DNA_NP, fill = group),
             shape = 21, color = "black", size = 1.5) +
  geom_text_repel(
    data = highlight_dt,
    aes(RNA_NP, DNA_NP, label = subfamily, color = group),
    size = 1.3, max.overlaps = Inf,
    box.padding = 0.08, point.padding = 0.03,
    force = 12
  ) +
  geom_vline(xintercept = c(-rna_cut, rna_cut), linetype = "dashed", color = "red") +
  geom_hline(yintercept = 0) +
  geom_vline(xintercept = 0) +
  scale_fill_manual(values = cols) +
  scale_color_manual(values = cols, guide = "none") +
  scale_x_continuous(
    limits = c(-6, 6),
    breaks = c(-6, -rna_cut, 0, rna_cut, 6),
    labels = x_label_fun
  ) +
  labs(
    title = "Naive - Primed",
    x = "RNA log2FC",
    y = "DNA methylation difference"
  ) +
  theme_classic(base_size = 12) +
  theme(
    plot.title = element_text(face = "bold"),
    legend.position = "right",
    axis.text.x = ggtext::element_markdown(),
    axis.text.y = ggtext::element_markdown()
  )

# ======================
# TSC - Naive
# ======================

plot_tn <- ggplot() +
  geom_point(data = dt, aes(RNA_TN, DNA_TN),
             color = "grey75", size = 0.8, alpha = 0.7) +
  geom_point(data = highlight_dt,
             aes(RNA_TN, DNA_TN, fill = group),
             shape = 21, color = "black", size = 1.5) +
  geom_text_repel(
    data = highlight_dt,
    aes(RNA_TN, DNA_TN, label = subfamily, color = group),
    size = 1.3, max.overlaps = Inf,
    box.padding = 0.08, point.padding = 0.03,
    force = 12
  ) +
  geom_vline(xintercept = c(-rna_cut, rna_cut), linetype = "dashed", color = "red") +
  geom_hline(yintercept = c(-dna_cut, dna_cut), linetype = "dashed", color = "red") +
  geom_hline(yintercept = 0) +
  geom_vline(xintercept = 0) +
  scale_fill_manual(values = cols) +
  scale_color_manual(values = cols, guide = "none") +
  scale_x_continuous(
    limits = c(-6, 6),
    breaks = c(-6, -rna_cut, 0, rna_cut, 6),
    labels = x_label_fun
  ) +
  scale_y_continuous(
    breaks = c(-0.2, -dna_cut, 0, dna_cut, 0.5),
    labels = y_label_fun
  ) +
  labs(
    title = "TSC - Naive",
    x = "RNA log2FC",
    y = "DNA methylation difference"
  ) +
  theme_classic(base_size = 12) +
  theme(
    plot.title = element_text(face = "bold"),
    legend.position = "right",
    axis.text.x = ggtext::element_markdown(),
    axis.text.y = ggtext::element_markdown()
  )

fig <- (plot_np | plot_tn) + plot_layout(guides = "collect")

ggsave(pdf_out, fig, width = 12, height = 5.2, useDingbats = FALSE)

print(highlight_dt[, .N, by = group][order(group)])
message("done")
message("[OK] wrote: ", pdf_out)
message("[OK] wrote: ", tsv_out)