#!/usr/bin/env Rscript

suppressPackageStartupMessages({
  library(data.table)
})

setDTthreads(8)

# =========================================================
# Settings
# =========================================================

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

TE_BED_GZ <- "/BLUES/eric/ONT/hg38_TE_noY.bed.gz"

# MAPQ >= 10, primary alignments, chrY removed
METHYLC_FILES <- list(
  "Primed" = paste0(
    "/BLUES/eric/ONT/MAPQ_bam/",
    "ONT_Primed.MAPQ10.primary.noY.onC.methylC.gz"
  ),
  "Naive" = paste0(
    "/BLUES/eric/ONT/MAPQ_bam/",
    "ONT_Naive.MAPQ10.primary.noY.onC.methylC.gz"
  ),
  "TSC" = paste0(
    "/BLUES/eric/ONT/MAPQ_bam/",
    "ONT_TSC.MAPQ10.primary.noY.onC.methylC.gz"
  )
)

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

MIN_CPG_PER_SUBFAM <- 50L
MIN_TOTAL_BP_SUBFAM <- 1000L
REMOVE_RRNA_TRNA <- TRUE

OUT_SUMMARY <- file.path(
  OUTDIR,
  "W3_TEsubfamily_DNA_methylation_summary_MAPQ10_primary_noY_no_rRNA_tRNA.tsv"
)

OUT_MATRIX <- file.path(
  OUTDIR,
  "DNA_methylation_matrix_FULL_MAPQ10_primary_noY_no_rRNA_tRNA.tsv"
)

OUT_FILTER_STATS <- file.path(
  OUTDIR,
  "DNA_methylation_FULL_MAPQ10_filter_stats.tsv"
)

# =========================================================
# Helpers
# =========================================================

parse_te_subfamily <- function(x) {
  repeat_id <- sub(
    '.*repeat_id\\s+"([^"]+)".*',
    '\\1',
    x
  )

  sub(",.*$", "", repeat_id)
}

is_rrna_trna <- function(x) {
  grepl(
    "(\\b|_)(r+rna|t+rna)(\\b|_)",
    x,
    ignore.case = TRUE
  ) |
    grepl(
      "rrna|trna",
      x,
      ignore.case = TRUE
    )
}

read_te_bed <- function(file) {
  if (!file.exists(file)) {
    stop("Missing TE annotation: ", file)
  }

  command <- sprintf(
    "zcat %s",
    shQuote(file)
  )

  te <- fread(
    cmd = command,
    sep = "\t",
    header = FALSE,
    showProgress = FALSE
  )

  if (ncol(te) < 5) {
    stop(
      "Expected at least 5 TE columns but found ",
      ncol(te)
    )
  }

  setnames(
    te,
    1:5,
    c("chr", "start", "end", "strand", "info")
  )

  te[, subfamily := parse_te_subfamily(info)]

  te[, `:=`(
    chr = as.character(chr),
    start = as.integer(start),
    end = as.integer(end)
  )]

  te <- te[
    !is.na(subfamily) &
    subfamily != "" &
    !is.na(start) &
    !is.na(end) &
    end > start
  ]

  te[, width := end - start]

  setkey(te, chr, start, end)
  te
}

read_methylC <- function(file, sample_name) {
  if (!file.exists(file)) {
    stop("Missing methylC file: ", file)
  }

  command <- sprintf(
    "zcat %s",
    shQuote(file)
  )

  dt <- fread(
    cmd = command,
    sep = "\t",
    header = FALSE,
    showProgress = FALSE
  )

  if (ncol(dt) < 7) {
    stop(
      "Expected at least 7 methylC columns but found ",
      ncol(dt),
      " in ",
      file
    )
  }

  setnames(
    dt,
    1:7,
    c(
      "chr",
      "start",
      "end",
      "context",
      "value",
      "strand",
      "cov"
    )
  )

  # CpG methylation only
  dt <- dt[context == "CG"]

  dt[, `:=`(
    chr = as.character(chr),
    start = as.integer(start),
    end = as.integer(end),
    value = as.numeric(value),
    cov = as.numeric(cov),
    sample = sample_name
  )]

  dt <- dt[
    !is.na(chr) &
    !is.na(start) &
    !is.na(end) &
    is.finite(value) &
    is.finite(cov) &
    cov > 0
  ]

  setkey(dt, chr, start, end)
  dt
}

summarize_methylation <- function(te, methyl_dt) {
  overlaps <- foverlaps(
    te,
    methyl_dt,
    type = "any",
    nomatch = 0L
  )

  overlaps[, .(
    nCpG = .N,
    mean_value = mean(
      value,
      na.rm = TRUE
    ),
    wmean_value = weighted.mean(
      value,
      w = cov,
      na.rm = TRUE
    )
  ), by = .(subfamily, sample)]
}

make_matrix <- function(summary_dt) {
  wide <- dcast(
    summary_dt,
    subfamily ~ sample,
    value.var = "wmean_value",
    fill = NA_real_
  )

  for (sample_name in COL_ORDER) {
    if (!sample_name %in% names(wide)) {
      wide[, (sample_name) := NA_real_]
    }
  }

  setcolorder(
    wide,
    c("subfamily", COL_ORDER)
  )

  wide
}

