library(tidyverse)

# ----------------------------------------------------------------------
# INPUT FILES -----------------------------------------------------------
# WGBS samples now provided as *.bg (columns: chr start end meth_rate methC unmethC)
# ONT samples remain *.methylC.gz (columns: chr start end context meth_rate strand cov)
# ----------------------------------------------------------------------
files <- c(
  # H9 / AN WGBS (BG)
  "/BLUES/eric/WGBS/TWFG-H9-primed-lib2/filtered_TWFG-H9-primed-lib2.bg",   # H9 primed (BG)
  "/BLUES/eric/ONT/Primed_W3MECP2_aln.methylC.gz",                          # ONT primed
  "/BLUES/eric/WGBS/TWFG-AN-naive/filtered_TWFG-AN-naive.bg",              # AN naive (BG)
  "/BLUES/eric/WGBS/TWFG-H9-naive-lib1/filtered_TWFG-H9-naive-lib1.bg",    # H9 naive (BG)
  "/BLUES/eric/ONT/091724_Naive_W3MECPC2_aln.methylC.gz",                  # ONT naive
  "/BLUES/eric/WGBS/TWFG-AN-naive_TSC-lib1/filtered_TWFG-AN-naive_TSC-lib1.bg", # AN TSC (BG)
  "/BLUES/eric/WGBS/TWFG-H9-naive_TSC-lib1/filtered_TWFG-H9-naive_TSC-lib1.bg", # H9 TSC (BG)
  "/BLUES/eric/ONT/TSC_GFP_W3MECP2_merged_aln.methylC.gz"                      # ONT TSC
)

labels <- c(
  "H9 primed", "W3MECP2 primed",
  "AN naive", "H9 naive", "W3MECP2 naive",
  "AN TSC", "H9 TSC", "W3MECP2 TSC"
)

# Helper: detect file type --------------------------------------------------
read_meth_file <- function(file, label) {
  if (grepl(".bg$", file)) {
    # BG format: chr start end meth_rate methC unmethC
    df <- read_tsv(file, col_names = c("chr","start","end","meth_rate","meth","unmeth"), show_col_types = FALSE) %>%
      mutate(coverage = meth + unmeth)
  } else {
    # methylC.gz format: chr start end context meth_rate strand coverage
    df <- read_tsv(file, comment = "#", col_names = c("chr","start","end","context","meth_rate","strand","coverage"), show_col_types = FALSE) %>%
      filter(context == "CG")
  }
  df %>%
    filter(!is.na(meth_rate)) %>%
    mutate(
      category = case_when(
        meth_rate < 0.2           ~ "0-0.2",
        meth_rate < 0.8           ~ "0.2-0.8",
        meth_rate <= 1            ~ "0.8-1",
        TRUE                      ~ NA_character_
      ),
      sample = label
    ) %>%
    filter(!is.na(category))
}

all_data <- map2_dfr(files, labels, read_meth_file)

# Summarise percent ---------------------------------------------------------
plot_df <- all_data %>%
  count(sample, category) %>%
  group_by(sample) %>%
  mutate(percent = 100 * n / sum(n)) %>%
  ungroup()

plot_df$sample   <- factor(plot_df$sample, levels = rev(labels))
plot_df$category <- factor(plot_df$category, levels = c("0.8-1","0.2-0.8","0-0.2"))

col_map <- c("0-0.2"  = scales::alpha("#8ecae6", 0.85),
             "0.2-0.8" = scales::alpha("#a6e3a1", 0.85),
             "0.8-1"   = scales::alpha("#f4a261", 0.85))

plot_df$is_ONT <- grepl("W3MECP2", plot_df$sample)

library(ggplot2)
p <- ggplot(plot_df, aes(x = sample, y = percent, fill = category)) +
  geom_col(width = 0.8, color = ifelse(plot_df$is_ONT, "black", NA), linewidth = ifelse(plot_df$is_ONT, 1.2, NA)) +
  coord_flip() +
  geom_text(aes(label = sprintf("%.1f%%", percent)), position = position_stack(vjust = 0.5), size = 3.3) +
  scale_y_continuous(expand = c(0,0)) +
  scale_fill_manual(values = col_map, name = "Methylation range", breaks = c("0.8-1","0.2-0.8","0-0.2")) +
  labs(title = "Percent methylation composition (WGBS BG + ONT methylC)", x = NULL, y = "Percentage") +
  theme_minimal(base_size = 13) +
  theme(panel.grid = element_blank(), axis.title.y = element_blank(), legend.position = "right")

print(p)

ggsave("Combined_WGBS_BG_ONT_barplot.png", p, width = 9, height = 6, dpi = 300, bg = "white")