if (!requireNamespace("BiocManager", quietly = TRUE))
  install.packages("BiocManager")

library(rtracklayer)
library(GenomicRanges)
library(ChIPseeker)
library(TxDb.Hsapiens.UCSC.hg38.knownGene)
library(org.Hs.eg.db)
library(ggplot2)
library(pheatmap)
library(dplyr)
library(tibble)
library(purrr)
library(ComplexHeatmap)
library(UpSetR)
BiocManager::install("TxDb.Hsapiens.UCSC.hg38.knownGene")


# Import dataset ----------------------------------------------------------
peak_dir <- "/gpfs/data/aifantislab/public/mingjun_cutrun/peaks/sns_based/hg38/peaks/macs3_narrow"

files <- c(
  NC_KMT2A  = file.path(peak_dir, "NC_KMT2A_peaks.narrowPeak"),
  NC_Menin  = file.path(peak_dir, "NC_Menin_peaks.narrowPeak"),
  NC_UBTF   = file.path(peak_dir, "NC_UBTF_peaks.narrowPeak"),
  NC_PolII  = file.path(peak_dir, "NC_Pol_II_peaks.narrowPeak"),
  NC_H3K4me3 = file.path(peak_dir, "NC_H3K4me3_peaks.narrowPeak"),
  NC_H3K27Ac = file.path(peak_dir, "NC_H3K27Ac_peaks.narrowPeak"),
  TD_KMT2A  = file.path(peak_dir, "TD_KMT2A_peaks.narrowPeak"),
  TD_Menin  = file.path(peak_dir, "TD_Menin_peaks.narrowPeak"),
  TD_UBTF   = file.path(peak_dir, "TD_UBTF_peaks.narrowPeak"),
  TD_PolII  = file.path(peak_dir, "TD_Pol_II_peaks.narrowPeak"),
  TD_H3K4me3 = file.path(peak_dir, "TD_H3K4me3_peaks.narrowPeak"),
  TD_H3K27Ac = file.path(peak_dir, "TD_H3K27Ac_peaks.narrowPeak")
  )

read_peak <- function(f) {
  df <- read.table(f, sep = "\t", header = FALSE, stringsAsFactors = FALSE)
  colnames(df) <- c("chrom", "start", "end", "name", "score", "strand",
                    "signalValue", "pValue", "qValue", "peak")
  
  gr <- GRanges(
    seqnames = df$chrom,
    ranges = IRanges(start = df$start + 1, end = df$end),
    strand = "*"
  )
  
  mcols(gr) <- df[, c("name", "score", "signalValue", "pValue", "qValue", "peak")]
  sort(gr)
}

peak_list <- lapply(files, read_peak)
peak_list <- GRangesList(peak_list)


# PeakAnnotation with ChIPseeker ------------------------------------------
txdb <- TxDb.Hsapiens.UCSC.hg38.knownGene
#Annotate a specific sample
anno_NC_Menin <- annotatePeak(
  peak_list$NC_Menin,
  TxDb = txdb,
  tssRegion = c(-3000, 3000),
  annoDb = "org.Hs.eg.db"
)
head(as.data.frame(anno_NC_Menin))

#Annotate all samples
anno_list <- lapply(peak_list, function(gr) {
  annotatePeak(
    gr,
    TxDb = txdb,
    tssRegion = c(-3000, 3000),
    annoDb = "org.Hs.eg.db"
  )
})

#Save the annotated peaks
for (nm in names(anno_list)) {
  df <- as.data.frame(anno_list[[nm]])
  write.csv(df, paste0(nm, "_macs3_narrow_ChIPseeker_annotation.csv"), row.names = FALSE)
}

#Visualization of annotation
plotAnnoBar(anno_list)
ggsave("annotation_distribution_macs3_narrow.png", width = 5, height = 4, dpi = 600)
plotDistToTSS(anno_list,
              title = "Distribution of CUT&RUN peaks relative to TSS")
#Extract annotated genes
annotated_genes <- lapply(anno_list, function(x) {
  df <- as.data.frame(x)
  unique(df$SYMBOL[!is.na(df$SYMBOL)])
})

lapply(annotated_genes, length)

