#!/usr/bin/env Rscript

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

base_dir <- "/BLUES/eric/ONT_WGBS/ChIP-seq/RSEM/DET_isoform"

tpm_file <- file.path(base_dir, "transcript_TPM.tsv")

te_files <- list(
  TSC_vs_Naive   = file.path(base_dir, "TSC_vs_Naive/DET_TE_full_TSC_vs_Naive.tsv"),
  Primed_vs_Naive = file.path(base_dir, "Primed_vs_Naive/DET_TE_full_Primed_vs_Naive.tsv"),
  TSC_vs_Primed  = file.path(base_dir, "TSC_vs_Primed/DET_TE_full_TSC_vs_Primed.tsv")
)

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

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
log2_expr <- log2(expr_mat + 1)

extract_field <- function(pattern, x) {
  m <- regexec(pattern, x)
  sapply(regmatches(x, m), function(z) if (length(z) > 1) z[2] else NA_character_)
}

load_te_table <- function(file) {
  lines <- readLines(file, warn = FALSE)
  dt <- fread(file, header = FALSE)
  out <- data.table(
    transcript_id = trimws(as.character(dt[[ncol(dt) - 1]])),
    class         = extract_field('repeat_class "([^"]+)"', lines)
  )
  out[is.na(class) | class == "", class := "Unknown"]
  unique(out, by = "transcript_id")
}

te_tables <- lapply(te_files, load_te_table)
te_ids    <- lapply(te_tables, function(x) unique(x$transcript_id))
te_annot  <- unique(rbindlist(te_tables, fill = TRUE), by = "transcript_id")

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")
)

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

make_heatmap <- function(ids, name, out_prefix) {

  ids_in_expr <- intersect(unique(ids), rownames(log2_expr))
  if (length(ids_in_expr) < 2) { cat("Skipping", name, "- too few transcripts\n"); return(NULL) }

  mat <- log2_expr[ids_in_expr, , drop = FALSE]
  mat <- mat[apply(mat, 1, function(x) all(is.finite(x))), , drop = FALSE]
  mat <- mat[apply(mat, 1, var) > 1e-6, , drop = FALSE]
  if (nrow(mat) < 2) { cat("Skipping", name, "- no variance\n"); return(NULL) }

  n_rows <- nrow(mat)
  subtitle <- paste0("n = ", n_rows, " | padj < 0.1, |log2FC| > 1")

  row_annot <- data.frame(
    class = te_annot[match(rownames(mat), transcript_id)]$class,
    row.names = rownames(mat)
  )

  # raw log2(TPM+1) heatmap
  pdf(paste0(out_prefix, "_TPM_heatmap.pdf"), width = 10, height = 12)
  p <- pheatmap(mat, cluster_cols = FALSE, show_rownames = FALSE,
                annotation_col = col_annot, annotation_row = row_annot,
                annotation_colors = ann_colors, silent = TRUE)
  grid.arrange(
    textGrob(name, gp = gpar(fontsize = 14, fontface = "bold")),
    textGrob(paste0(subtitle, " | log2(TPM + 1)"), gp = gpar(fontsize = 10)),
    p$gtable, ncol = 1, heights = c(0.05, 0.05, 0.90)
  )
  dev.off()

  tpm_order <- rownames(mat)[p$tree_row$order]
  tpm_out <- as.data.table(mat[tpm_order, ], keep.rownames = "transcript_id")
  tpm_out[, te_class := te_annot[match(tpm_out$transcript_id, te_annot$transcript_id)]$class]
  fwrite(tpm_out, paste0(out_prefix, "_TPM_matrix.tsv"), sep = "\t")

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

  row_annot_z <- data.frame(
    class = te_annot[match(rownames(mat_z), transcript_id)]$class,
    row.names = rownames(mat_z)
  )

  pdf(paste0(out_prefix, "_Zscore_heatmap.pdf"), width = 10, height = 12)
  p_z <- pheatmap(mat_z, cluster_cols = FALSE, show_rownames = FALSE,
                  annotation_col = col_annot, annotation_row = row_annot_z,
                  annotation_colors = ann_colors,
                  color = colors, breaks = breaks, silent = TRUE)
  grid.arrange(
    textGrob(name, gp = gpar(fontsize = 14, fontface = "bold")),
    textGrob(paste0(subtitle, " | row Z-score, clamped [-2, 2]"), gp = gpar(fontsize = 10)),
    p_z$gtable, ncol = 1, heights = c(0.05, 0.05, 0.90)
  )
  dev.off()

  zscore_order <- rownames(mat_z)[p_z$tree_row$order]
  zscore_out <- as.data.table(mat_z[zscore_order, ], keep.rownames = "transcript_id")
  zscore_out[, te_class := te_annot[match(zscore_out$transcript_id, te_annot$transcript_id)]$class]
  fwrite(zscore_out, paste0(out_prefix, "_Zscore_matrix.tsv"), sep = "\t")

  cat("Done:", name, "| n =", n_rows, "\n")
}

for (nm in names(te_ids)) {
  out_dir <- file.path(base_dir, nm)
  make_heatmap(
    ids        = te_ids[[nm]],
    name       = paste0("DET TE heatmap - ", gsub("_vs_", " vs ", nm)),
    out_prefix = file.path(out_dir, nm)
  )
}

cat("All done.\n")