#!/usr/bin/env Rscript
suppressPackageStartupMessages({
  library(ggplot2)
  library(ggrepel)
})

# ============================================================
# 0. CONFIG
# ============================================================
fdr_cutoff <- 0.01
fc_cutoff  <- log2(1.5)
outdir     <- "/BRC/yan/heart/analysis/"

files <- list(
  "Bulk heart" = "/BRC/yan/heart/analysis/sn_ATAC_analysis/heart_DAR_female_vs_male_RUVr_k6.txt",
  "snATAC CM"  = "/BRC/yan/heart/snACAT_CM/DAR_results/snACAT_CM_DAR_female_vs_male_RUVr_k6.txt"
)

# ============================================================
# 1. PLOT FUNCTION
# ============================================================
plot_volcano <- function(dat, title) {

  # -log10 FDR; cap zeros at machine epsilon to avoid -Inf
  dat$FDR[dat$FDR == 0] <- .Machine$double.xmin
  dat$neglog10FDR <- -log10(dat$FDR)

  # direction labels
  dat$color <- "non-DAR"
  dat$color[dat$FDR < fdr_cutoff & dat$logFC >  fc_cutoff] <- "MORE"
  dat$color[dat$FDR < fdr_cutoff & dat$logFC < -fc_cutoff] <- "LESS"
  dat$color <- factor(dat$color, levels = c("MORE", "LESS", "non-DAR"))

  n_more <- sum(dat$color == "MORE")
  n_less <- sum(dat$color == "LESS")

  # cap y-axis at 99th percentile to avoid extreme outliers collapsing the plot
  y_max <- quantile(dat$neglog10FDR, 0.999)

  colors <- c(MORE = "#E64B35", LESS = "#4DBBD5", "non-DAR" = "grey75")

  ggplot(dat, aes(x = logFC, y = neglog10FDR, color = color)) +
    geom_point(size = 0.6, alpha = 0.5) +
    geom_hline(yintercept = -log10(fdr_cutoff),
               linetype = "dashed", linewidth = 0.5, color = "grey40") +
    geom_vline(xintercept = c(-fc_cutoff, fc_cutoff),
               linetype = "dashed", linewidth = 0.5, color = "grey40") +
    scale_color_manual(values = colors,
                       labels = c(
                         MORE     = paste0("MORE (n=", n_more, ")"),
                         LESS     = paste0("LESS (n=", n_less, ")"),
                         "non-DAR" = paste0("non-DAR (n=", sum(dat$color == "non-DAR"), ")")
                       ),
                       name = "Direction") +
    coord_cartesian(ylim = c(0, y_max)) +
    labs(
      title = paste0(title, "\n(female vs male, RUVr k=6)"),
      x     = "Log2FC (female / male)",
      y     = "-log10(FDR)"
    ) +
    theme_bw() +
    theme(
      plot.title      = element_text(hjust = 0.5, face = "bold", size = 12),
      axis.title      = element_text(size = 10),
      legend.position = "right",
      legend.title    = element_text(face = "bold"),
      panel.grid.minor = element_blank()
    )
}

# ============================================================
# 2. MAKE AND SAVE PLOTS
# ============================================================
plots <- list()

for (label in names(files)) {
  cat("Processing:", label, "\n")
  dat <- read.table(files[[label]], header = TRUE, sep = "\t")
  cat("  Total peaks:", nrow(dat), "\n")
  plots[[label]] <- plot_volcano(dat, label)
}

# individual PDFs
for (label in names(plots)) {
  fname <- gsub(" ", "_", label)
  out   <- paste0(outdir, "volcano_DAR_", fname, ".pdf")
  ggsave(out, plots[[label]], width = 6, height = 5)
  cat("Saved:", out, "\n")
}

# combined PDF (side by side)
library(patchwork)
combined <- plots[["Bulk heart"]] + plots[["snATAC CM"]] +
  plot_layout(ncol = 2)
out_combined <- paste0(outdir, "volcano_DAR_combined.pdf")
ggsave(out_combined, combined, width = 12, height = 5)
cat("Saved:", out_combined, "\n")

cat("[Done]\n")