#!/usr/bin/env Rscript
# =============================================================================
# GSEA — BPA all tissue (blood / brain / liver) × sex × dose
# Usage: Rscript run_gsea_all_tissue.R
# =============================================================================

suppressPackageStartupMessages({
  library(dplyr)
  library(tidyverse)
  library(clusterProfiler)
  library(org.Mm.eg.db)
  library(enrichplot)
  library(ggplot2)
  library(pheatmap)
})

# ── paths ─────────────────────────────────────────────────────────────────────
infile <- "/home/yan/TaRGET_II/BPA_brian/RNA/merged_all_gene_all_tissue.txt"
outdir <- "/home/yan/TaRGET_II/BPA_brian/GSEA/"
dir.create(outdir, recursive = TRUE, showWarnings = FALSE)

# ── load data ─────────────────────────────────────────────────────────────────
message("Loading: ", infile)
DEG <- read_tsv(infile, show_col_types = FALSE)
message("Loaded: ", nrow(DEG), " rows")

# ── add metadata columns ──────────────────────────────────────────────────────
meta <- DEG %>%
  mutate(
    dose       = str_extract(cond1, "10mg|10ug"),
    sex        = str_extract(cond1, "Female|Male"),
    gsea_group = paste0("Bartolomei_BPA", dose, "_", sex, "_adult_", tisssue)
  )

groups <- meta %>% distinct(gsea_group) %>% pull(gsea_group) %>% sort()
message("Total groups: ", length(groups))
message(paste(groups, collapse = "\n"))

# ── GSEA function ─────────────────────────────────────────────────────────────
run_gsea_go <- function(df, group_name) {

  message("\n========================================")
  message("Processing: ", group_name)

  dat <- df %>%
    filter(gsea_group == group_name) %>%
    filter(!is.na(log2FoldChange), !is.na(gene)) %>%
    group_by(gene) %>%
    slice_max(order_by = abs(log2FoldChange), n = 1, with_ties = FALSE) %>%
    ungroup()

  message("  Genes before ID mapping: ", nrow(dat))

  if (nrow(dat) == 0) {
    message("  Skipping — no data for this group")
    return(NULL)
  }

  gene_map <- bitr(
    dat$gene,
    fromType = "SYMBOL",
    toType   = "ENTREZID",
    OrgDb    = org.Mm.eg.db
  )

  dat2 <- dat %>%
    inner_join(gene_map, by = c("gene" = "SYMBOL")) %>%
    group_by(ENTREZID) %>%
    slice_max(order_by = abs(log2FoldChange), n = 1, with_ties = FALSE) %>%
    ungroup()

  message("  Genes after ID mapping: ", nrow(dat2))

  gene_list <- sort(setNames(dat2$log2FoldChange, dat2$ENTREZID), decreasing = TRUE)
  gene_list <- gene_list[!is.na(gene_list)]

  message("  Final gene_list length: ", length(gene_list))

  gse_go <- tryCatch(
    gseGO(
      geneList      = gene_list,
      OrgDb         = org.Mm.eg.db,
      keyType       = "ENTREZID",
      ont           = "BP",
      minGSSize     = 10,
      maxGSSize     = 500,
      pvalueCutoff  = 0.05,
      pAdjustMethod = "BH",
      verbose       = FALSE
    ),
    error = function(e) {
      message("  GSEA failed: ", e$message)
      return(NULL)
    }
  )

  if (is.null(gse_go) || nrow(as.data.frame(gse_go)) == 0) {
    message("  No significant terms found")
    return(NULL)
  }

  message("  Significant terms: ", nrow(as.data.frame(gse_go)))

  safe_name <- gsub("[^A-Za-z0-9_]", "_", group_name)

  # results table
  write.csv(
    as.data.frame(gse_go),
    file      = file.path(outdir, paste0(safe_name, "_GO_GSEA.csv")),
    row.names = FALSE
  )

  # dotplot
  pdf(file.path(outdir, paste0(safe_name, "_GO_dotplot.pdf")),
      width = 10, height = 20)
  print(
    dotplot(gse_go, showCategory = 20, split = ".sign") +
      facet_grid(. ~ .sign) +
      labs(title = group_name) +
      theme(plot.title = element_text(hjust = 0.5))
  )
  dev.off()

  # ridgeplot
  pdf(file.path(outdir, paste0(safe_name, "_GO_ridgeplot.pdf")),
      width = 10, height = 15)
  print(
    ridgeplot(gse_go) +
      labs(x = "log2FC", title = group_name) +
      theme(plot.title = element_text(hjust = 0.5))
  )
  dev.off()

  message("  Saved: ", safe_name)
  return(gse_go)
}

