#!/usr/bin/env Rscript

suppressPackageStartupMessages({
  library(DESeq2)
  library(data.table)
})

args <- commandArgs(trailingOnly = TRUE)

if (length(args) < 5) {
  stop("Usage: Rscript DET_pairwise_clean.R <count.tsv> <design.tsv> <cpm_cutoff> <padj> <log2fc>")
}

count_file <- args[1]
design_file <- args[2]
cpm_cutoff <- as.numeric(args[3])
padj_cutoff <- as.numeric(args[4])
log2fc_cutoff <- as.numeric(args[5])

# -----------------------------
# LOAD DATA
# -----------------------------
counts <- fread(count_file)
design <- fread(design_file)

# count matrix
count_mat <- as.matrix(counts[, -(1:2)])
rownames(count_mat) <- counts$transcript_id

# -----------------------------
# ALIGN DESIGN TO COUNTS
# -----------------------------
samples <- colnames(count_mat)

design <- design[match(samples, name)]

if (any(is.na(design$name))) {
  stop("Sample mismatch between count matrix and design file")
}

design_df <- data.frame(condition = factor(design$condition))
rownames(design_df) <- design$name

cat("Samples order:\n")
print(rownames(design_df))

# -----------------------------
# CPM FILTER (like old script)
# -----------------------------
lib_size <- colSums(count_mat)
cpm <- t(t(count_mat) / lib_size * 1e6)

keep <- rowSums(cpm >= cpm_cutoff) >= 2
count_mat <- count_mat[keep, ]

cat("Retained after CPM filter:", nrow(count_mat), "\n")

# -----------------------------
# DESEQ2
# -----------------------------
dds <- DESeqDataSetFromMatrix(
  countData = round(count_mat),
  colData = design_df,
  design = ~ condition
)

dds <- DESeq(dds)

# -----------------------------
# DEFINE PAIRWISE COMPARISONS
# -----------------------------
comparisons <- list(
  c("condition", "Primed", "Naive"),
  c("condition", "TSC", "Naive"),
  c("condition", "TSC", "Primed")
)

names(comparisons) <- c(
  "Primed_vs_Naive",
  "TSC_vs_Naive",
  "TSC_vs_Primed"
)

# -----------------------------
# RUN COMPARISONS
# -----------------------------
for (nm in names(comparisons)) {

  cat("\nRunning:", nm, "\n")

  res <- results(dds, contrast = comparisons[[nm]])
  res_df <- as.data.frame(res)
  res_df$transcript_id <- rownames(res_df)

  # full output
  write.table(
    res_df,
    file = paste0(nm, "_DET_full_clean.tsv"),
    sep = "\t",
    quote = FALSE,
    row.names = FALSE
  )

  # filtered (significant)
  sig <- subset(
    res_df,
    !is.na(padj) &
      padj < padj_cutoff &
      abs(log2FoldChange) >= log2fc_cutoff
  )

  write.table(
    sig,
    file = paste0(nm, "_DET_significant_clean.tsv"),
    sep = "\t",
    quote = FALSE,
    row.names = FALSE
  )

  cat("Significant:", nrow(sig), "\n")
}

cat("\nAll comparisons done.\n")