#!/usr/bin/env Rscript

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

setDTthreads(8)

BASE <- "/BLUES/eric/ONT_WGBS/Allele_analysis"

OUTDIR <- file.path(BASE, "Comprehensive_Trend_Figure")
dir.create(OUTDIR, showWarnings = FALSE)

# =========================================================
# Helper
# =========================================================

read_noX <- function(f) {
  dt <- fread(f)

  if ("chr" %in% names(dt)) {
    dt[, chr := gsub('"', '', chr)]
    dt <- dt[chr != "chrX"]
  } else if ("chr_p" %in% names(dt)) {
    dt[, chr_p := gsub('"', '', chr_p)]
    dt <- dt[chr_p != "chrX"]
  } else {
    stop("No chr or chr_p column found in: ", f)
  }

  dt
}

# =========================================================
# LOAD FILES, chrX removed
# =========================================================

prime_pos <- read_noX(file.path(
  BASE,
  "window_delta_prime_d0.5_cpg5_meancov10",
  "prime_windows_1000bp_delta_ge_0.5_cpg5_meancov10.tsv"
))

prime_neg <- read_noX(file.path(
  BASE,
  "window_delta_prime_d0.5_cpg5_meancov10",
  "prime_windows_1000bp_delta_le_-0.5_cpg5_meancov10.tsv"
))

naive_pos <- read_noX(file.path(
  BASE,
  "window_delta_naive_in_prime_windows_cpg5_meancov10",
  "naive_in_prime_windows_1000bp_prime_delta_ge_0.5_cpg5_meancov10.tsv"
))

naive_neg <- read_noX(file.path(
  BASE,
  "window_delta_naive_in_prime_windows_cpg5_meancov10",
  "naive_in_prime_windows_1000bp_prime_delta_le_-0.5_cpg5_meancov10.tsv"
))

tsc_pos <- read_noX(file.path(
  BASE,
  "window_delta_TSC_in_prime_windows_cpg5_meancov10",
  "TSC_in_prime_windows_1000bp_prime_delta_ge_0.5_cpg5_meancov10.tsv"
))

tsc_neg <- read_noX(file.path(
  BASE,
  "window_delta_TSC_in_prime_windows_cpg5_meancov10",
  "TSC_in_prime_windows_1000bp_prime_delta_le_-0.5_cpg5_meancov10.tsv"
))

same_pos <- read_noX(file.path(
  BASE,
  "TSC_in_prime_same_trend_cpg5_meancov10",
  "same_pos.tsv"
))

same_neg <- read_noX(file.path(
  BASE,
  "TSC_in_prime_same_trend_cpg5_meancov10",
  "same_neg.tsv"
))

rev_pos <- read_noX(file.path(
  BASE,
  "TSC_in_prime_reverse_trend_cpg5_meancov10",
  "flip_pos_to_neg.tsv"
))

rev_neg <- read_noX(file.path(
  BASE,
  "TSC_in_prime_reverse_trend_cpg5_meancov10",
  "flip_neg_to_pos.tsv"
))

# =========================================================
# COMMON THEME
# =========================================================

common_theme <- theme_bw(base_size = 12) +
  theme(
    plot.title = element_text(size = 12, face = "bold"),
    plot.subtitle = element_text(size = 10),
    axis.title.x = element_text(size = 11),
    axis.title.y = element_text(size = 11),
    axis.text = element_text(size = 10),
    panel.grid.major.y = element_line(color = "grey85", linewidth = 0.4),
    panel.grid.minor.y = element_blank(),
    panel.grid.major.x = element_blank()
  )

# =========================================================
# PLOT FUNCTION
# =========================================================

make_plot <- function(dt, h1_col, h2_col, title_label, subtitle_label) {

  d <- rbind(
    data.table(hap = "h1", meth = dt[[h1_col]]),
    data.table(hap = "h2", meth = dt[[h2_col]])
  )

  d[, hap := factor(hap, levels = c("h1", "h2"))]

  ggplot(d, aes(hap, meth)) +
    geom_boxplot(width = 0.6) +
    scale_y_continuous(
      limits = c(0, 1),
      breaks = seq(0, 1, 0.1)
    ) +
    labs(
      title = title_label,
      subtitle = subtitle_label,
      x = "Haplotype",
      y = "Mean methylation (coverage-weighted)"
    ) +
    common_theme
}

# =========================================================
# PRIME ROW
# =========================================================