#Visualization of identified peaks for each factor
# Peak count summary -------------------------------------------------------
peak_counts_NC <- tibble(
  factor = c("KMT2A", "Menin", "UBTF", "PolII", "H3K4me3", "H3K27Ac"),
  group  = "NC",
  peaks  = c(
    length(peak_list$NC_KMT2A),
    length(peak_list$NC_Menin),
    length(peak_list$NC_UBTF),
    length(peak_list$NC_PolII),
    length(peak_list$NC_H3K4me3),
    length(peak_list$NC_H3K27Ac)
  )
)

peak_counts_TD <- tibble(
  factor = c("KMT2A", "Menin", "UBTF", "PolII", "H3K4me3", "H3K27Ac"),
  group  = "TD",
  peaks  = c(
    length(peak_list$TD_KMT2A),
    length(peak_list$TD_Menin),
    length(peak_list$TD_UBTF),
    length(peak_list$TD_PolII),
    length(peak_list$TD_H3K4me3),
    length(peak_list$TD_H3K27Ac)
  )
)

peak_counts <- bind_rows(peak_counts_NC, peak_counts_TD)

write.csv(peak_counts, "cutrun_peak_counts.csv", row.names = FALSE)

peak_counts$factor <- factor(
  peak_counts$factor,
  levels = c("KMT2A", "Menin", "UBTF", "PolII", "H3K4me3", "H3K27Ac")
)

peak_counts$factor <- factor(
  peak_counts$factor,
  levels = c("KMT2A", "Menin", "UBTF", "PolII", "H3K4me3", "H3K27Ac")
)

#Visualization
library(tidyverse)
library(patchwork)

peak_counts$factor <- factor(
  peak_counts$factor,
  levels = c("KMT2A", "Menin", "UBTF", "PolII", "H3K4me3", "H3K27Ac")
)

make_panel <- function(fac) {
  df <- peak_counts %>% filter(factor == fac)
  
  y_max <- max(df$peaks) * 1.25
  
  ggplot(df, aes(x = group, y = peaks, fill = group)) +
    geom_col(width = 0.65, color = "black", linewidth = 0.2) +
    geom_text(
      aes(label = scales::comma(peaks)),
      vjust = -0.35,
      size = 3
    ) +
    scale_fill_manual(values = c("NC" = "#4C72B0", "TD" = "#CB4C4E")) +
    coord_cartesian(ylim = c(0, y_max), clip = "off") +
    labs(x = NULL, y = "Peaks", title = fac) +
    theme_classic(base_size = 11) +
    theme(
      legend.position = "none",
      plot.title = element_text(face = "bold", hjust = 0.5),
      axis.text.x = element_text(face = "bold")
    )
}

p1 <- make_panel("KMT2A")
p2 <- make_panel("Menin")
p3 <- make_panel("UBTF")
p4 <- make_panel("PolII")
p5 <- make_panel("H3K4me3")
p6 <- make_panel("H3K27Ac")

p_all <- (p1 + p2 + p3) / (p4 + p5 + p6) +
  plot_layout(guides = "collect") &
  theme(legend.position = "top")
p_all

ggsave("NC_vs_TD_peak_counts_3x2.png", p_all, width = 6, height = 5, dpi = 600)

# Peak Overlap Analysis ---------------------------------------------------
library(UpSetR)
BiocManager::install("ComplexUpset")
library(ComplexUpset)
library(ggplot2)
library(dplyr)

multi <- read.table(
  "/gpfs/data/aifantislab/public/mingjun_cutrun/peaks/nextflow/hg38/peaks/macs2_narrow/TD_overlap_analysis/multiinter/TD_multiinter.tsv",
  sep = "\t",
  header = FALSE,
  stringsAsFactors = FALSE
)

head(multi)
dim(multi)
colnames(multi) <- c(
  "chr", "start", "end", "n_sets", "set_list",
  "H3K4me3", "H3K27Ac", "KMT2A", "Menin", "PolII", "UBTF"
)

head(multi)
upset_df <- multi %>%
  mutate(region_id = paste0(chr, ":", start, "-", end)) %>%
  dplyr::select(region_id, H3K4me3, H3K27Ac, KMT2A, Menin, PolII, UBTF)

head(upset_df)

set_cols <- c("H3K4me3", "H3K27Ac", "KMT2A", "Menin", "PolII", "UBTF")

