#!/usr/bin/env Rscript
suppressPackageStartupMessages({
  library(RUVSeq)
  library(edgeR)
  library(DESeq2)
  library(ggplot2)
  library(ggrepel)
  library(pheatmap)
})

# ============================================================
# 0. CONFIG
# ============================================================
outdir   <- "/BRC/yan/heart/analysis/DEG_LV/female_v2/"
infile   <- "/BRC/yan/heart/getx_LV/filtered_count_table/Heart_-_Left_Ventricle_sex2_filter.txt"
metafile <- "/BRC/yan/heart/analysis/meta_sex2.txt"
k_range  <- 1:6

fdr_cutoff <- 0.05
fc_cutoff  <- 0.58

# ============================================================
# 1. LOAD DATA
# ============================================================
cat("[1] Loading count matrix...\n")
counts_raw           <- read.table(infile, header = TRUE, sep = "\t", check.names = FALSE)
rownames(counts_raw) <- counts_raw$Name
counts_raw           <- counts_raw[, -(1:2)]
cat("    Genes:", nrow(counts_raw), "  Total samples:", ncol(counts_raw), "\n")

# ============================================================
# 2. METADATA & SUBSET
# ============================================================
cat("[2] Loading metadata...\n")
meta           <- read.table(metafile, header = TRUE, sep = "\t")
meta           <- meta[meta$sample_id %in% colnames(counts_raw), ]
counts         <- counts_raw[, meta$sample_id]
x              <- factor(meta$age_group, levels = c("young", "old"))
rownames(meta) <- meta$sample_id
cat("    young:", sum(x == "young"), "  old:", sum(x == "old"), "\n")

# ============================================================
# 3. SKIP FILTER (already filtered)
# ============================================================
cat("[3] Skipping filter step (matrix pre-filtered).\n")
y <- DGEList(counts = as.matrix(counts), group = x)

# ============================================================
# 4. PRE-BATCH PCA + HEATMAP
# ============================================================
cat("[4] Pre-batch PCA and heatmap...\n")
log_pre  <- log1p(as.matrix(counts))

# PCA
pca_pre <- prcomp(t(log_pre), scale. = FALSE)
pct_pre <- round(summary(pca_pre)$importance[2, 1:2] * 100, 1)
df_pre  <- data.frame(PC1       = pca_pre$x[, 1],
                      PC2       = pca_pre$x[, 2],
                      age_group = meta[rownames(pca_pre$x), "age_group"],
                      sample    = rownames(pca_pre$x))
p_pre <- ggplot(df_pre, aes(PC1, PC2, color = age_group, label = sample)) +
  geom_point(size = 3) +
  geom_text_repel(size = 2.5, max.overlaps = 20) +
  scale_color_manual(values = c(young = "#4DBBD5", old = "#E64B35")) +
  labs(title = "Sex2 LV PCA (pre-batch)",
       x = paste0("PC1 (", pct_pre[1], "%)"),
       y = paste0("PC2 (", pct_pre[2], "%)")) +
  theme_bw() +
  theme(plot.title = element_text(hjust = 0.5, face = "bold"))
ggsave(paste0(outdir, "sex2_v2_PCA_pre.pdf"), p_pre, width = 8, height = 6)

# Heatmap
ann_col <- data.frame(age_group = meta[colnames(counts), "age_group"],
                      row.names = colnames(counts))
pdf(paste0(outdir, "sex2_v2_heatmap_pre.pdf"), width = 9, height = 8)
pheatmap(cor(log_pre), annotation_col = ann_col, show_colnames = FALSE,
         main = "Sex2 LV Sample Correlation (pre-batch)")
dev.off()

# ============================================================
# 5. RUVr: edgeR model → deviance residuals
# ============================================================
cat("[5] Computing deviance residuals for RUVr (edgeR)...\n")
set    <- newSeqExpressionSet(as.matrix(y$counts),
           phenoData = data.frame(x, row.names = colnames(y)))
design <- model.matrix(~ x, data = pData(set))

y_tmp  <- calcNormFactors(y, method = "RLE")
y_tmp  <- estimateDisp(y_tmp, design)
fit    <- glmFit(y_tmp, design)
res    <- residuals(fit, type = "deviance")
seqUQ  <- betweenLaneNormalization(set, which = "upper")