# ── run all groups ────────────────────────────────────────────────────────────
gsea_results <- lapply(groups, function(g) run_gsea_go(meta, g))
names(gsea_results) <- groups

saveRDS(gsea_results, file = file.path(outdir, "gsea_results_all_tissue.rds"))
message("\nSaved RDS: gsea_results_all_tissue.rds")

# ── merge summary ─────────────────────────────────────────────────────────────
gsea_summary <- bind_rows(
  lapply(names(gsea_results), function(g) {
    if (is.null(gsea_results[[g]])) return(NULL)
    df <- as.data.frame(gsea_results[[g]])
    if (nrow(df) == 0) return(NULL)
    df$group <- g
    df
  })
) %>%
  mutate(
    direction = ifelse(NES > 0, "Up_enriched", "Down_enriched"),
    tissue    = str_extract(group, "blood|brain|liver"),
    sex       = str_extract(group, "Female|Male"),
    dose      = str_extract(group, "10mg|10ug")
  )

write.csv(
  gsea_summary,
  file      = file.path(outdir, "All_groups_GO_GSEA_summary.csv"),
  row.names = FALSE
)
message("Saved: All_groups_GO_GSEA_summary.csv  (", nrow(gsea_summary), " terms)")

# ── top 20 per group × direction ──────────────────────────────────────────────
top_gsea <- gsea_summary %>%
  group_by(group, direction) %>%
  arrange(p.adjust, .by_group = TRUE) %>%
  slice_head(n = 20) %>%
  ungroup()

write.csv(
  top_gsea,
  file      = file.path(outdir, "Top20_per_group_GO_GSEA.csv"),
  row.names = FALSE
)
message("Saved: Top20_per_group_GO_GSEA.csv")

message("\n=== All done. Output: ", outdir, " ===")

# ── heatmap per tissue ────────────────────────────────────────────────────────
message("\nGenerating per-tissue NES heatmaps...")

for (tis in c("blood", "brain", "liver")) {

  tis_data <- top_gsea %>% filter(tissue == tis)

  if (nrow(tis_data) == 0) {
    message("  No data for tissue: ", tis, " — skipping")
    next
  }

  # NES matrix: term × group
  nes_mat <- tis_data %>%
    dplyr::select(Description, group, NES) %>%
    pivot_wider(names_from = group, values_from = NES) %>%
    column_to_rownames("Description") %>%
    as.matrix()

  nes_mat[is.na(nes_mat)] <- 0

  # remove rows with all zeros
  nes_mat <- nes_mat[rowSums(nes_mat != 0) > 0, , drop = FALSE]

  if (nrow(nes_mat) == 0) {
    message("  Empty NES matrix for tissue: ", tis, " — skipping")
    next
  }

  # column annotation: sex + dose
  col_anno <- data.frame(
    Sex  = ifelse(grepl("Female", colnames(nes_mat)), "Female", "Male"),
    Dose = ifelse(grepl("10mg",   colnames(nes_mat)), "10mg",   "10ug"),
    row.names = colnames(nes_mat)
  )

  anno_colors <- list(
    Sex  = c(Female = "#E87D72", Male   = "#56B4E9"),
    Dose = c(`10mg` = "#F0A500", `10ug` = "#7DC97D")
  )

  # clean column names for display
  colnames(nes_mat) <- colnames(nes_mat) %>%
    str_remove("Bartolomei_BPA") %>%
    str_remove(paste0("_adult_", tis))

  pdf(
    file.path(outdir, paste0("Heatmap_NES_", tis, ".pdf")),
    width = 8, height = max(6, nrow(nes_mat) * 0.2 + 3)
  )
  pheatmap(
    nes_mat,
    clustering_method        = "ward.D2",
    clustering_distance_rows = "euclidean",
    clustering_distance_cols = "euclidean",
    annotation_col           = col_anno,
    annotation_colors        = anno_colors,
    color  = colorRampPalette(c("#313695", "white", "#A50026"))(100),
    breaks = seq(-3, 3, length.out = 101),
    fontsize_row = 8,
    fontsize_col = 9,
    angle_col    = 45,
    border_color = NA,
    main = paste0("GSEA NES — ", str_to_title(tis), " (top 20 per group × direction)")
  )
  dev.off()

  message("  Saved: Heatmap_NES_", tis, ".pdf  (",
          nrow(nes_mat), " terms × ", ncol(nes_mat), " groups)")
}