#!/usr/bin/env Rscript

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

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_TE_allcomparisons")
dir.create(out_dir, showWarnings = FALSE)

te_id_files <- list(
  TSC_vs_Naive    = file.path(base_dir, "TSC_vs_Naive/DET_TE_ids.txt"),
  Primed_vs_Naive = file.path(base_dir, "Primed_vs_Naive/DET_TE_ids.txt"),
  TSC_vs_Primed   = file.path(base_dir, "TSC_vs_Primed/DET_TE_ids.txt")
)

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

k <- 3
set.seed(42)

# ── load TE IDs - union of all 3 comparisons ─────────────────────────────────
te_ids <- lapply(te_id_files, function(f) trimws(readLines(f)))
all_te_ids <- unique(unlist(te_ids))
cat("TE transcripts per comparison:\n")
cat("  TSC vs Naive:   ", length(te_ids[["TSC_vs_Naive"]]), "\n")
cat("  Primed vs Naive:", length(te_ids[["Primed_vs_Naive"]]), "\n")
cat("  TSC vs Primed:  ", length(te_ids[["TSC_vs_Primed"]]), "\n")
cat("  Union (total):  ", length(all_te_ids), "\n")

# ── load TPM, subset to TE 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_te_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 ────────────────────────────────────────────────────────────────
cat("Computing elbow plot (k=2..12) ...\n")
wss <- sapply(2:12, function(ki) {
  kmeans(mat_z, centers = ki, 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 - TE transcripts (all comparisons)",
       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

# ── silhouette plot ───────────────────────────────────────────────────────────
cat("Computing silhouette scores ...\n")
sil     <- silhouette(clusters, dist(mat_z, method = "euclidean"))
avg_sil <- round(mean(sil[, 3]), 3)
cat("Average silhouette width:", avg_sil, "\n")

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

sil_df <- data.frame(
  transcript_id = names(clusters),
  cluster       = factor(clusters),
  sil_width     = sil[, 3]
)[order(clusters, -sil[, 3]), ]

pdf(file.path(out_dir, paste0("silhouette_k", k, ".pdf")), width = 8, height = 5)
print(ggplot(sil_df, aes(x = seq_along(sil_width), y = sil_width, fill = cluster)) +
  geom_col(width = 1) +
  scale_fill_manual(values = cluster_palette) +
  geom_hline(yintercept = avg_sil, linetype = "dashed", color = "black") +
  labs(title = paste0("Silhouette plot (k=", k, ") - TE transcripts"),
       subtitle = paste0("Average silhouette width: ", avg_sil),
       x = "Transcripts", y = "Silhouette width", fill = "Cluster") +
  theme_classic() +
  theme(axis.text.x = element_blank(), axis.ticks.x = element_blank()))
dev.off()
cat("Silhouette plot saved.\n")

# ── source annotation (which comparison each TE transcript came from) ─────────
source_annot <- data.frame(
  transcript_id = all_te_ids,
  in_TvsN       = all_te_ids %in% te_ids[["TSC_vs_Naive"]],
  in_PvsN       = all_te_ids %in% te_ids[["Primed_vs_Naive"]],
  in_TvsP       = all_te_ids %in% te_ids[["TSC_vs_Primed"]]
)
source_annot$source <- with(source_annot, case_when(
  in_TvsN & in_PvsN & in_TvsP ~ "All_3",
  in_TvsN & in_PvsN            ~ "TvsN+PvsN",
  in_TvsN & in_TvsP            ~ "TvsN+TvsP",
  in_PvsN & in_TvsP            ~ "PvsN+TvsP",
  in_TvsN                      ~ "TSC_vs_Naive",
  in_PvsN                      ~ "Primed_vs_Naive",
  in_TvsP                      ~ "TSC_vs_Primed"
))

# ── 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
)

ann_colors <- list(
  condition = c(Primed = "#3C5488", Naive = "#00A087", TSC = "#E64B35"),
  Cluster   = cluster_palette
)

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, ") - all changed TE transcripts"),
           gp = gpar(fontsize = 13, fontface = "bold")),
  textGrob(paste0("n = ", nrow(mat_plot), " | union of 3 comparisons | 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")