#!/usr/bin/env Rscript
###############################################################################
# 功能富集与可视化模板（clusterProfiler）
# 输入：deseq2_de.R 产出的 DE_genes_significant.tsv
# 输出：GO / KEGG 富集结果表 + 气泡图 + 条形图，附带 GSEA 示例
###############################################################################

suppressPackageStartupMessages({
  library(clusterProfiler)
  library(org.Hs.eg.db)      # 人类注释，换物种需换对应的 org.*.eg.db
  library(ggplot2)
  library(enrichplot)
})

## ------------------------------ 参数 ---------------------------------------
deg_file  <- "results/06_DE/DE_genes_significant.tsv"
all_file  <- "results/06_DE/DE_results_all.tsv"
out_dir   <- "results/07_enrichment"
padj_cut  <- 0.05
dir.create(out_dir, showWarnings = FALSE, recursive = TRUE)

## ---------------------------- 读入基因列表 ---------------------------------
deg <- read.delim(deg_file, stringsAsFactors = FALSE)
up_genes   <- deg$gene[deg$log2FoldChange > 0]
down_genes <- deg$gene[deg$log2FoldChange < 0]
message("上调基因 ", length(up_genes), "，下调基因 ", length(down_genes))

## 基因 ID 转换（Ensembl -> Entrez）
convert <- function(ids) {
  mapIds(org.Hs.eg.db, keys = ids, column = "ENTREZID",
         keytype = "ENSEMBL", multiVals = "first")
}

## ------------------------- 1. GO 富集（过表征分析） ------------------------
run_go <- function(ids, tag) {
  eg <- bitr(ids, fromType = "ENSEMBL", toType = "ENTREZID",
             OrgDb = org.Hs.eg.db)
  if (nrow(eg) == 0) return(NULL)

  for (ont in c("BP", "CC", "MF")) {
    res <- enrichGO(gene = eg$ENTREZID, OrgDb = org.Hs.eg.db,
                    ont = ont, keyType = "ENTREZID",
                    pAdjustMethod = "BH", pvalueCutoff = 0.05,
                    qvalueCutoff = 0.2, readable = TRUE)
    if (is.null(res) || nrow(as.data.frame(res)) == 0) next

    write.table(as.data.frame(res),
                file = file.path(out_dir, paste0("GO_", ont, "_", tag, ".tsv")),
                sep = "\t", quote = FALSE, row.names = FALSE)

    p <- dotplot(res, showCategory = 15) +
      ggtitle(paste0("GO ", ont, " enrichment (", tag, ")"))
    ggsave(file.path(out_dir, paste0("GO_", ont, "_", tag, ".pdf")),
           p, width = 8, height = 7)
  }
}

run_go(up_genes,   "up")
run_go(down_genes, "down")

## ---------------------------- 2. KEGG 富集 ---------------------------------
run_kegg <- function(ids, tag) {
  eg <- bitr(ids, fromType = "ENSEMBL", toType = "ENTREZID",
             OrgDb = org.Hs.eg.db)
  if (nrow(eg) == 0) return(NULL)

  kk <- enrichKEGG(gene = eg$ENTREZID, organism = "hsa",
                   keyType = "kegg", pvalueCutoff = 0.05,
                   qvalueCutoff = 0.2)
  if (is.null(kk) || nrow(as.data.frame(kk)) == 0) return(NULL)
  kk <- setReadable(kk, OrgDb = org.Hs.eg.db, keyType = "ENTREZID")

  write.table(as.data.frame(kk),
              file = file.path(out_dir, paste0("KEGG_", tag, ".tsv")),
              sep = "\t", quote = FALSE, row.names = FALSE)

  p <- barplot(kk, showCategory = 15, x = "GeneRatio") +
    ggtitle(paste0("KEGG enrichment (", tag, ")"))
  ggsave(file.path(out_dir, paste0("KEGG_", tag, ".pdf")),
         p, width = 8, height = 7)
}

run_kegg(up_genes,   "up")
run_kegg(down_genes, "down")

## --------------------------- 3. GSEA（进阶） -------------------------------
## GSEA 用全部基因按 log2FC 排序，不设阈值，优势是不丢信息
all_res <- read.delim(all_file, stringsAsFactors = FALSE)
all_res <- all_res[!is.na(all_res$padj) & !is.na(all_res$entrez),
                   c("gene", "log2FoldChange")]
gene_list <- all_res$log2FoldChange
names(gene_list) <- all_res$gene
gene_list <- sort(gene_list, decreasing = TRUE)

## 把 Ensembl 换成 Entrez
eg_all <- bitr(names(gene_list), fromType = "ENSEMBL", toType = "ENTREZID",
               OrgDb = org.Hs.eg.db)
gene_list <- gene_list[eg_all$ENSEMBL]
names(gene_list) <- eg_all$ENTREZID

gse <- gseGO(geneList = gene_list, OrgDb = org.Hs.eg.db, ont = "BP",
             keyType = "ENTREZID", pvalueCutoff = 0.05,
             minGSSize = 10, maxGSSize = 500)
if (!is.null(gse) && nrow(as.data.frame(gse)) > 0) {
  write.table(as.data.frame(gse),
              file = file.path(out_dir, "GSEA_GO_BP.tsv"),
              sep = "\t", quote = FALSE, row.names = FALSE)
  pdf(file.path(out_dir, "GSEA_ridgeplot.pdf"), width = 8, height = 8)
  print(ridgeplot(gse, showCategory = 15))
  dev.off()
}

message("富集分析完成，结果见：", normalizePath(out_dir))

###############################################################################
# 依赖安装：
#   BiocManager::install(c("clusterProfiler","org.Hs.eg.db","enrichplot"))
# 换物种：小鼠用 org.Mm.eg.db + organism = "mmu"；其余同理
# 运行：Rscript enrichment_plot.R
# 常见坑：
#   1) 基因 ID 类型必须说清（ENSEMBL / SYMBOL / ENTREZID），搞错会返回空结果
#   2) 网络不通时 enrichKEGG 会失败，可改 enrichGO（纯本地注释，不联网）
#   3) 富集结果没有条目数，通常是基因太少或 ID 没对上，先检查 bitr 的返回
###############################################################################