library(memes)
library(universalmotif)
library(tidyverse)
library(plyranges)
library(extraChIPs)
library(rtracklayer)
library(magrittr)
library(BSgenome.Hsapiens.UCSC.hg19)
library(scales)
library(pander)
library(glue)
library(cowplot)
library(parallel)
library(gt)
library(parallel)
library(readxl)
library(UpSetR)
library(reactable)
library(htmltools)
library(multcomp)

options(meme_bin = "/opt/meme/bin/")
check_meme_install()
theme_set(theme_bw())
panderOptions("big.mark", ",")
n_cores <- 12

getBestMatch <- function(x, subject, min.score = "80%", rc = TRUE, ...) {
  
  pwm <- NULL
  if (is.matrix(x)) pwm <- x
  if (is(x, "universalmotif")) pwm <- x@motif
  stopifnot(!is.null(pwm))
  stopifnot(is(subject, "XStringViews"))
  
  matches <- matchPWM(
    pwm, subject, min.score = min.score, with.score = TRUE, ...
  )
  df <- as.data.frame(matches)
  df$score <- mcols(matches)$score
  df$strand <- "+"
  
  if (rc) {
    rc_matches <-  matchPWM(
      reverseComplement(pwm), subject ,
      min.score = min.score, with.score = TRUE, 
      ...
    )
    rc_df <- as.data.frame(rc_matches)
    rc_df$score <- mcols(rc_matches)$score
    rc_df$strand <- "-"
    df <- bind_rows(df, rc_df)
  }
  
  grp_cols <- c("start", "end", "width", "seq", "score")
  grp_df <- group_by(df, !!!syms(grp_cols))
  summ_df <- summarise(
    grp_df, strand = paste(strand, collapse = " "), .groups=  "drop"
  )
  summ_df$strand <- str_replace_all(summ_df$strand, ". .", "*")
  summ_df$strand <- factor(summ_df$strand, levels = c("+", "-", "*"))
  summ_df$view <- findInterval(summ_df$start, start(subject))
  summ_df <- arrange(summ_df, view, desc(score))
  
  best_df <- distinct(summ_df, view, .keep_all = TRUE)
  best_df <- mutate(
    best_df,
    start_in_view = start - start(subject)[view] + 1,
    end_in_view = start_in_view + width
  )
  ## This will only be added if the views have names
  best_df$view_name <- names(subject[best_df$view])
  
  best_df
  
}
bar_style <- function(width = 1, fill = "#e6e6e6", height = "75%", align = c("left", "right"), color = NULL) {
  align <- match.arg(align)
  if (align == "left") {
    position <- paste0(width * 100, "%")
    image <- sprintf("linear-gradient(90deg, %1$s %2$s, transparent %2$s)", fill, position)
  } else {
    position <- paste0(100 - width * 100, "%")
    image <- sprintf("linear-gradient(90deg, transparent %1$s, %2$s %1$s)", position, fill)
  }
  list(
    backgroundImage = image,
    backgroundSize = paste("100%", height),
    backgroundRepeat = "no-repeat",
    backgroundPosition = "center",
    color = color
  )
}

hg19 <- BSgenome.Hsapiens.UCSC.hg19
sq19 <- seqinfo(hg19)

all_gr <- here::here("data", "annotations", "all_gr.rds") %>%
  read_rds() %>% 
  lapply(
    function(x) {
      seqinfo(x) <- sq19
      x
    }
  )
counts <- here::here("data", "rnaseq", "counts.out.gz") %>%
  read_tsv(skip = 1) %>%
  dplyr::select(Geneid, ends_with("bam")) %>%
  dplyr::rename(gene_id = Geneid)
detected <- all_gr$gene %>%
  as_tibble() %>%
  distinct(gene_id, gene_name) %>%
  left_join(counts) %>%
  pivot_longer(
    cols = ends_with("bam"), names_to = "sample", values_to = "counts"
  ) %>%
  mutate(detected = counts > 0) %>%
  group_by(gene_id, gene_name) %>%
  summarise(detected = mean(detected) > 0.25, .groups = "drop") %>%
  dplyr::filter(detected) %>%
  dplyr::select(starts_with("gene"))

features <- here::here("data", "h3k27ac") %>%
  list.files(full.names = TRUE, pattern = "bed$") %>%
  sapply(import.bed, seqinfo = sq19) %>%
  lapply(granges) %>%
  setNames(basename(names(.))) %>%
  setNames(str_remove_all(names(.), "s.bed")) %>%
  GRangesList() %>%
  unlist() %>%
  names_to_column("feature") %>%
  sort()