upset_df[set_cols] <- lapply(upset_df[set_cols], as.integer)

str(upset_df)

UpSetR::upset(
  upset_df,
  sets = c("H3K4me3", "H3K27Ac", "KMT2A", "Menin", "PolII", "UBTF"),
  keep.order = TRUE,
  order.by = "freq",
  mb.ratio = c(0.6, 0.4),
  text.scale = 1.2
)


# Peak_annotation analysis ------------------------------------------------
#For UBTF-TD peaks
library(tidyverse)
library(ggrepel)

ubtf <- read_csv("TD_UBTF_macs3_narrow_ChIPseeker_annotation.csv")

ubtf_volcano <- ubtf %>%
  arrange(desc(qValue)) %>%                # highest -log10(FDR) first
  mutate(
    rank_q = row_number(),
    top40  = rank_q <= 40                  # logical flag for top 30 peaks
  )

ggplot(ubtf_volcano,
       aes(x = signalValue, y = qValue)) +
  geom_point(aes(color = top40), alpha = 0.7, size = 1.2) +
  geom_text_repel(
    data = ubtf_volcano %>% filter(top40),
    aes(label = SYMBOL),
    size = 2.5,
    max.overlaps = 100
  ) +
  scale_color_manual(
    values = c(`FALSE` = "grey70", `TRUE` = "#CB4C4E"),
    labels = c("Other peaks", "Top 30 qValue"),
    name   = NULL
  ) +
  labs(
    x = "MACS3 signalValue (fold-enrichment at summit)",
    y = expression(-log[10]("FDR (qValue)")),
    title = "UBTF-TD peaks: top 30 most significant (labeled)"
  ) +
  theme_classic(base_size = 11) +
  theme(
    legend.position = "right",
    plot.title = element_text(face = "bold", hjust = 0.5)
  )
ggsave("UBTF-TD_Peaks.png", width = 5, height = 5, dpi = 300)

# Gene level overlapping --------------------------------------------------
library(tidyverse)

# Helper to load and make a gene set (optionally promoter-only)
load_gene_set <- function(file, promoter_only = TRUE) {
  df <- read_csv(file)
  if (promoter_only) {
    df <- df %>% filter(grepl("Promoter", annotation, ignore.case = TRUE))
  }
  df %>%
    filter(!is.na(SYMBOL), SYMBOL != "") %>%
    pull(SYMBOL) %>%
    unique()
}

genes_UBTF    <- load_gene_set("TD_UBTF_macs3_narrow_ChIPseeker_annotation.csv")
genes_UBTF
genes_KMT2A   <- load_gene_set("TD_KMT2A_macs3_narrow_ChIPseeker_annotation.csv")
genes_Menin   <- load_gene_set("TD_Menin_macs3_narrow_ChIPseeker_annotation.csv")
genes_PolII   <- load_gene_set("TD_PolII_macs3_narrow_ChIPseeker_annotation.csv")
genes_H3K4me3 <- load_gene_set("TD_H3K4me3_macs3_narrow_ChIPseeker_annotation.csv")
genes_H3K27ac <- load_gene_set("TD_H3K27Ac_macs3_narrow_ChIPseeker_annotation.csv")

# Global overlaps
ubtf_kmt2a      <- intersect(genes_UBTF, genes_KMT2A)
ubtf_menin      <- intersect(genes_UBTF, genes_Menin)
ubtf_kmt2a_menin <- Reduce(intersect, list(genes_UBTF, genes_KMT2A, genes_Menin))

ubtf_polii_h3k4 <- Reduce(intersect, list(genes_UBTF, genes_PolII, genes_H3K4me3))

lengths_list <- list(
  UBTF              = length(genes_UBTF),
  KMT2A             = length(genes_KMT2A),
  Menin             = length(genes_Menin),
  PolII             = length(genes_PolII),
  H3K4me3           = length(genes_H3K4me3),
  UBTF_KMT2A        = length(ubtf_kmt2a),
  UBTF_Menin        = length(ubtf_menin),
  UBTF_KMT2A_Menin  = length(ubtf_kmt2a_menin),
  UBTF_PolII_H3K4me3 = length(ubtf_polii_h3k4)
)

lengths_list

library(VennDiagram)

