###############################################################################
# main_OA_pipeline.R
# Multi-Omics Pipeline for Osteoarthritis (OA)
###############################################################################

rm(list = ls())
options(stringsAsFactors = FALSE)

suppressPackageStartupMessages({
  library(dplyr); library(readxl); library(ggplot2); library(limma)
  library(WGCNA); library(sva); library(clusterProfiler); library(org.Hs.eg.db)
  library(glmnet); library(pROC); library(PRROC); library(pheatmap)
})

dir.create("results", showWarnings = FALSE)

###############################################################################
# Utility
###############################################################################

volcano_plot <- function(df, fc="logFC", p="adj.P.Val", out){
  df$status <- "no"
  df$status[df[[p]]<0.05 & df[[fc]]>0.5]  <- "up"
  df$status[df[[p]]<0.05 & df[[fc]]<(-0.5)] <- "down"
  p1 <- ggplot(df, aes(.data[[fc]], -log10(.data[[p]]), color=status)) +
    geom_point() + theme_bw() +
    scale_color_manual(values=c("grey","#0072B5","#BC3C28")) +
    geom_hline(yintercept=-log10(0.05), lty=2) +
    geom_vline(xintercept=c(-0.5,0.5), lty=2)
  ggsave(out, p1, width=6, height=6)
}

simple_heatmap <- function(expr, genes, out){
  mat <- expr[genes, , drop=FALSE]
  mat <- t(scale(t(mat)))
  pheatmap(mat, show_rownames=FALSE, show_colnames=FALSE,
           filename=out, color=colorRampPalette(rev(RColorBrewer::brewer.pal(9,"RdYlBu")))(100))
}

###############################################################################
# PART 1 — Batch Correction + PCA
###############################################################################

batch_correct_pca <- function(expr1, expr2, info1, info2){
  common <- intersect(rownames(expr1), rownames(expr2))
  expr <- cbind(expr1[common,], expr2[common,])
  info <- rbind(info1, info2); info <- info[match(colnames(expr), info$id),]
  
  p_before <- prcomp(t(expr), scale.=TRUE)
  
  mod <- model.matrix(~ info$group)
  expr_bc <- ComBat(expr, batch=factor(info$dataset), mod=mod)
  
  p_after <- prcomp(t(expr_bc), scale.=TRUE)
  
  write.csv(expr_bc, "results/Expression_Integrated.csv", row.names=TRUE)
  list(expr=expr_bc, pca_before=p_before, pca_after=p_after)
}

###############################################################################
# PART 2 — Sensitivity Analysis (Compressed)
###############################################################################

sensitivity_analysis <- function(expr, traits, gene_sets, target_genes){
  design <- model.matrix(~traits)
  fit <- eBayes(lmFit(expr, design))
  deg <- rownames(topTable(fit, coef=2, number=Inf, p.value=0.05))
  mat <- sapply(gene_sets, function(gs) target_genes %in% intersect(gs, deg))
  pheatmap(mat*1, filename="results/Sensitivity_Heatmap.pdf")
}

###############################################################################
# PART 3 — DEG (mRNA / miRNA / circRNA)
###############################################################################

run_deg <- function(expr_file, group_file, prefix){
  expr <- read.table(expr_file, header=TRUE, row.names=1)
  grp  <- read.table(group_file, header=TRUE)
  expr <- log2(expr+1)
  
  design <- model.matrix(~ factor(grp$Group, levels=c("normal","OA")))
  fit <- eBayes(lmFit(expr, design))
  res <- topTable(fit, coef=2, n=Inf)
  write.csv(res, paste0("results/", prefix, "_DEG_full.csv"))
  
  volcano_plot(res, out=paste0("results/", prefix, "_volcano.pdf"))
  simple_heatmap(expr, rownames(res)[1:30], paste0("results/", prefix, "_heatmap.pdf"))
  
  return(res)
}

###############################################################################
# PART 4 — WGCNA (Compressed Version)
###############################################################################

run_wgcna <- function(expr_file, trait_file){
  expr <- read.table(expr_file, header=TRUE, row.names=1)
  expr <- t(expr)
  traits <- read.table(trait_file, header=TRUE, row.names=1)
  
  expr <- expr[, apply(expr,2,var) > 0]
  
  sft <- pickSoftThreshold(expr, powerVector=1:20, verbose=0)
  power <- sft$powerEstimate %||% 6
  
  net <- blockwiseModules(expr, power=power, minModuleSize=30,
                          mergeCutHeight=0.25, numericLabels=TRUE,
                          verbose=0)
  
  MEs <- moduleEigengenes(expr, net$colors)$eigengenes
  corM <- cor(MEs, traits, use="p")
  
  pheatmap(corM, filename="results/WGCNA_Module_Trait.pdf")
  
  return(list(net=net, MEs=MEs))
}

###############################################################################
# PART 5 — Enrichment (GO + KEGG)
###############################################################################

run_enrich <- function(genes, prefix){
  eg <- bitr(genes, "SYMBOL","ENTREZID", OrgDb=org.Hs.eg.db)$ENTREZID
  ego <- enrichGO(eg, OrgDb=org.Hs.eg.db, ont="BP")
  ekegg <- enrichKEGG(eg)
  
  pdf(paste0("results/", prefix, "_GO.pdf")); barplot(ego); dev.off()
  pdf(paste0("results/", prefix, "_KEGG.pdf")); dotplot(ekegg); dev.off()
}

###############################################################################
# PART 6 — Spearman Correlation (Gene–Metabolite)
###############################################################################

run_corr <- function(df, out="results/Correlation_Heatmap.pdf"){
  M <- cor(df, method="spearman")
  pdf(out); corrplot::corrplot(M, method="color"); dev.off()
}

###############################################################################
# PART 7 — LASSO Model
###############################################################################

build_lasso <- function(expr, y){
  X <- scale(expr)
  fit <- cv.glmnet(X, y, family="binomial", alpha=1)
  save(fit, file="results/LASSO_model.RData")
  return(fit)
}

###############################################################################
# PART 8 — LASSO Validation
###############################################################################

validate_lasso <- function(expr, y, model){
  probs <- predict(model, newx=scale(expr), s="lambda.min", type="response")
  roc_obj <- roc(y, probs)
  pdf("results/LASSO_ROC.pdf"); plot(roc_obj); dev.off()
}

###############################################################################
# PART 9 — Multivariable Regression
###############################################################################

multi_regression <- function(df, genes){
  df_OA <- df[df$Group=="OA",]
  mets <- setdiff(names(df), c("SampleID","Group", genes))
  results <- list()
  for(m in mets){
    vars <- genes[sapply(genes, function(g) cor(df_OA[[m]], df_OA[[g]], method="spearman") |> is.finite)]
    if(length(vars)>=2){
      form <- as.formula(paste(m, "~", paste(vars, collapse="+")))
      results[[m]] <- summary(lm(form, data=df_OA))
    }
  }
  save(results, file="results/Multivariate_Regression.RData")
}

###############################################################################
# PART 10 — SCI Plots
###############################################################################

plot_lollipop <- function(stats){
  p <- ggplot(stats, aes(Pair, Rho, color=Group)) +
    geom_segment(aes(xend=Pair, y=0, yend=Rho)) +
    geom_point(size=4) +
    theme_bw()
  ggsave("results/Lollipop.pdf", p, width=6, height=5)
}

###############################################################################
cat("Pipeline loaded. Fill in your data paths and run modules as needed.\n")