#!/usr/bin/env Rscript

suppressPackageStartupMessages({
  library(data.table)
  library(ggplot2)
  library(scales)
  library(patchwork)
})

setDTthreads(8)

# =========================================================
# SETTINGS
# =========================================================

outdir <- "/BLUES/eric/ONT_WGBS/Figure_3/Heatmap"
dir.create(outdir, showWarnings = FALSE, recursive = TRUE)

infile <- "/BLUES/eric/ONT_WGBS/Heatmap/DNA_methylation_matrix_FULL_no_rRNA_tRNA.tsv"
te_bed <- "/BLUES/eric/ONT/hg38_TE_noY.bed.gz"

out_pdf_beta <- file.path(
  outdir,
  "TE_DNA_methylation_heatmap_beta_trueRepeatClass_compact.pdf"
)

out_pdf_zscore <- file.path(
  outdir,
  "TE_DNA_methylation_heatmap_zscore_trueRepeatClass_compact.pdf"
)

out_class_summary <- file.path(
  outdir,
  "TE_DNA_methylation_heatmap_trueRepeatClass_summary.tsv"
)

col_order <- c("Primed", "Naive", "TSC")

# =========================================================
# LOAD TE ANNOTATION FROM BED
# =========================================================

load_te_annotation <- function(te_bed) {

  te <- fread(cmd = paste("zcat", shQuote(te_bed)), header = FALSE)

  # Expected BED-like format:
  # V1 chr
  # V2 start
  # V3 end
  # V4 strand
  # V5 attributes:
  # repeat_id "L1MC5a,chr1,11485,11676";repeat_class "LINE";repeat_family "L1";

  te[, repeat_id := sub('.*repeat_id "([^"]+)".*', '\\1', V5)]
  te[, repeat_class := sub('.*repeat_class "([^"]+)".*', '\\1', V5)]
  te[, repeat_family := sub('.*repeat_family "([^"]+)".*', '\\1', V5)]

  # Subfamily is the part before the first comma in repeat_id
  te[, subfamily := tstrsplit(repeat_id, ",", fixed = TRUE, keep = 1L)]

  anno <- unique(te[, .(subfamily, repeat_class, repeat_family)])

  # Check whether any subfamily maps to more than one class/family
  ambig <- anno[, .(
    n_class = uniqueN(repeat_class),
    n_family = uniqueN(repeat_family)
  ), by = subfamily][n_class > 1 | n_family > 1]

  if (nrow(ambig) > 0) {
    warning(sprintf(
      "%d subfamilies map to more than one repeat_class/repeat_family. Keeping first sorted annotation per subfamily.",
      nrow(ambig)
    ))

    fwrite(
      ambig,
      file.path(dirname(out_class_summary), "ambiguous_subfamily_repeatClass_repeatFamily.tsv"),
      sep = "\t"
    )
  }

  setorder(anno, subfamily, repeat_class, repeat_family)
  anno <- anno[, .SD[1], by = subfamily]

  return(anno)
}

anno <- load_te_annotation(te_bed)

# =========================================================
# LOAD METHYLATION MATRIX
# =========================================================

dt <- fread(infile)

# Remove known non-TE rows if present
dt <- dt[!grepl("^5S$|^7SK$|^7SLRNA$", subfamily)]

# Keep only complete rows for samples used in heatmap
dt <- dt[complete.cases(dt[, ..col_order])]

# Join true RepeatMasker annotation
dt <- merge(dt, anno, by = "subfamily", all.x = TRUE)

# =========================================================
# USE TRUE REPEATMASKER CLASS
# =========================================================

dt[, plot_class := repeat_class]
dt[is.na(plot_class) | plot_class == "", plot_class := "Unannotated"]

# Preferred display order.
# Any additional repeat_class values present in your BED will be appended alphabetically.
preferred_class_order <- c(
  "LINE",
  "SINE",
  "LTR",
  "Retroposon",
  "DNA",
  "RC",
  "Satellite",
  "Simple_repeat",
  "Low_complexity",
  "RNA",
  "Unknown",
  "Unannotated"
)

observed_classes <- unique(dt$plot_class)

class_order <- c(
  preferred_class_order[preferred_class_order %in% observed_classes],
  sort(setdiff(observed_classes, preferred_class_order))
)

dt[, plot_class := factor(plot_class, levels = class_order)]

# =========================================================
# CLASS COLORS
# =========================================================

fixed_class_colors <- c(
  "LINE"           = "#d73027",
  "SINE"           = "#4575b4",
  "LTR"            = "#1a9850",
  "Retroposon"     = "#f46d43",
  "DNA"            = "#984ea3",
  "RC"             = "#66c2a5",
  "Satellite"      = "#a65628",
  "Simple_repeat"  = "#999999",
  "Low_complexity" = "#bdbdbd",
  "RNA"            = "#e78ac3",
  "Unknown"        = "#666666",
  "Unannotated"    = "#000000"
)

missing_color_classes <- setdiff(class_order, names(fixed_class_colors))

if (length(missing_color_classes) > 0) {
  extra_colors <- setNames(
    hue_pal()(length(missing_color_classes)),
    missing_color_classes
  )
} else {
  extra_colors <- character(0)
}

