#!/usr/bin/env Rscript

suppressPackageStartupMessages({
  library(data.table)
})

setDTthreads(8)

# =====================================================
# INPUT / OUTPUT
# =====================================================

base_dir <- "/BLUES/eric/ONT_WGBS/Figure_3/Dynamic_copies_table"

infile <- file.path(
  base_dir,
  "TE_table_W3_H9_min5CpG.tsv"
)

outfile <- file.path(
  base_dir,
  "W3_TE_methylation_change_over_genome_change.tsv"
)

# =====================================================
# FILTER THRESHOLD
# =====================================================

ratio_cutoff <- 1.5

# =====================================================
# LOAD
# =====================================================

dt <- fread(infile)

cat("Loaded rows:", nrow(dt), "\n")
# =====================================================
# CALCULATE GENOME-WIDE CHANGES
# =====================================================

genome_PN_methylation_change <- mean(
  dt$W3_Naive_avg_meth -
    dt$W3_Primed_avg_meth,
  na.rm = TRUE
)

genome_TN_methylation_change <- mean(
  dt$W3_TSC_avg_meth -
    dt$W3_Naive_avg_meth,
  na.rm = TRUE
)

cat("Genome-wide P?N change:",
    round(genome_PN_methylation_change, 4), "\n")

cat("Genome-wide N?T change:",
    round(genome_TN_methylation_change, 4), "\n")
    
# =====================================================
# CALCULATE COPY-LEVEL METHYLATION CHANGES
# =====================================================

dt[, W3_PN_TE_methylation_change :=
      W3_Naive_avg_meth - W3_Primed_avg_meth]

dt[, W3_TN_TE_methylation_change :=
      W3_TSC_avg_meth - W3_Naive_avg_meth]

# =====================================================
# CALCULATE RATIO OVER GENOME CHANGE
# =====================================================

# direction-aware ratio:
# negative / negative = positive
# positive / positive = positive

dt[, PN_TE_change_over_genome_change :=
      W3_PN_TE_methylation_change /
      genome_PN_methylation_change]

dt[, TN_TE_change_over_genome_change :=
      W3_TN_TE_methylation_change /
      genome_TN_methylation_change]

# =====================================================
# SUMMARIZE BY SUBFAMILY
# =====================================================

summary_dt <- dt[
  ,
  .(
    mean_PN_TE_methylation_change = round(
      mean(W3_PN_TE_methylation_change,
           na.rm = TRUE),
      4
    ),

    mean_TN_TE_methylation_change = round(
      mean(W3_TN_TE_methylation_change,
           na.rm = TRUE),
      4
    ),

    genome_PN_methylation_change =
      genome_PN_methylation_change,

    genome_TN_methylation_change =
      genome_TN_methylation_change,

    PN_TE_change_over_genome_change = round(
      mean(PN_TE_change_over_genome_change,
           na.rm = TRUE),
      4
    ),

    TN_TE_change_over_genome_change = round(
      mean(TN_TE_change_over_genome_change,
           na.rm = TRUE),
      4
    )
  ),

  by = .(
    subfamily = TE_name,
    family,
    class
  )
]

# =====================================================
# KEEP ONLY SUBFAMILIES
# WITH RATIO > 1.5
# =====================================================

summary_dt <- summary_dt[
  PN_TE_change_over_genome_change > ratio_cutoff |
    TN_TE_change_over_genome_change > ratio_cutoff
]

# =====================================================
# SORT
# =====================================================

setorder(
  summary_dt,
  -PN_TE_change_over_genome_change,
  -TN_TE_change_over_genome_change
)

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

fwrite(
  summary_dt,
  outfile,
  sep = "\t",
  quote = FALSE
)

cat("Saved:\n")
cat(outfile, "\n")

cat("Rows retained:", nrow(summary_dt), "\n")