#!/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_P_to_N_methylation_change <- mean(
  dt$W3_Naive_avg_meth - dt$W3_Primed_avg_meth,
  na.rm = TRUE
)

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

cat("Genome-wide P_to_N change:",
    round(genome_P_to_N_methylation_change, 4), "\n")

cat("Genome-wide N_to_T change:",
    round(genome_N_to_T_methylation_change, 4), "\n")

# =====================================================
# CALCULATE COPY-LEVEL METHYLATION CHANGES
# =====================================================

dt[, W3_P_to_N_TE_methylation_change :=
     W3_Naive_avg_meth - W3_Primed_avg_meth]

dt[, W3_N_to_T_TE_methylation_change :=
     W3_TSC_avg_meth - W3_Naive_avg_meth]

# =====================================================
# CALCULATE COPY-LEVEL RATIO OVER GENOME CHANGE
# =====================================================

dt[, P_to_N_TE_change_over_genome_change :=
     W3_P_to_N_TE_methylation_change /
     genome_P_to_N_methylation_change]

dt[, N_to_T_TE_change_over_genome_change :=
     W3_N_to_T_TE_methylation_change /
     genome_N_to_T_methylation_change]

# =====================================================
# DEFINE DYNAMIC COPIES
# =====================================================
# A dynamic copy means its copy-level change is >1.5x
# the genome-wide change in the same direction.

dt[, P_to_N_dynamic_copy :=
     P_to_N_TE_change_over_genome_change > ratio_cutoff]

dt[, N_to_T_dynamic_copy :=
     N_to_T_TE_change_over_genome_change > ratio_cutoff]

# =====================================================
# GENOME-WIDE BACKGROUND COUNTS
# =====================================================

total_TE_copies <- nrow(dt)

total_P_to_N_dynamic_copies <- dt[
  P_to_N_dynamic_copy == TRUE,
  .N
]

total_N_to_T_dynamic_copies <- dt[
  N_to_T_dynamic_copy == TRUE,
  .N
]

background_P_to_N_dynamic_fraction <-
  total_P_to_N_dynamic_copies / total_TE_copies

background_N_to_T_dynamic_fraction <-
  total_N_to_T_dynamic_copies / total_TE_copies

cat("Total TE copies:", total_TE_copies, "\n")
cat("Total P_to_N dynamic copies:",
    total_P_to_N_dynamic_copies, "\n")
cat("Total N_to_T dynamic copies:",
    total_N_to_T_dynamic_copies, "\n")

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

summary_dt <- dt[
  ,
  .(
    total_TE_copy_count = .N,

    P_to_N_dynamic_copy_count = sum(
      P_to_N_dynamic_copy,
      na.rm = TRUE
    ),

    N_to_T_dynamic_copy_count = sum(
      N_to_T_dynamic_copy,
      na.rm = TRUE
    ),

    mean_P_to_N_TE_methylation_change = round(
      mean(W3_P_to_N_TE_methylation_change, na.rm = TRUE),
      4
    ),

    mean_N_to_T_TE_methylation_change = round(
      mean(W3_N_to_T_TE_methylation_change, na.rm = TRUE),
      4
    ),

    genome_P_to_N_methylation_change =
      genome_P_to_N_methylation_change,

    genome_N_to_T_methylation_change =
      genome_N_to_T_methylation_change,

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

    N_to_T_TE_change_over_genome_change = round(
      mean(N_to_T_TE_change_over_genome_change, na.rm = TRUE),
      4
    )
  ),
  by = .(
    subfamily = TE_name,
    family,
    class
  )
]

# =====================================================
# CALCULATE FRACTION + ENRICHMENT
# =====================================================

summary_dt[, P_to_N_dynamic_fraction :=
             P_to_N_dynamic_copy_count / total_TE_copy_count]

summary_dt[, N_to_T_dynamic_fraction :=
             N_to_T_dynamic_copy_count / total_TE_copy_count]

summary_dt[, P_to_N_dynamic_enrichment :=
             P_to_N_dynamic_fraction /
             background_P_to_N_dynamic_fraction]

summary_dt[, N_to_T_dynamic_enrichment :=
             N_to_T_dynamic_fraction /
             background_N_to_T_dynamic_fraction]

# =====================================================
# FISHER TEST FOR ENRICHMENT
# =====================================================

summary_dt[, P_to_N_fisher_p := mapply(
  function(k, n) {
    mat <- matrix(
      c(
        k,
        n - k,
        total_P_to_N_dynamic_copies - k,
        total_TE_copies - n - total_P_to_N_dynamic_copies + k
      ),
      nrow = 2
    )
    fisher.test(mat, alternative = "greater")$p.value
  },
  P_to_N_dynamic_copy_count,
  total_TE_copy_count
)]

summary_dt[, N_to_T_fisher_p := mapply(
  function(k, n) {
    mat <- matrix(
      c(
        k,
        n - k,
        total_N_to_T_dynamic_copies - k,
        total_TE_copies - n - total_N_to_T_dynamic_copies + k
      ),
      nrow = 2
    )
    fisher.test(mat, alternative = "greater")$p.value
  },
  N_to_T_dynamic_copy_count,
  total_TE_copy_count
)]

summary_dt[, P_to_N_fisher_FDR :=
             p.adjust(P_to_N_fisher_p, method = "BH")]

summary_dt[, N_to_T_fisher_FDR :=
             p.adjust(N_to_T_fisher_p, method = "BH")]

# =====================================================
# ROUND NUMERIC COLUMNS
# =====================================================

round_cols <- c(
  "P_to_N_dynamic_fraction",
  "N_to_T_dynamic_fraction",
  "P_to_N_dynamic_enrichment",
  "N_to_T_dynamic_enrichment",
  "P_to_N_fisher_p",
  "N_to_T_fisher_p",
  "P_to_N_fisher_FDR",
  "N_to_T_fisher_FDR"
)

summary_dt[, (round_cols) := lapply(.SD, round, 4), .SDcols = round_cols]

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

summary_dt <- summary_dt[
  P_to_N_TE_change_over_genome_change > ratio_cutoff |
    N_to_T_TE_change_over_genome_change > ratio_cutoff
]

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

setorder(
  summary_dt,
  -P_to_N_TE_change_over_genome_change,
  -N_to_T_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")