gene_lists <- list(
  UBTF  = genes_UBTF,
  Menin = genes_Menin,
  KMT2A = genes_KMT2A,
  H3K27ac = genes_H3K27ac
)

venn.plot <- venn.diagram(
  x = gene_lists,
  filename = NULL,
  fill = c("#4C72B0", "#CB4C4E", "#55A868", "#C4B24D"),
  alpha = 0.5,
  cex = 1.0,
  cat.cex = 1.2,
  cat.col = "black",
  main = "Promoter-level co-bound genes\nUBTF-TD, Menin, KMT2A, PolII"
)

# Save to file
png("Venn_UBTF_Menin_KMT2A_H3K27ac.png", width = 1200, height = 1000, res = 300)
grid::grid.draw(venn.plot)
dev.off()

##Check specific genes 
ubtf_kmt2a_h3k27 <- Reduce(
  intersect,
  list(genes_UBTF, genes_KMT2A, genes_H3K27ac)
)

ubtf_kmt2a_h3k27_no_menin <- setdiff(
  ubtf_kmt2a_h3k27,
  genes_Menin
)

length(ubtf_kmt2a_h3k27_no_menin)
ubtf_kmt2a_h3k27_no_menin

ubtf_menin_h3k27 <- Reduce(
  intersect,
  list(genes_UBTF, genes_Menin, genes_H3K27ac)
)

ubtf_menin_h3k27_no_kmt2a <- setdiff(
  ubtf_menin_h3k27,
  genes_KMT2A
)

length(ubtf_menin_h3k27_no_kmt2a)
ubtf_menin_h3k27_no_kmt2a

ubtf_h3k27 <- intersect(genes_UBTF, genes_H3K27ac)

ubtf_h3k27_only <- ubtf_h3k27 %>%
  setdiff(genes_KMT2A) %>%
  setdiff(genes_Menin)

length(ubtf_h3k27_only)
ubtf_h3k27_only

# Comparing NC and TD peaks -----------------------------------------------
library(tidyverse)
# Helper: gene-level promoter set
load_promoter_genes <- function(file) {
  read_csv(file) %>%
    filter(
      !is.na(SYMBOL),
      SYMBOL != "",
      grepl("Promoter", annotation, ignore.case = TRUE)
    ) %>%
    pull(SYMBOL) %>%
    unique()
}
genes_KMT2A_NC <- load_promoter_genes("NC_KMT2A_macs3_narrow_ChIPseeker_annotation.csv")
genes_KMT2A_TD <- load_promoter_genes("TD_KMT2A_macs3_narrow_ChIPseeker_annotation.csv")

kmt2a_shared   <- intersect(genes_KMT2A_NC, genes_KMT2A_TD)
kmt2a_TD_only  <- setdiff(genes_KMT2A_TD, genes_KMT2A_NC)
kmt2a_NC_only  <- setdiff(genes_KMT2A_NC, genes_KMT2A_TD)

length(kmt2a_TD_only)
length(kmt2a_NC_only)
length(kmt2a_shared)

genes_UBTF_TD <- load_promoter_genes("TD_UBTF_macs3_narrow_ChIPseeker_annotation.csv")

kmt2a_TD_only_UBTF <- intersect(kmt2a_TD_only, genes_UBTF_TD)
length(genes_UBTF_TD)
length(kmt2a_TD_only_UBTF)
head(kmt2a_TD_only_UBTF)

summary_df <- tibble(
  category = c("KMT2A TD-only", "KMT2A TD-only + UBTF"),
  n_genes  = c(length(kmt2a_TD_only), length(kmt2a_TD_only_UBTF))
)

p <- ggplot(summary_df, aes(x = category, y = n_genes)) +
  geom_col(fill = "#4C72B0") +
  geom_text(aes(label = n_genes), vjust = -0.3) +
  labs(
    x = NULL,
    y = "Number of promoter genes",
    title = "TD-specific KMT2A promoter binding\nand overlap with UBTF-TD"
  ) +
  theme_classic(base_size = 11) +
  theme(axis.text.x = element_text(angle = 20, hjust = 1))
p

#Venn Diagram
library(VennDiagram)

gene_lists_3 <- list(
  KMT2A_NC = genes_KMT2A_NC,
  KMT2A_TD = genes_KMT2A_TD,
  UBTF_TD  = genes_UBTF_TD
)