# =========================================================
# Validate inputs
# =========================================================

missing_files <- unlist(METHYLC_FILES)[
  !file.exists(unlist(METHYLC_FILES))
]

if (length(missing_files) > 0) {
  stop(
    "Missing MAPQ10 methylC files:\n",
    paste(missing_files, collapse = "\n")
  )
}

# =========================================================
# Read and filter TE annotation
# =========================================================

message("[1/6] Reading TE annotation...")

te <- read_te_bed(TE_BED_GZ)

initial_subfamilies <- uniqueN(te$subfamily)

te_bp <- te[, .(
  total_bp = sum(width),
  n_loci = .N
), by = subfamily]

keep_by_bp <- te_bp[
  total_bp >= MIN_TOTAL_BP_SUBFAM,
  subfamily
]

te <- te[subfamily %in% keep_by_bp]

after_bp_subfamilies <- uniqueN(te$subfamily)

if (REMOVE_RRNA_TRNA) {
  te <- te[!is_rrna_trna(subfamily)]
}

after_rrna_subfamilies <- uniqueN(te$subfamily)

message(
  "  Initial TE subfamilies: ",
  initial_subfamilies
)

message(
  "  After total bp >= ",
  MIN_TOTAL_BP_SUBFAM,
  ": ",
  after_bp_subfamilies
)

message(
  "  After rRNA/tRNA removal: ",
  after_rrna_subfamilies
)

# =========================================================
# Read MAPQ10 methylation
# =========================================================

message("[2/6] Reading MAPQ10 methylC files...")

methyl_list <- lapply(
  names(METHYLC_FILES),
  function(sample_name) {
    message("  Reading ", sample_name, "...")

    read_methylC(
      METHYLC_FILES[[sample_name]],
      sample_name
    )
  }
)

methyl_dt <- rbindlist(
  methyl_list,
  use.names = TRUE
)

setkey(methyl_dt, chr, start, end)

# =========================================================
# Overlap CpGs with TE copies and summarize
# =========================================================

message("[3/6] Overlapping CpGs with TE annotations...")

methyl_summary <- summarize_methylation(
  te,
  methyl_dt
)

before_cpg_filter <- uniqueN(
  methyl_summary$subfamily
)

# Require >=50 CpGs in every condition
passing_samples <- methyl_summary[
  nCpG >= MIN_CPG_PER_SUBFAM,
  .(n_passing_samples = uniqueN(sample)),
  by = subfamily
]

keep_by_cpg <- passing_samples[
  n_passing_samples == length(COL_ORDER),
  subfamily
]

methyl_summary <- methyl_summary[
  subfamily %in% keep_by_cpg
]

after_cpg_filter <- uniqueN(
  methyl_summary$subfamily
)

# Defensive removal
if (REMOVE_RRNA_TRNA) {
  methyl_summary <- methyl_summary[
    !is_rrna_trna(subfamily)
  ]
}

message(
  "  Subfamilies before CpG filter: ",
  before_cpg_filter
)

message(
  "  Subfamilies with >= ",
  MIN_CPG_PER_SUBFAM,
  " CpGs in all three conditions: ",
  after_cpg_filter
)

# =========================================================
# Write long-format summary
# =========================================================

message("[4/6] Writing MAPQ10 summary...")

# Sort conditions as Primed, Naive, TSC
methyl_summary[
  ,
  sample_order := match(sample, COL_ORDER)
]

setorder(
  methyl_summary,
  subfamily,
  sample_order
)

methyl_summary[
  ,
  sample_order := NULL
]

fwrite(
  methyl_summary,
  OUT_SUMMARY,
  sep = "\t"
)

# =========================================================
# Create and write FULL matrix
# =========================================================

message("[5/6] Writing FULL MAPQ10 matrix...")

methyl_matrix <- make_matrix(
  methyl_summary
)

fwrite(
  methyl_matrix,
  OUT_MATRIX,
  sep = "\t"
)

# =========================================================
# Write filter statistics
# =========================================================

filter_stats <- data.table(
  step = c(
    "Initial TE annotation",
    paste0("Total TE bp >= ", MIN_TOTAL_BP_SUBFAM),
    "After rRNA/tRNA removal",
    "Observed in methylC overlap",
    paste0(
      "CpGs >= ",
      MIN_CPG_PER_SUBFAM,
      " in all 3 conditions"
    )
  ),
  n_subfamilies = c(
    initial_subfamilies,
    after_bp_subfamilies,
    after_rrna_subfamilies,
    before_cpg_filter,
    after_cpg_filter
  )
)

fwrite(
  filter_stats,
  OUT_FILTER_STATS,
  sep = "\t"
)

message("[6/6] Done.")

cat("\nWritten:\n")
cat(OUT_SUMMARY, "\n")
cat(OUT_MATRIX, "\n")
cat(OUT_FILTER_STATS, "\n")