#!/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 <- file.path(
  outdir,
  "TE_DNA_methylation_heatmap_beta_trueRepeatClass_deltaTSCminusNaive_top10_boxes.pdf"
)

out_png <- file.path(
  outdir,
  "TE_DNA_methylation_heatmap_beta_trueRepeatClass_deltaTSCminusNaive_top10_boxes.png"
)

out_top10_table <- file.path(
  outdir,
  "TE_DNA_methylation_heatmap_beta_trueRepeatClass_deltaTSCminusNaive_top10_boxes.tsv"
)

out_class_summary <- file.path(
  outdir,
  "TE_DNA_methylation_heatmap_beta_trueRepeatClass_deltaTSCminusNaive_class_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)

  # BED-like format:
  # chr start end strand 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 = part before first comma
  te[, subfamily := tstrsplit(repeat_id, ",", fixed = TRUE, keep = 1L)]

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

  # sanity check: if any subfamily maps to >1 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 >1 repeat_class / repeat_family. Keeping first sorted annotation.",
      nrow(ambig)
    ))

    fwrite(
      ambig,
      file.path(outdir, "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
dt <- dt[complete.cases(dt[, ..col_order])]

# merge in 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_class_order <- c(
  "LINE",
  "SINE",
  "LTR",
  "Retroposon",
  "DNA",
  "RC",
  "Satellite",
  "Simple_repeat",
  "Low_complexity",
  "RNA",
  "Unknown",
  "DNA?",
  "LTR?",
  "RC?",
  "scRNA",
  "snRNA",
  "srpRNA",
  "tRNA",
  "rRNA",
  "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",
  "DNA?"           = "#fb8072",
  "LTR?"           = "#b3de69",
  "RC?"            = "#00bfc4",
  "scRNA"          = "#ff7f00",
  "snRNA"          = "#c77cff",
  "srpRNA"         = "#8dd3c7",
  "tRNA"           = "#ffffb3",
  "rRNA"           = "#bebada",
  "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]

# =========================================================
# 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)]
)

# =========================================================
# RANKING METRIC: TSC - NAIVE
# =========================================================

dt[, delta_TSC_minus_Naive := TSC - Naive]

# order by class, then delta high -> low
setorder(dt, plot_class, -delta_TSC_minus_Naive, subfamily)

# row positions
dt[, row_id := .I]
dt[, y := .N - row_id + 1]

# rank within class
dt[, rank_in_class := seq_len(.N), by = plot_class]

# =========================================================
# LONG FORMAT FOR HEATMAP
# =========================================================

plot_dt <- melt(
  dt,
  id.vars = c(
    "subfamily",
    "repeat_class",
    "repeat_family",
    "plot_class",
    "delta_TSC_minus_Naive",
    "row_id",
    "y",
    "rank_in_class"
  ),
  measure.vars = col_order,
  variable.name = "sample",
  value.name = "beta"
)

plot_dt[, sample := factor(sample, levels = col_order)]
plot_dt[, x := as.numeric(sample)]

# =========================================================
# CLASS STRIP DATA
# =========================================================

class_df <- unique(dt[, .(subfamily, plot_class, y)])

# =========================================================
# TOP 10 TABLE PER CLASS
# =========================================================

top10_dt <- dt[, head(.SD, 10), by = plot_class]

top10_dt[, repeat_family := fifelse(
  is.na(repeat_family) | repeat_family == "",
  "NA",
  repeat_family
)]

top10_dt[, display_line := paste0(
  rank_in_class, ". ", subfamily, " (", repeat_family, ")"
)]

fwrite(
  top10_dt[, .(
    repeat_class = as.character(plot_class),
    rank_in_class,
    subfamily,
    repeat_family,
    Primed,
    Naive,
    TSC,
    delta_TSC_minus_Naive
  )][order(repeat_class, rank_in_class)],
  out_top10_table,
  sep = "\t"
)

box_dt <- top10_dt[, .(
  label_text = paste(display_line, collapse = "\n")
), by = plot_class]

# =========================================================
# MAIN HEATMAP
# =========================================================

heat <- ggplot(plot_dt, aes(x = x, y = y, fill = beta)) +

  # width=1 removes the white gap between Primed/Naive/TSC
  geom_tile(width = 1, height = 1) +

  scale_fill_gradientn(
    colors = c("white", "#fff7bc", "#d7301f"),
    values = rescale(c(0, 0.5, 1)),
    limits = c(0, 1),
    name = "Beta"
  ) +

  scale_x_continuous(
    breaks = seq_along(col_order),
    labels = col_order,
    limits = c(0.5, 3.5),
    expand = expansion(mult = c(0, 0))
  ) +

  scale_y_continuous(
    limits = range(plot_dt$y),
    expand = expansion(mult = c(0, 0))
  ) +

  labs(
    title = "TE DNA methylation heatmap (beta values)",
    subtitle = "Grouped by true RepeatMasker repeat_class; within each class ranked by (TSC - Naive)",
    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 (THINNER)
# =========================================================

strip <- ggplot(
  class_df,
  aes(x = 1, y = y, fill = plot_class)
) +
  geom_tile(width = 1, height = 1) +
  scale_fill_manual(
    values = class_colors,
    drop = FALSE,
    name = "Repeat class"
  ) +
  scale_y_continuous(
    limits = range(plot_dt$y),
    expand = expansion(mult = c(0, 0))
  ) +
  labs(x = "Class", y = NULL) +
  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)
  )

# =========================================================
# TOP-10 BOX PANELS
# =========================================================

make_class_box <- function(cls, txt, border_col) {

  fill_col <- alpha(border_col, 0.08)

  ggplot(
    data.table(x = 0, y = 1, label = txt),
    aes(x = x, y = y, label = label)
  ) +
    annotate(
      "text",
      x = 0,
      y = 1.08,
      label = cls,
      hjust = 0,
      vjust = 0,
      fontface = "bold",
      size = 4,
      color = border_col
    ) +
    geom_label(
      hjust = 0,
      vjust = 1,
      size = 2.6,
      lineheight = 0.95,
      label.size = 0.35,
      label.r = grid::unit(0.12, "lines"),
      label.padding = grid::unit(0.18, "lines"),
      fill = fill_col,
      color = "black"
    ) +
    coord_cartesian(
      xlim = c(-0.02, 1),
      ylim = c(0, 1.12),
      clip = "off"
    ) +
    theme_void()
}

box_levels <- as.character(box_dt$plot_class)

box_plots <- lapply(box_levels, function(cls) {
  txt <- box_dt[plot_class == cls, label_text][1]
  make_class_box(cls, txt, class_colors[cls])
})

# 2-column boxed annotation panel
anno_panel <- wrap_plots(box_plots, ncol = 2)

# =========================================================
# COMBINE
# =========================================================

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

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

ggsave(
  out_pdf,
  combined,
  width = 14,
  height = 10,
  useDingbats = FALSE
)

ggsave(
  out_png,
  combined,
  width = 14,
  height = 10,
  dpi = 300
)

cat("\nSaved:\n")
cat(out_pdf, "\n")
cat(out_png, "\n")
cat(out_top10_table, "\n")
cat(out_class_summary, "\n\n")