#!/usr/bin/env Rscript

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

# =========================
# FILES
# =========================
det_files <- list(
  TSC_vs_Naive = "TSC_vs_Naive_DEG_full_list.csv",
  Primed_vs_Naive = "Primed_vs_Naive_DEG_full_list.csv",
  TSC_vs_Primed = "TSC_vs_Primed_DEG_full_list.csv"
)

te_files <- list(
  TSC_vs_Naive = "TSC_vs_Naive_DEG_TSS_TE.tsv",
  Primed_vs_Naive = "Primed_vs_Naive_DEG_TSS_TE.tsv",
  TSC_vs_Primed = "TSC_vs_Primed_DEG_TSS_TE.tsv"
)

comparison_order <- c("Primed_vs_Naive", "TSC_vs_Naive", "TSC_vs_Primed")
cap_value <- 5

# =========================
# LOAD FC
# =========================
load_fc <- function(file, name){
  dt <- fread(file)
  ids <- gsub('"', '', dt[[1]])
  fc <- as.numeric(dt$log2FoldChange)
  out <- data.table(transcript_id = ids, value = fc)
  setnames(out, "value", name)
  unique(out, by="transcript_id")
}

fc_list <- lapply(names(det_files), function(nm){
  load_fc(det_files[[nm]], nm)
})
names(fc_list) <- names(det_files)

fc_matrix <- Reduce(function(x,y) merge(x,y, by="transcript_id", all=TRUE), fc_list)
fc_matrix <- fc_matrix[, c("transcript_id", comparison_order), with=FALSE]

# =========================
# LOAD TE IDS
# =========================
te_ids <- lapply(te_files, function(f){
  dt <- fread(f, header=FALSE)
  unique(trimws(dt[[ncol(dt)]]))
})

# =========================
# LOAD TE ANNOTATION
# =========================
extract_one <- function(pattern, x) {
  m <- regexec(pattern, x)
  r <- regmatches(x, m)
  sapply(r, function(z) if (length(z) > 1) z[2] else NA_character_)
}

all_lines <- unlist(lapply(te_files, readLines))

te_annot <- data.table(
  transcript_id = extract_one('transcript_id "([^"]+)"', all_lines),
  class = extract_one('repeat_class "([^"]+)"', all_lines),
  family = extract_one('repeat_family "([^"]+)"', all_lines)
)

te_annot <- unique(te_annot, by="transcript_id")
te_annot[is.na(class), class := "Unknown"]
te_annot[is.na(family), family := "Unknown"]

# =========================
# MATRIX BUILDER
# =========================
make_matrix <- function(ids){

  mat <- merge(data.table(transcript_id=ids), fc_matrix, by="transcript_id", all.x=TRUE)

  m <- as.matrix(mat[, ..comparison_order])
  rownames(m) <- mat$transcript_id

  m[is.na(m)] <- 0

  keep <- apply(m,1,function(x) all(is.finite(x)))
  m <- m[keep,,drop=FALSE]

  vars <- apply(m,1,var)
  m <- m[vars>0,,drop=FALSE]

  m[m>cap_value] <- cap_value
  m[m< -cap_value] <- -cap_value

  return(m)
}

# =========================
# HEATMAP FUNCTION (FIXED)
# =========================
plot_heatmap <- function(m, title, outfile){

  if (nrow(m) < 2) return(NULL)

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

  pdf(outfile, width=7, height=10)
  pheatmap(
    m,
    scale="none",
    cluster_cols=TRUE,
    annotation_row = row_annot,   # ? NO manual color mapping
    show_rownames=FALSE,
    main=title
  )
  dev.off()
}

# =========================
# A: INTERSECTION
# =========================
intersection_ids <- Reduce(intersect, te_ids)
cat("Intersection:", length(intersection_ids), "\n")

mA <- make_matrix(intersection_ids)
cat("Intersection matrix:", dim(mA), "\n")

plot_heatmap(mA, "Intersection TE-DEG", "Heatmap_A.pdf")

# =========================
# C: SEPARATE
# =========================
for (nm in names(te_ids)){
  m <- make_matrix(te_ids[[nm]])
  cat(nm, "matrix:", dim(m), "\n")

  plot_heatmap(
    m,
    paste("Selected by", nm),
    paste0("Heatmap_", nm, ".pdf")
  )
}

cat("DONE\n")