# ============================================================
# 04_enrichment_analysis.R
# GO/KEGG富集分析 + PPI网络
# ============================================================

library(clusterProfiler)
library(org.Hs.eg.db)
library(enrichplot)
library(STRINGdb)
library(igraph)
library(ggraph)
library(ggplot2)
library(patchwork)

# ---- 读取核心基因 ----
core_genes <- read.csv("output/core_genes_final.csv")$Gene

cat("核心基因:", paste(core_genes, collapse = ", "), "\n")

# ---- Gene Symbol 转 ENTREZ ID ----
gene_df <- bitr(core_genes, fromType = "SYMBOL", toType = "ENTREZID", OrgDb = org.Hs.eg.db)
cat("成功转换:", nrow(gene_df), "个基因\n")

# ---- GO富集分析 ----
# BP
ego_bp <- enrichGO(gene = gene_df$ENTREZID, OrgDb = org.Hs.eg.db, 
                   ont = "BP", pAdjustMethod = "BH", 
                   pvalueCutoff = 0.05, qvalueCutoff = 0.2, readable = TRUE)

# CC  
ego_cc <- enrichGO(gene = gene_df$ENTREZID, OrgDb = org.Hs.eg.db,
                   ont = "CC", pAdjustMethod = "BH",
                   pvalueCutoff = 0.05, qvalueCutoff = 0.2, readable = TRUE)

# MF
ego_mf <- enrichGO(gene = gene_df$ENTREZID, OrgDb = org.Hs.eg.db,
                   ont = "MF", pAdjustMethod = "BH",
                   pvalueCutoff = 0.05, qvalueCutoff = 0.2, readable = TRUE)

# 保存结果
if (!is.null(ego_bp) && nrow(as.data.frame(ego_bp)) > 0) {
  write.csv(as.data.frame(ego_bp), "output/GO_BP_results.csv", row.names = FALSE)
  dotplot(ego_bp, showCategory = 15, title = "GO BP")
  ggsave("output/12_GO_BP_dotplot.png", width = 10, height = 8)
}

if (!is.null(ego_cc) && nrow(as.data.frame(ego_cc)) > 0) {
  write.csv(as.data.frame(ego_cc), "output/GO_CC_results.csv", row.names = FALSE)
}

if (!is.null(ego_mf) && nrow(as.data.frame(ego_mf)) > 0) {
  write.csv(as.data.frame(ego_mf), "output/GO_MF_results.csv", row.names = FALSE)
}

# ---- KEGG富集分析 ----
kk <- enrichKEGG(gene = gene_df$ENTREZID, organism = 'hsa',
                 pAdjustMethod = "BH", pvalueCutoff = 0.05, qvalueCutoff = 0.2)

if (!is.null(kk) && nrow(as.data.frame(kk)) > 0) {
  write.csv(as.data.frame(kk), "output/KEGG_results.csv", row.names = FALSE)
  dotplot(kk, showCategory = 15, title = "KEGG Pathway")
  ggsave("output/16_KEGG_dotplot.png", width = 10, height = 8)
}

# ---- GO三合一图 ----
if (!is.null(ego_bp) && !is.null(ego_cc) && !is.null(ego_mf)) {
  p1 <- dotplot(ego_bp, showCategory = 10, title = "BP") + theme(legend.position = "none")
  p2 <- dotplot(ego_cc, showCategory = 10, title = "CC") + theme(legend.position = "none")
  p3 <- dotplot(ego_mf, showCategory = 10, title = "MF")
  p1 + p2 + p3 + plot_layout(ncol = 1)
  ggsave("output/15_GO_combined.png", width = 10, height = 14)
}

# ---- PPI网络分析（STRING） ----
cat("\n正在进行PPI分析...\n")

string_db <- STRINGdb$new(version = "11.5", species = 9606, score_threshold = 700)
mapped_genes <- string_db$map(data.frame(gene = core_genes), "gene", removeUnmappedRows = TRUE)

if (nrow(mapped_genes) > 1) {
  interactions <- string_db$get_interactions(mapped_genes$STRING_id)
  degree <- string_db$get_degree(mapped_genes$STRING_id)
  mapped_genes$degree <- degree[match(mapped_genes$STRING_id, names(degree))]
  
  hub_genes <- mapped_genes %>% arrange(desc(degree))
  write.csv(hub_genes, "output/PPI_hub_genes.csv", row.names = FALSE)
  cat("Hub基因:\n")
  print(head(hub_genes[, c("gene", "degree")], 10))
  
  # igraph可视化
  if (nrow(interactions) > 0) {
    edges <- data.frame(
      from = mapped_genes$gene[match(interactions$from, mapped_genes$STRING_id)],
      to = mapped_genes$gene[match(interactions$to, mapped_genes$STRING_id)],
      stringsAsFactors = FALSE
    )
    edges <- edges[!is.na(edges$from) & !is.na(edges$to), ]
    
    nodes <- data.frame(name = mapped_genes$gene, degree = mapped_genes$degree)
    
    g <- graph_from_data_frame(edges[, 1:2], vertices = nodes, directed = FALSE)
    
    ggraph(g, layout = "fr") +
      geom_edge_link(alpha = 0.3, color = "grey70") +
      geom_node_point(aes(size = degree, color = degree)) +
      geom_node_text(aes(label = name), repel = TRUE, size = 4) +
      scale_color_gradient(low = "lightblue", high = "red") +
      scale_size_continuous(range = c(3, 15)) +
      theme_graph() +
      labs(title = "PPI Network of Core Genes")
    
    ggsave("output/18_PPI_network.png", width = 10, height = 10)
  }
}

cat("\n富集分析完成！\n")