# ============================================================
# 6. RUVr k=1~6: diagnostics + DESeq2 DEG per k
# ============================================================
cat("[6] Running RUVr k=1~6 with DESeq2 DEG each...\n")
colors          <- ifelse(meta[colnames(y), "age_group"] == "young", "#4DBBD5", "#E64B35")
results_summary <- data.frame()
deg_lists       <- list(all  = vector("list", length(k_range)),
                        up   = vector("list", length(k_range)),
                        down = vector("list", length(k_range)))

for (k in k_range) {
  cat("    k =", k, "\n")
  ruv    <- RUVr(seqUQ, rownames(set), k = k, res)
  norm_c <- normCounts(ruv)

  # RLE plot
  suppressMessages({
    pdf(paste0(outdir, "sex2_v2_RUVr_k", k, "_RLE.pdf"))
    plotRLE(ruv, outline = FALSE, ylim = c(-2, 2),
            col = colors, main = paste0("Sex2 LV RUVr k=", k))
    dev.off()
  })

  # Post-batch PCA
  log_post <- log1p(norm_c)
  pca_k    <- prcomp(t(log_post), scale. = FALSE)
  pct_k    <- round(summary(pca_k)$importance[2, 1:2] * 100, 1)
  df_k     <- data.frame(PC1       = pca_k$x[, 1],
                         PC2       = pca_k$x[, 2],
                         age_group = meta[rownames(pca_k$x), "age_group"],
                         sample    = rownames(pca_k$x))
  p_k <- ggplot(df_k, aes(PC1, PC2, color = age_group, label = sample)) +
    geom_point(size = 3) +
    geom_text_repel(size = 2.5, max.overlaps = 20) +
    scale_color_manual(values = c(young = "#4DBBD5", old = "#E64B35")) +
    labs(title = paste0("Sex2 LV RUVr k=", k),
         x = paste0("PC1 (", pct_k[1], "%)"),
         y = paste0("PC2 (", pct_k[2], "%)")) +
    theme_bw() +
    theme(plot.title = element_text(hjust = 0.5, face = "bold"))
  ggsave(paste0(outdir, "sex2_v2_RUVr_k", k, "_PCA.pdf"), p_k, width = 8, height = 6)

  # Post-batch heatmap
  ann_k <- data.frame(age_group = meta[colnames(norm_c), "age_group"],
                      row.names = colnames(norm_c))
  pdf(paste0(outdir, "sex2_v2_RUVr_k", k, "_heatmap.pdf"), width = 9, height = 8)
  pheatmap(cor(log_post), annotation_col = ann_k, show_colnames = FALSE,
           main = paste0("Sex2 LV Sample Correlation RUVr k=", k))
  dev.off()

  # DESeq2 DEG on RUVr-corrected counts
  coldata_k  <- data.frame(age_group = x, row.names = colnames(norm_c))
  design_mat <- model.matrix(~ 0 + age_group, data = coldata_k)
  colnames(design_mat) <- levels(x)   # "young", "old"

  dds <- DESeqDataSetFromMatrix(countData = round(norm_c),
                                 colData   = coldata_k,
                                 design    = design_mat)
  dds <- DESeq(dds)

  res_deg      <- results(dds, contrast = list("old", "young"))
  res_deg$padj <- ifelse(is.na(res_deg$padj), 1, res_deg$padj)
  deg          <- as.data.frame(res_deg)
  deg$gene     <- rownames(deg)
  deg$sig      <- "non-DEG"
  deg$sig[deg$padj < fdr_cutoff & deg$log2FoldChange >  fc_cutoff] <- "UP"
  deg$sig[deg$padj < fdr_cutoff & deg$log2FoldChange < -fc_cutoff] <- "DOWN"

  cat("      UP  (old > young):", sum(deg$sig == "UP"),   "\n")
  cat("      DOWN(old < young):", sum(deg$sig == "DOWN"), "\n")

  write.table(deg,
    paste0(outdir, "sex2_v2_DEG_old_vs_young_RUVr_k", k, ".txt"),
    sep = "\t", quote = FALSE, row.names = FALSE)
  write.table(deg[deg$sig != "non-DEG", ],
    paste0(outdir, "sex2_v2_DEG_old_vs_young_RUVr_k", k, "_filter.txt"),
    sep = "\t", quote = FALSE, row.names = FALSE)

  ki <- which(k_range == k)
  deg_f <- deg[deg$sig != "non-DEG", ]
  deg_lists$all[[ki]]  <- deg_f$gene
  deg_lists$up[[ki]]   <- deg_f$gene[deg_f$sig == "UP"]
  deg_lists$down[[ki]] <- deg_f$gene[deg_f$sig == "DOWN"]

  results_summary <- rbind(results_summary, data.frame(
    k           = k,
    total_genes = nrow(deg),
    UP          = sum(deg$sig == "UP"),
    DOWN        = sum(deg$sig == "DOWN")
  ))
}