p_prime_neg <- make_plot(
  prime_neg, "h1_w", "h2_w",
  paste0("Prime h1 - h2 < -0.5 (n=", nrow(prime_neg), ")"),
  "Prime methylation (1000 bp; chrX removed)"
)

p_prime_pos <- make_plot(
  prime_pos, "h1_w", "h2_w",
  paste0("Prime h1 - h2 > 0.5 (n=", nrow(prime_pos), ")"),
  "Prime methylation (1000 bp; chrX removed)"
)

prime_row <- p_prime_neg + p_prime_pos

# =========================================================
# NAIVE ROW
# =========================================================

p_naive_neg <- make_plot(
  naive_neg, "h1_w", "h2_w",
  paste0("Prime h1 - h2 < -0.5 (n=", nrow(naive_neg), ")"),
  "Naive methylation (Prime-defined windows; 1000 bp; chrX removed)"
)

p_naive_pos <- make_plot(
  naive_pos, "h1_w", "h2_w",
  paste0("Prime h1 - h2 > 0.5 (n=", nrow(naive_pos), ")"),
  "Naive methylation (Prime-defined windows; 1000 bp; chrX removed)"
)

naive_row <- p_naive_neg + p_naive_pos

# =========================================================
# TSC ROW
# =========================================================

p_tsc_neg <- make_plot(
  tsc_neg, "h1_w", "h2_w",
  paste0("Prime h1 - h2 < -0.5 (n=", nrow(tsc_neg), ")"),
  "TSC methylation (Prime-defined windows; 1000 bp; chrX removed)"
)

p_tsc_pos <- make_plot(
  tsc_pos, "h1_w", "h2_w",
  paste0("Prime h1 - h2 > 0.5 (n=", nrow(tsc_pos), ")"),
  "TSC methylation (Prime-defined windows; 1000 bp; chrX removed)"
)

tsc_row <- p_tsc_neg + p_tsc_pos

# =========================================================
# SAME TREND PANELS
# =========================================================

bottom_same_neg <- make_plot(
  same_neg, "h1_w_t", "h2_w_t",
  paste0("Prime h1 - h2 < -0.5 & TSC h1 - h2 < 0 (n=", nrow(same_neg), ")"),
  "TSC methylation (chrX removed)"
)

bottom_same_pos <- make_plot(
  same_pos, "h1_w_t", "h2_w_t",
  paste0("Prime h1 - h2 > 0.5 & TSC h1 - h2 > 0 (n=", nrow(same_pos), ")"),
  "TSC methylation (chrX removed)"
)

same_row <- bottom_same_neg + bottom_same_pos

# =========================================================
# REVERSE TREND PANELS
# =========================================================

bottom_rev_neg <- make_plot(
  rev_neg, "h1_w_t", "h2_w_t",
  paste0("Prime h1 - h2 < -0.5 & TSC h1 - h2 > 0 (n=", nrow(rev_neg), ")"),
  "TSC methylation (chrX removed)"
)

bottom_rev_pos <- make_plot(
  rev_pos, "h1_w_t", "h2_w_t",
  paste0("Prime h1 - h2 > 0.5 & TSC h1 - h2 < 0 (n=", nrow(rev_pos), ")"),
  "TSC methylation (chrX removed)"
)

reverse_row <- bottom_rev_neg + bottom_rev_pos

blank <- plot_spacer()

# =========================================================
# FINAL FIGURE
# =========================================================

final_plot <-
  (
    prime_row /
      naive_row /
      tsc_row /
      blank /
      same_row /
      reverse_row
  ) +
  plot_layout(
    heights = c(1, 1, 1, 0.12, 1, 1)
  ) +
  plot_annotation(
    title = "Allele-specific methylation dynamics across cell states",
    subtitle = paste(
      "1000 bp windows; CpG >= 5 per hap;",
      "mean coverage >= 10 per hap; chrX removed"
    ),
    theme = theme(
      plot.title = element_text(size = 18, face = "bold"),
      plot.subtitle = element_text(size = 13)
    )
  )

ggsave(
  file.path(OUTDIR, "Comprehensive_prime_naive_TSC_trends_vertical_noX_cov10.pdf"),
  final_plot,
  width = 14,
  height = 28
)

cat("\n[DONE] Figure saved to:\n")
cat(file.path(OUTDIR, "Comprehensive_prime_naive_TSC_trends_vertical_noX_cov10.pdf"), "\n\n")