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

# ============================================================
# 0. CONFIG
# ============================================================
outdir   <- "/BRC/yan/heart/analysis/DEG_LV/male/"
infile   <- "/BRC/yan/heart/getx_LV/filtered_count_table/Heart_-_Left_Ventricle_sex1_filter.txt"
metafile <- "/BRC/yan/heart/analysis/meta_sex1.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)]   # drop Name, Description
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
# ============================================================
cat("[4] Pre-batch PCA...\n")
log_mat_pre <- log1p(t(as.matrix(counts)))
pca_pre     <- prcomp(log_mat_pre, scale. = FALSE)
pct_pre     <- round(summary(pca_pre)$importance[2, 1:2] * 100, 1)
df_pca_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_pca_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 = "Sex1 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, "sex1_PCA_pre.pdf"), p_pre, width = 8, height = 6)

# ============================================================
# 5. RUVr: compute residuals
# ============================================================
cat("[5] Computing deviance residuals for RUVr...\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 + full DEG per k
# ============================================================
cat("[6] Running RUVr k=1~6 with full 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")
  set2  <- RUVr(seqUQ, rownames(set), k = k, res)
  ndddd <- normCounts(set2)

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

  # PCA
  log_mat <- log1p(t(ndddd))
  pca     <- prcomp(log_mat, scale. = FALSE)
  pct     <- round(summary(pca)$importance[2, 1:2] * 100, 1)
  df_pca  <- data.frame(
    PC1       = pca$x[, 1],
    PC2       = pca$x[, 2],
    age_group = meta[rownames(pca$x), "age_group"],
    sample    = rownames(pca$x)
  )
  p <- ggplot(df_pca, 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("Sex1 LV RUVr k=", k),
         x = paste0("PC1 (", pct[1], "%)"),
         y = paste0("PC2 (", pct[2], "%)")) +
    theme_bw() +
    theme(plot.title = element_text(hjust = 0.5, face = "bold"))
  ggsave(paste0(outdir, "sex1_RUVr_k", k, "_PCA.pdf"), p, width = 8, height = 6)

  # DEG: exactTest (old vs young)
  dgList_ruv <- DGEList(counts = ndddd, genes = rownames(ndddd), group = x)
  design_mat <- model.matrix(~ 0 + y$samples$group)
  colnames(design_mat) <- levels(y$samples$group)
  dgList_ruv <- estimateCommonDisp(dgList_ruv, design = design_mat)
  dgList_ruv <- estimateTagwiseDisp(dgList_ruv)

  et  <- exactTest(dgList_ruv, pair = c("young", "old"))
  deg <- topTags(et, n = nrow(et))$table
  deg$gene <- rownames(deg)
  deg$sig  <- "non-DEG"
  deg$sig[deg$FDR < fdr_cutoff & deg$logFC >  fc_cutoff] <- "UP"
  deg$sig[deg$FDR < fdr_cutoff & deg$logFC < -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, "sex1_DEG_old_vs_young_RUVr_k", k, ".txt"),
    sep = "\t", quote = FALSE, row.names = FALSE)
  write.table(deg[deg$sig != "non-DEG", ],
    paste0(outdir, "sex1_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, "sex1_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, "sex1_DEG_old_vs_young_RUVr_k", k, ".txt"),
    header = TRUE, sep = "\t")

  p_vol <- ggplot(deg, aes(logFC, -log10(FDR), 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("Sex1 LV | old vs young (RUVr k=", k, ")"),
         x = "logFC (old - young)", y = "-log10(FDR)") +
    theme_bw() +
    theme(plot.title = element_text(hjust = 0.5, face = "bold"))
  ggsave(paste0(outdir, "sex1_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, "sex1_DEG_overlap_all_k.csv"))
write.csv(ov_up,   paste0(outdir, "sex1_DEG_overlap_UP_k.csv"))
write.csv(ov_down, paste0(outdir, "sex1_DEG_overlap_DOWN_k.csv"))

# Per-gene presence table: rows = genes, cols = k1..k6 + direction + n_k
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, "sex1_DEG_gene_presence_across_k.csv"))

cat("[Done]\n")