write.csv(results_summary,
  paste0(outdir, "sex2_v2_DEG_old_vs_young_k_comparison_summary.csv"),
  row.names = FALSE)
cat("[6] k summary saved.\n")

# ============================================================
# 7. VOLCANO PLOT for each k
# ============================================================
cat("[7] Plotting volcano plots for each k...\n")

for (k in k_range) {
  deg <- read.table(
    paste0(outdir, "sex2_v2_DEG_old_vs_young_RUVr_k", k, ".txt"),
    header = TRUE, sep = "\t")

  p_vol <- ggplot(deg, aes(log2FoldChange, -log10(padj), color = sig)) +
    geom_point(size = 1, alpha = 0.6) +
    geom_vline(xintercept = c(-fc_cutoff, fc_cutoff), linetype = "dashed") +
    geom_hline(yintercept = -log10(fdr_cutoff), linetype = "dashed") +
    scale_color_manual(values = c(UP = "#E64B35", DOWN = "#4DBBD5", "non-DEG" = "grey70")) +
    labs(title = paste0("Sex2 LV | old vs young (RUVr k=", k, ")"),
         x = "log2FC (old / young)", y = "-log10(padj)") +
    theme_bw() +
    theme(plot.title = element_text(hjust = 0.5, face = "bold"))
  ggsave(paste0(outdir, "sex2_v2_DEG_old_vs_young_RUVr_k", k, "_volcano.pdf"),
         p_vol, width = 7, height = 6)
}

# ============================================================
# 8. DEG OVERLAP ACROSS k VALUES
# ============================================================
cat("[8] Computing DEG overlap across k values...\n")
k_labels <- paste0("k", k_range)
n_k      <- length(k_range)

make_overlap_mat <- function(lst, labels) {
  mat <- matrix(0L, n_k, n_k, dimnames = list(labels, labels))
  for (i in 1:n_k) for (j in 1:n_k)
    mat[i, j] <- length(intersect(lst[[i]], lst[[j]]))
  mat
}

ov_all  <- make_overlap_mat(deg_lists$all,  k_labels)
ov_up   <- make_overlap_mat(deg_lists$up,   k_labels)
ov_down <- make_overlap_mat(deg_lists$down, k_labels)

write.csv(ov_all,  paste0(outdir, "sex2_v2_DEG_overlap_all_k.csv"))
write.csv(ov_up,   paste0(outdir, "sex2_v2_DEG_overlap_UP_k.csv"))
write.csv(ov_down, paste0(outdir, "sex2_v2_DEG_overlap_DOWN_k.csv"))

all_genes <- unique(unlist(deg_lists$all))
gene_tbl  <- as.data.frame(
  matrix(0L, length(all_genes), n_k, dimnames = list(all_genes, k_labels)))
for (i in 1:n_k) if (length(deg_lists$all[[i]]) > 0)
  gene_tbl[deg_lists$all[[i]], i] <- 1L
gene_tbl$n_k <- rowSums(gene_tbl[, k_labels])
gene_tbl$direction <- sapply(all_genes, function(g) {
  is_up   <- any(sapply(deg_lists$up,   function(v) g %in% v))
  is_down <- any(sapply(deg_lists$down, function(v) g %in% v))
  if (is_up & is_down) "mixed" else if (is_up) "UP" else "DOWN"
})
gene_tbl <- gene_tbl[order(-gene_tbl$n_k), ]
write.csv(gene_tbl, paste0(outdir, "sex2_v2_DEG_gene_presence_across_k.csv"))

cat("[Done]\n")