class_colors <- c(fixed_class_colors, extra_colors)
class_colors <- class_colors[class_order]

# =========================================================
# OUTPUT CLASS SUMMARY / SANITY CHECKS
# =========================================================

class_summary <- dt[, .N, by = .(plot_class)][order(plot_class)]
fwrite(class_summary, out_class_summary, sep = "\t")

cat("\nRepeat classes represented in heatmap:\n")
print(class_summary)

cat("\nExample SVA rows, if present:\n")
print(
  dt[grepl("^SVA", subfamily), .(
    subfamily,
    repeat_class,
    repeat_family,
    Primed,
    Naive,
    TSC
  )][order(subfamily)]
)

# =========================================================
# ORDER BY TRUE REPEAT CLASS
# =========================================================

setorder(dt, plot_class, subfamily)

# =========================================================
# MATRICES
# =========================================================

mat_beta <- as.matrix(dt[, ..col_order])

mat_z <- t(scale(t(mat_beta)))
mat_z[is.na(mat_z)] <- 0

rownames(mat_beta) <- dt$subfamily
rownames(mat_z) <- dt$subfamily

# =========================================================
# HEATMAP FUNCTION
# =========================================================

make_heatmap <- function(mat, mode = "beta", class_vec) {

  # -------------------------
  # MAIN HEATMAP DATA
  # -------------------------

  df <- as.data.table(as.table(mat))
  setnames(df, c("subfamily", "sample", "value"))

  df[, sample := factor(sample, levels = col_order)]
  df[, subfamily := factor(subfamily, levels = rev(rownames(mat)))]

  # -------------------------
  # CLASS STRIP DATA
  # -------------------------

  class_df <- data.table(
    subfamily = factor(rev(rownames(mat)), levels = rev(rownames(mat))),
    plot_class = factor(rev(as.character(class_vec)), levels = class_order)
  )

  # -------------------------
  # HEATMAP COLORS
  # -------------------------

  if (mode == "beta") {

    heat <- ggplot(df, aes(sample, subfamily, fill = value)) +
      geom_tile() +
      scale_fill_gradientn(
        colors = c("white", "#fff7bc", "#d7301f"),
        values = rescale(c(0, 0.5, 1)),
        limits = c(0, 1),
        name = "Beta"
      )

    plot_title <- "TE DNA methylation heatmap (beta values)"

  } else {

    heat <- ggplot(df, aes(sample, subfamily, fill = value)) +
      geom_tile() +
      scale_fill_gradient2(
        low = "#4575b4",
        mid = "white",
        high = "#d7301f",
        midpoint = 0,
        limits = c(-2.5, 2.5),
        name = "Z-score"
      )

    plot_title <- "TE DNA methylation heatmap (row z-score)"
  }

  # -------------------------
  # MAIN HEATMAP STYLE
  # -------------------------

  heat <- heat +
    labs(
      title = plot_title,
      subtitle = "Ordered by true RepeatMasker repeat_class",
      x = NULL,
      y = NULL
    ) +
    theme_minimal(base_size = 11) +
    theme(
      panel.grid = element_blank(),

      axis.text.y = element_blank(),
      axis.ticks.y = element_blank(),

      axis.text.x = element_text(
        size = 12,
        face = "bold"
      ),

      plot.title = element_text(
        size = 15,
        face = "bold",
        hjust = 0.5
      ),

      plot.subtitle = element_text(
        size = 10,
        hjust = 0.5
      ),

      legend.title = element_text(size = 10),
      legend.text = element_text(size = 9)
    )

  # -------------------------
  # CLASS STRIP
  # -------------------------

  strip <- ggplot(
    class_df,
    aes(x = 1, y = subfamily, fill = plot_class)
  ) +
    geom_tile() +
    scale_fill_manual(
      values = class_colors,
      drop = FALSE,
      name = "Repeat class"
    ) +
    labs(x = "Class") +
    theme_void() +
    theme(
      axis.title.x = element_text(
        size = 10,
        face = "bold",
        margin = margin(t = 6)
      ),

      legend.position = "right",

      legend.title = element_text(size = 10),
      legend.text = element_text(size = 9)
    )

  # -------------------------
  # COMBINE
  # -------------------------

  combined <- strip + heat +
    plot_layout(
      widths = c(0.035, 1),
      guides = "collect"
    ) &
    theme(
      legend.position = "right"
    )

  return(combined)
}

# =========================================================
# GENERATE
# =========================================================

p_beta <- make_heatmap(mat_beta, "beta", dt$plot_class)
p_z <- make_heatmap(mat_z, "zscore", dt$plot_class)

# =========================================================
# SAVE
# =========================================================

ggsave(
  out_pdf_beta,
  p_beta,
  width = 6,
  height = 8,
  useDingbats = FALSE
)

ggsave(
  out_pdf_zscore,
  p_z,
  width = 6,
  height = 8,
  useDingbats = FALSE
)

cat("\nSaved:\n")
cat(out_pdf_beta, "\n")
cat(out_pdf_zscore, "\n")
cat(out_class_summary, "\n\n")