venn3 <- venn.diagram(
  x = gene_lists_3,
  filename = NULL,
  fill = c("#4C72B0", "#55A868", "#CB4C4E"),
  alpha = 0.6,
  cex = 1.1,
  cat.cex = 1.3,
  cat.col = "black",
  main = "Promoter-level overlap\nNC KMT2A, TD KMT2A, UBTF-TD"
)

png("Venn_KMT2A_NC_TD_UBTF_TD.png", width = 1000, height = 1000, res = 150)
grid::grid.draw(venn3)
dev.off()


# Overlapping of KMT2A/Menin in WT and TD mouse ---------------------------
genes_KMT2A_NC <- load_promoter_genes("NC_KMT2A_macs3_narrow_ChIPseeker_annotation.csv")
genes_KMT2A_TD <- load_promoter_genes("TD_KMT2A_macs3_narrow_ChIPseeker_annotation.csv")
genes_Menin_TD <- load_promoter_genes("TD_Menin_macs3_narrow_ChIPseeker_annotation.csv")
genes_Menin_NC <- load_promoter_genes("NC_Menin_macs3_narrow_ChIPseeker_annotation.csv")
# Shared KMT2A promoters between NC and TD
kmt2a_shared <- intersect(genes_KMT2A_NC, genes_KMT2A_TD)

# TD-only KMT2A promoters
kmt2a_TD_only <- setdiff(genes_KMT2A_TD, genes_KMT2A_NC)

# KMT2A_NC with Menin in TD
kmt2a_NC_Menin <- intersect(genes_KMT2A_NC, genes_Menin_TD)

# KMT2A_TD with Menin in TD
kmt2a_TD_Menin <- intersect(genes_KMT2A_TD, genes_Menin_TD)

# TD-only KMT2A + Menin (new Menin-dependent sites in TD)
kmt2a_TD_only_Menin <- intersect(kmt2a_TD_only, genes_Menin_TD)

length(kmt2a_TD_only_Menin)
head(kmt2a_TD_only_Menin)

gene_lists_3 <- list(
  KMT2A_NC = genes_KMT2A_NC,
  KMT2A_TD = genes_KMT2A_TD,
  Menin_TD = genes_Menin_TD,
  Menin_NC = genes_Menin_NC
)

venn3 <- venn.diagram(
  x = gene_lists_3,
  filename = NULL,
  fill = c("#4C72B0", "#CB4C4E", "#55A868", "#C4B24D"),
  alpha = 0.6,
  cex = 1.1,
  cat.cex = 1.3,
  cat.col = "black",
  main = "Promoter-level overlap KMT2A_NC, TD, Menin_NC, TD"
)

png("Venn_KMT2A_NC_TD_Menin_TD.png", width = 1000, height = 900, res = 150)
grid::grid.draw(venn3)
dev.off()


# Analysis of UBTF-TD specific peaks --------------------------------------
kmt2a_TD_only <- setdiff(genes_KMT2A_TD, genes_KMT2A_NC)

length(kmt2a_TD_only)
head(kmt2a_TD_only)

write.table(
  kmt2a_TD_only,
  "KMT2A_TD_specific_promoter_genes.txt",
  quote = FALSE, row.names = FALSE, col.names = FALSE
)

#Pathway analysis of KMT2A specific TD genes
library(msigdbr)
library(clusterProfiler)
library(dplyr)

# Hallmark sets for human
msig_h <- msigdbr(species = "Homo sapiens", category = "H")

# Convert to TERM2GENE format (pathway ↔ gene symbol)
hallmark_t2g <- msig_h %>%
  dplyr::select(gs_name, gene_symbol) %>%
  distinct()

# KMT2A TD-specific genes as symbols
genes_kmt2a_td_only <- kmt2a_TD_only  # character vector of symbols

# Optional: define a universe (background)
# e.g., all genes with any KMT2A peak in NC or TD:
gene_universe <- union(genes_KMT2A_NC, genes_KMT2A_TD)

hallmark_enrich <- enricher(
  gene      = genes_kmt2a_td_only,
  TERM2GENE = hallmark_t2g,
  pAdjustMethod = "BH",
  qvalueCutoff  = 0.05
)

# Inspect top pathways
head(as.data.frame(hallmark_enrich))