uniprot_to_ensg <- here::here("data", "external", "uniprot_to_ensg.txt.gz") %>% 
  read_tsv(col_names = c("uniprot_id", "protein", "gene_id")) %>% 
  mutate(
    gene_id = lapply(gene_id, str_split, pattern = "; ") %>% 
      lapply(unlist)
  ) %>% 
  unnest(gene_id) %>% 
  mutate(gene_id = str_replace_all(gene_id, "(ENSG[0-9]+)\\..*", "\\1")) %>% 
  left_join(
    as_tibble(all_gr$gene) %>% 
      dplyr::select(gene_id, gene_name)
  ) %>% 
  dplyr::filter(!is.na(gene_id)) %>% 
  dplyr::select(-protein)
rime_detected <- here::here("data", "external", "MDA-MB-453_RIME_processed.xlsx") %>% 
  read_excel() %>% 
  dplyr::select(-starts_with("AVG"), -ends_with("Coverage2"), -contains("delta")) %>% 
  pivot_longer(cols = contains("_")) %>% 
  separate(name, c("sample", "measure"), sep = "_", extra = "merge") %>% 
  mutate(
    IP = str_remove(sample, "[0-9]$"),
    replicate = str_extract(sample, "[0-9]$"),
    measure = paste(measure, IP, sep = "_")
  ) %>% 
  pivot_wider(
    names_from = "measure", values_from = "value", 
    id_cols = c("Accession", "Description", "Protein", "replicate")
  ) %>% 
  mutate(
    logfc_cov = log2(Coverage_AR + 1) - log2(Coverage_IgG + 1),
    logfc_peptides = log2(Unique_peptides_AR + 1) - log2(Unique_peptides_IgG + 1)
  ) %>% 
  dplyr::filter(logfc_cov > 0) %>% 
  group_by(Accession, Description, Protein) %>% 
  summarise(
    n_detected = dplyr::n(),
    mn_logfc_coverage = mean(logfc_cov),
    mn_coverage_AR = mean(Coverage_AR),
    mn_coverage_IgG = mean(Coverage_IgG),
    mn_peptides_AR = mean(Unique_peptides_AR),
    mn_peptides_IgG = mean(Unique_peptides_IgG),
    .groups = "drop"
  ) %>% 
  dplyr::filter(n_detected > 1) %>% 
  arrange(desc(mn_logfc_coverage)) %>% 
  left_join(uniprot_to_ensg, by = c("Accession" = "uniprot_id")) %>% 
  dplyr::select(
    Accession, Description, Protein, starts_with("gene"), everything()
  ) %>% 
  dplyr::filter(!is.na(gene_name))

meme_db <- here::here("data", "external", "JASPAR2022_CORE_Hsapiens.meme") %>%
  read_meme() %>%
  to_df() %>%
  as_tibble() %>%
  mutate(
    across(everything(), vctrs::vec_proxy),
    altname = str_remove(altname, "MA[0-9]+\\.[0-9]*\\."),
    organism = "Homo sapiens"
  )
meme_altname <- setNames(meme_db$altname, meme_db$name)
meme_db_detected <- meme_db %>%
  mutate(
    gene_name = str_split(altname, pattern = "::", n = 2)
  ) %>%
  dplyr::filter(
    vapply(
      gene_name,
      function(x) all(x %in% detected$gene_name),
      logical(1)
    )
  )
meme_info <- meme_db_detected %>% 
  mutate(motif_length = vapply(motif, function(x){ncol(x@motif)}, integer(1))) %>% 
  dplyr::select(name, altname, icscore, motif_length) 

dht_peaks <- here::here("data", "peaks") %>%
  list.files(recursive = TRUE, pattern = "oracle", full.names = TRUE) %>%
  sapply(read_rds, simplify = FALSE) %>%
  lapply(function(x) x[["DHT"]]) %>%
  lapply(setNames, nm = c()) %>%
  setNames(str_extract_all(names(.), "AR|FOXA1|GATA3|TFAP2B")) %>% 
  lapply(
    function(x) {
      seqinfo(x) <- sq19
      x
    }
  )
targets <- names(dht_peaks)
dht_consensus <- dht_peaks %>%
  lapply(granges) %>%
  GRangesList() %>%
  unlist() %>%
  reduce() %>%
  mutate(
    AR = overlapsAny(., dht_peaks$AR),
    FOXA1 = overlapsAny(., dht_peaks$FOXA1),
    GATA3 = overlapsAny(., dht_peaks$GATA3),
    TFAP2B = overlapsAny(., dht_peaks$TFAP2B),
    promoter = bestOverlap(., features, var = "feature", missing = "None") == "promoter",
    enhancer = bestOverlap(., features, var = "feature", missing = "None") == "enhancer",
    H3K27ac = promoter | enhancer
  )

a <- dht_peaks %>% 
  lapply(mutate, w = width) %>% 
  lapply(as_tibble) %>% 
  bind_rows(.id = "target") %>% 
  ggplot(aes(w, stat(density), colour = target)) +
  geom_density() +
  scale_x_log10(breaks = c(100, 300, 600, 1e3, 2e3, 3e3), labels = comma) +
  labs(
    x = "Peak Width (bp)", y = "Density", colour = "Target"
  )
