#!/usr/bin/env Rscript

suppressPackageStartupMessages({
  library(data.table)
  library(pheatmap)
  library(gridExtra)
  library(grid)
  library(ggplot2)
  library(cluster)
})

base_dir    <- "/BLUES/eric/ONT_WGBS/ChIP-seq/RSEM/DET_isoform"
tpm_file    <- file.path(base_dir, "transcript_TPM.tsv")
out_dir     <- file.path(base_dir, "kmeans_PvsN_NvsT")
dir.create(out_dir, showWarnings = FALSE)

det_files <- list(
  Primed_vs_Naive = file.path(base_dir, "Primed_vs_Naive/Naive_vs_Primed_DET_significant_only.csv"),
  TSC_vs_Naive    = file.path(base_dir, "TSC_vs_Naive/TSC_vs_Naive_DET_significant_only.csv")
)

sample_order <- c(
  "Primed_Rep1", "Primed_Rep2",
  "Naive_Rep1",  "Naive_Rep2",
  "GFP_TSC_Rep1", "GFP_TSC_Rep2"
)

k         <- 3   # k=3 chosen based on elbow plot
set.seed(42)

# ── load and merge DET IDs ──────────────────────────────────────────────────
det_ids <- lapply(det_files, function(f) {
  d <- fread(f)
  trimws(as.character(d[[1]]))
})
all_ids <- unique(unlist(det_ids))
cat("Transcripts in Primed_vs_Naive:", length(det_ids[[1]]), "\n")
cat("Transcripts in TSC_vs_Naive:   ", length(det_ids[[2]]), "\n")
cat("Union (total unique):          ", length(all_ids), "\n")

# ── load TPM, filter to DET union ───────────────────────────────────────────
expr <- fread(tpm_file)
expr[, transcript_id := trimws(as.character(transcript_id))]
expr_mat <- as.matrix(expr[, ..sample_order])
mode(expr_mat) <- "numeric"
rownames(expr_mat) <- expr$transcript_id

mat <- log2(expr_mat[intersect(all_ids, rownames(expr_mat)), ] + 1)
mat <- mat[apply(mat, 1, function(x) all(is.finite(x))), ]
mat <- mat[apply(mat, 1, var) > 1e-6, ]
cat("Transcripts after filtering:   ", nrow(mat), "\n")

# ── row Z-score ──────────────────────────────────────────────────────────────
mat_z <- t(scale(t(mat)))
mat_z <- mat_z[complete.cases(mat_z), ]
mat_z[mat_z >  2] <-  2
mat_z[mat_z < -2] <- -2

# ── elbow plot to help choose k ──────────────────────────────────────────────
cat("Computing elbow plot (k=2..12) ...\n")
wss <- sapply(2:12, function(k) {
  kmeans(mat_z, centers = k, nstart = 25, iter.max = 100)$tot.withinss
})
elbow_df <- data.frame(k = 2:12, wss = wss)
pdf(file.path(out_dir, "elbow_plot.pdf"), width = 6, height = 4)
print(ggplot(elbow_df, aes(x = k, y = wss)) +
  geom_line() + geom_point(size = 2) +
  scale_x_continuous(breaks = 2:12) +
  labs(title = "K-means elbow plot", x = "Number of clusters (k)",
       y = "Total within-cluster SS") +
  theme_classic())
dev.off()
cat("Elbow plot saved.\n")

# ── k-means ──────────────────────────────────────────────────────────────────
cat("Running k-means with k =", k, "...\n")
km <- kmeans(mat_z, centers = k, nstart = 25, iter.max = 100)
clusters <- km$cluster

# ── source annotation ────────────────────────────────────────────────────────
source_annot <- data.frame(
  transcript_id = all_ids,
  in_PvsN       = all_ids %in% det_ids[["Primed_vs_Naive"]],
  in_TvsN       = all_ids %in% det_ids[["TSC_vs_Naive"]]
)
source_annot$source <- with(source_annot, ifelse(in_PvsN & in_TvsN, "Both",
                              ifelse(in_PvsN, "Primed_vs_Naive", "TSC_vs_Naive")))

# ── heatmap ──────────────────────────────────────────────────────────────────
row_order <- order(clusters)
mat_plot  <- mat_z[row_order, ]

row_annot <- data.frame(
  Cluster = factor(clusters[row_order]),
  Source  = source_annot$source[match(names(clusters[row_order]), source_annot$transcript_id)],
  row.names = names(clusters[row_order])
)

col_annot <- data.frame(
  condition = factor(c("Primed","Primed","Naive","Naive","TSC","TSC"),
                     levels = c("Primed","Naive","TSC")),
  row.names = sample_order
)

cluster_colors <- setNames(
  colorRampPalette(c("#E41A1C","#377EB8","#4DAF4A","#984EA3","#FF7F00","#A65628","#F781BF","#999999"))(k),
  as.character(1:k)
)

ann_colors <- list(
  condition = c(Primed = "#3C5488", Naive = "#00A087", TSC = "#E64B35"),
  Cluster   = cluster_colors,
  Source    = c(Both = "#444444", Primed_vs_Naive = "#3C5488", TSC_vs_Naive = "#E64B35")
)

breaks <- seq(-2, 2, length.out = 101)
colors <- colorRampPalette(c("#3B4CC0","white","#B40426"))(100)

pdf(file.path(out_dir, paste0("kmeans_k", k, "_heatmap.pdf")), width = 11, height = 14)
p <- pheatmap(mat_plot,
              cluster_rows   = FALSE,
              cluster_cols   = FALSE,
              show_rownames  = FALSE,
              annotation_col = col_annot,
              annotation_row = row_annot,
              annotation_colors = ann_colors,
              color  = colors,
              breaks = breaks,
              gaps_row = cumsum(table(clusters[row_order]))[1:(k-1)],
              silent = TRUE)
grid.arrange(
  textGrob(paste0("K-means clustering (k=", k, ") — merged Primed vs Naive + TSC vs Naive DET"),
           gp = gpar(fontsize = 13, fontface = "bold")),
  textGrob(paste0("n = ", nrow(mat_plot), " transcripts | row Z-score, clamped [-2,2]"),
           gp = gpar(fontsize = 10)),
  p$gtable, ncol = 1, heights = c(0.04, 0.03, 0.93)
)
dev.off()
cat("Heatmap saved.\n")

# ── export matrix TSV ────────────────────────────────────────────────────────
out_dt <- as.data.table(mat_plot, keep.rownames = "transcript_id")
out_dt[, cluster := clusters[match(transcript_id, names(clusters))]]
out_dt[, source  := source_annot$source[match(transcript_id, source_annot$transcript_id)]]
setcolorder(out_dt, c("transcript_id", "cluster", "source", sample_order))
fwrite(out_dt, file.path(out_dir, paste0("kmeans_k", k, "_matrix.tsv")), sep = "\t")

# ── cluster summary ──────────────────────────────────────────────────────────
summary_dt <- out_dt[, .N, by = cluster][order(cluster)]
cat("\nCluster sizes:\n")
print(summary_dt)
fwrite(summary_dt, file.path(out_dir, paste0("kmeans_k", k, "_cluster_sizes.tsv")), sep = "\t")

cat("\nAll done. Output in:", out_dir, "\n")