b <- dht_peaks %>% 
  lapply(mutate, w = width) %>% 
  lapply(as_tibble) %>% 
  bind_rows(.id = "target") %>% 
  mutate(skew = 0.5 - peak / w) %>% 
  ggplot(aes(skew, stat(density), colour = target)) +
  geom_density() +
  scale_x_continuous(labels = percent) +
  labs(
    x = "% Width Between Peak Summit and Peak Centre",
    y = "Density",
    colour = "ChIP Target"
  )
plot_grid(
  a + theme(legend.position = "none"), 
  b + theme(legend.position = "none"),
  get_legend(b),
  rel_widths = c(1, 1, 0.2), nrow = 1,
  labels = c("A", "B")
)

grl_by_targets <- dht_consensus %>% 
  as_tibble() %>% 
  pivot_longer(
    all_of(targets), names_to = "target", values_to = "detected"
  ) %>% 
  dplyr::filter(detected) %>% 
  group_by(range, H3K27ac) %>% 
  summarise(
    targets = paste(sort(target), collapse = " "),
    .groups = "drop"
  ) %>% 
  colToRanges(var = "range", seqinfo = sq19) %>% 
  sort() %>% 
  splitAsList(.$targets)

grl_by_targets %>% 
  lapply(
    function(x) {
      tibble(
        `N Sites` = length(x), 
        `% H3K27ac` = mean(x$H3K27ac),
        `Median Width` = floor(median(width(x)))
      )
    }
  ) %>% 
  bind_rows(.id = "Targets") %>% 
  mutate(
    `N Targets` = str_count(Targets, " ") + 1,
    `% H3K27ac` = percent(`% H3K27ac`, 0.1)
  ) %>% 
  dplyr::select(contains("Targets"), everything()) %>% 
  arrange(`N Targets`, Targets) %>% 
  pander(
    justify = "lrrrr",
    caption = "Summary of union ranges and which combinations of the targets were detected."
  )

seq_by_targets <- grl_by_targets %>% 
  lapply(function(x) split(x, x$H3K27ac)) %>% 
  lapply(GRangesList) %>% 
  lapply(get_sequence, genome = hg19)

## All bound sites
rds_ame_by_targets <- here::here("output", "ame_by_targets.rds")
if (!file.exists(rds_ame_by_targets)) {
  ame_by_targets <- seq_by_targets %>% 
    lapply(unlist) %>% 
    mclapply(
      runAme, 
      database = to_list(meme_db_detected), 
      mc.cores = min(n_cores, length(.))
    )
  write_rds(ame_by_targets, rds_ame_by_targets, compress = "gz")
} else {
  ame_by_targets <- read_rds(rds_ame_by_targets)
}

## Only sites marked as active by H3K27ac
rds_ame_by_targets_active <- here::here("output", "ame_by_targets_active.rds")
if (!file.exists(rds_ame_by_targets_active)) {
  ame_by_targets_active <- seq_by_targets %>% 
    lapply(function(x) x[["TRUE"]]) %>% 
    mclapply(
      runAme, 
      database = to_list(meme_db_detected), 
      mc.cores = min(n_cores, length(.))
    )
  write_rds(ame_by_targets_active, rds_ame_by_targets_active, compress = "gz")
} else {
  ame_by_targets_active <- read_rds(rds_ame_by_targets_active)
}

## Compare the active and inactive sites for each set of targets
rds_ame_active_vs_inactive <- here::here("output", "ame_active_vs_inactive.rds")
if (!file.exists(rds_ame_active_vs_inactive)) {
  ame_active_vs_inactive <- seq_by_targets %>% 
    mclapply(
      function(x) {
        runAme(x[["TRUE"]], x[["FALSE"]], database = to_list(meme_db_detected))
      },
      mc.cores = min(n_cores, length(.))
    )
  write_rds(ame_active_vs_inactive, rds_ame_active_vs_inactive, compress = "gz")
} else {
  ame_active_vs_inactive <- read_rds(rds_ame_active_vs_inactive)
}

## Compare the inactive and active sites for each set of targets
rds_ame_inactive_vs_active <- here::here("output", "ame_inactive_vs_active.rds")
if (!file.exists(rds_ame_inactive_vs_active)) {
  ame_inactive_vs_active <- seq_by_targets %>% 
    mclapply(
      function(x) {
        runAme(x[["FALSE"]], x[["TRUE"]], database = to_list(meme_db_detected))
      },
      mc.cores = min(n_cores, length(.))
    )
  write_rds(ame_inactive_vs_active, rds_ame_inactive_vs_active, compress = "gz")
} else {
  ame_inactive_vs_active <- read_rds(rds_ame_inactive_vs_active)
}