---
title: "Transposon subclass expansions"
output: html_document
---
  
```{r setup, include=FALSE}
knitr::opts_chunk$set(
  echo = TRUE, message = FALSE, warning = FALSE,
  cache = TRUE, cache.path = '.cache_knitr/te_subclass_expansion/',
  fig.width = 7, fig.height = 4, dpi = 350,
  fig.path='figure/6-te_subclass/',
  dev = c('png', 'pdf')
)
```

In this document, we identify the TE subclasses intersecting with SVs.

## Setup
```{r, cache=FALSE}
library(AnnotationDbi)
library(GenomicRanges)
library(ggplot2)
library(tibble)
library(dplyr)
library(tidyr)
library(readr)
library(glue)
library(reactable)

# color schema for species
color_schema_species <- readr::read_csv(
  glue(Sys.getenv("HOME"), "/code/malawi_transposon/metadata/species_colorscheme.csv"), comment='#') %>% 
  mutate(color = paste0('#', color))

theme_set(theme_classic())
theme_update(
  panel.background = element_rect(fill = "transparent", colour = NA),
  plot.background = element_rect(fill = "transparent", colour = NA),
  legend.background = element_rect(fill = "transparent", colour = NA)
)

setwd(paste0(Sys.getenv("HOME"), "/code/malawi_transposon/script/R_SCRIPT/"))

source("modules/frequently_used_items.R")
source("modules/pseudogenome_functions.R")
```

## Import bubble information
```{r}
# import sv data frame
df_sv_raw <- readr::read_csv(
  paste0(Sys.getenv("HOME"), "/code/malawi_transposon/local/script/malawi_haplochromines-summary-preprocessed.csv"),
  col_types = paste0('cciiiiiffiiicilliidf', paste0(rep('ciicic', num_of_genomes), collapse=''), collapse=''),
  na = c("", "NA", ".")
)
```

Get the longest allele for each bubble, and the species containing it following a hierarchy of contiguousness.
```{r}
# list of species ordered by genome contiguity
species_list_contiguity <- unlist(strsplit('astCal,mayZeb,troMau,rhaChi,aulStu,otoArg,rhaChi2,copChr', split=','))

# select matrices of the paths and lengths
matrix_paths <- select(df_sv_raw, paste0(species_list_contiguity, "_path"))
matrix_lengths <- select(df_sv_raw, paste0(species_list_contiguity, "_len"))

# get the index of the longest allele (first occurrence) row-wise
idx_longest_allele <- apply(matrix_lengths, MARGIN=1, which.max) 

# use this index to retrieve the longest alleles and the species containing it
longest_alleles <- as.matrix(matrix_paths)[cbind(1:length(idx_longest_allele), idx_longest_allele)]
species_with_longest_alleles <- species_list_contiguity[idx_longest_allele]

# add to main dataframe
df_sv_raw$longest_allele <- longest_alleles
df_sv_raw$species_longest <- species_with_longest_alleles

# prepare new columns showing whether the longest allele is present in each species
df_sv_raw <- df_sv_raw %>% 
  mutate(
    across(
      .cols = all_of(paste0(color_schema_species$species_abrv, "_path")),
      .fns = function(x) x == longest_allele,
      .names = "{.col}_exist"
    )
  )

# count the number of times the longest allele shows up (takes a while to run)
df_sv_raw <- df_sv_raw %>% 
  rowwise() %>%
  mutate(longest_count = sum(c_across(ends_with("_exist")), na.rm=TRUE)) %>% 
  ungroup()
```

Filter for simple bubbles.
```{r}
# filter for simple bubbles only
df_sv <- df_sv_raw %>% filter(n_paths_bio <= 3, lenbio_min < 10, lenbio_minmaxratio <= 0.1)
glue("retained {nrow(df_sv)} out of {nrow(df_sv_raw)} bubbles")
```

## Import pseudogenome ranges (for each species)
For TEs, flexible nodes and chrom sizes.
```{r}
# import coordinates of flexible nodes, dropping zero-length ones
gr_flexnode_list <- lapply(color_schema_species$species_abrv, import_flex_coord_into_granges)
names(gr_flexnode_list) <- color_schema_species$species_abrv

# import coordinates of flexible nodes, dropping Simple_repeats, 
gr_te_list <- lapply(color_schema_species$species_abrv, import_transposons_into_granges_pseudo)
names(gr_te_list) <- color_schema_species$species_abrv

# get chrom sizes for pseudogenomes
chrom_info_list <- lapply(color_schema_species$species_abrv, import_chrom_sizes)
names(chrom_info_list) <- color_schema_species$species_abrv
```

## Intersection function
This is a (overly long) function to get the TE intersecting with the nodes of a given species for the longest allele.
```{r}
get_TE_intersect_in_longest_allele <- function(curr_species) {
  # get the longest allele for the current species
  tmp_df <- df_sv %>% filter(species_longest == curr_species)
  
  # get a list of alleles where the given species is the longest at the loci
  # but we parse them to remove the direction, and separate them as commas
  longest_alleles_for_species_parsed <- gsub('^,', '', gsub('[><]', ',',  tmp_df$longest_allele))
  
  # get a data frame of nodes in the longest alleles
  tmp_df2 <- tmp_df %>% 
    select(id, allele = longest_allele, freq = longest_count) %>% 
    mutate(node = gsub('^,', '', gsub('[><]', ',',  allele)) ) %>%
    separate_rows(node, sep=',')
  
  ####
  
  # we can create GRanges with the nodes separately
  # or "manually reduce" them before turning them into GRanges (to keep metadata information)
  # the latter is easier for our purposes
  # but will not calculate nested variation as clearly
  tmp_df2_reduced <- tmp_df2 %>% 
    left_join(
      as.data.frame(gr_flexnode_list[[curr_species]]) %>% select(chr = seqnames, start, end, node = node_name)
    ) %>% 
    group_by(allele, chr, id, freq) %>% 
    summarise(start = min(start), end = max(end), total_length = sum(end-start+1)) %>% 
    arrange(chr, start, end) %>% 
    ungroup()
  
  # create GRanges
  gr_sv_curr <- GRanges(
    seqnames = tmp_df2_reduced$chr, 
    ranges = IRanges(start = tmp_df2_reduced$start, end = tmp_df2_reduced$end),
    freq = tmp_df2_reduced$freq, bubble_id = tmp_df2_reduced$id, #total_length = tmp_df2_reduced$total_length
    seqinfo = chrom_info_list[[curr_species]]
  )
  
  ######
  
  # intersect with TEs, keeping those that intersect more than 90% of TE within the allele
  gr_a <- gr_sv_curr
  
  gr_b <- gr_te_list[[curr_species]]
  gr_b$te_len <- width(gr_b)
  
  
  gr_intersect <- findOverlapPairs(gr_a, gr_b, ignore.strand=TRUE) %>% pintersect(ignore.strand=FALSE)
  gr_intersect_reciprocal <- findOverlapPairs(gr_b, gr_a, ignore.strand=TRUE) %>% pintersect(ignore.strand=FALSE)
  
  
  df_intersect <- bind_cols(
    as.data.frame(gr_intersect), 
    as.data.frame(gr_intersect_reciprocal) %>% select(name, score, repeat_class, repeat_subclass, te_len)
    ) %>% 
    mutate(te_perc_within_sv = width / te_len)
  
  return(df_intersect)
}
```

## Running it on the species
We run this function on each species, afterwards we add information from the original SV dataframe.
```{r}
df_intersect <- bind_rows(
    lapply(color_schema_species$species_abrv, get_TE_intersect_in_longest_allele)
  ) %>% 
  left_join(df_sv_raw %>% select(bubble_id = id, ends_with('_exist')) ) %>% 
  separate(bubble_id, c('backbone_chr', 'backbone_start', 'backbone_end'), remove = FALSE, convert = TRUE) %>% 
  arrange(backbone_chr, backbone_start)

head(df_intersect)
```

Filter for TE hits.
```{r}
df_intersect_filter <- df_intersect %>% filter(te_perc_within_sv > 0.8)

glue("{nrow(df_intersect_filter)} hits within {length(unique(df_intersect_filter$bubble_id))} unique bubbles")
```


## Plots
Get average count by subclass.
```{r}
df_count_threshold <- df_intersect_filter %>% 
  group_by(repeat_subclass, repeat_class) %>% 
  summarise(n_entries = n())

count_threshold <- filter(df_count_threshold, repeat_subclass != 'Unknown') %>% .$n_entries %>% quantile(0.5)
repeat_subclasses_to_plot <- filter(df_count_threshold, n_entries > count_threshold) %>% .$repeat_subclass

count_threshold
```


```{r te_frequency_by_class}
df_intersect_filter %>% 
  ggplot(aes(freq)) +
  geom_bar() +
  facet_wrap(.~repeat_class, scales = 'free_y')
```

```{r teinsert_count_by_subclass-PAPER, fig.width=3.8, fig.height=4.6}
df_intersect_filter %>% 
  filter(repeat_subclass %in% repeat_subclasses_to_plot, repeat_subclass != 'DNA') %>% 
  ggplot(aes(repeat_subclass, fill = repeat_class)) +
  geom_bar(color='black') +
  geom_text(
    aes(x = repeat_subclass, y = n_entries, label = n_entries), 
    data = filter(df_count_threshold, repeat_subclass %in% repeat_subclasses_to_plot, repeat_subclass != 'DNA'), 
    angle = 0, size = 2.5, hjust = -0.15 
  ) +
  xlab(NULL) +
  scale_x_discrete(limits = rev) +
  scale_y_continuous(
    "Number of TE insertions in bubbles", expand = expansion(mult = c(.0, .2)) 
  ) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  coord_flip() +
  theme(
    legend.position = 'none',
    legend.key.size=unit(0.02, "npc"),
    legend.title = element_text(size=10), 
    legend.text = element_text(size=9)
  )
```

Paper version, plotting only the most noticable ones.
```{r teinsert_count_by_subclass, fig.width=6, fig.height=3.8}
df_intersect_filter %>% 
  ggplot(aes(repeat_subclass)) +
  geom_bar(aes(fill = score <= 10)) +
  geom_text(
    aes(x = repeat_subclass, y = n_entries, label = n_entries), 
    data = filter(df_count_threshold, n_entries > count_threshold), 
    angle = 90, size = 2.5, hjust = -0.15 
  ) +
  xlab(NULL) +
  scale_y_continuous(
    "Number of TE fragments in bubbles", expand = expansion(mult = c(.05, .2)) 
  ) +
  scale_fill_manual("sequence divergence <= 10%", values = c('darkorange', 'royalblue')) +
  theme(
    legend.position = 'top',
    legend.key.size=unit(0.02, "npc"),
    legend.title = element_text(size=10), 
    legend.text = element_text(size=9),
    axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1)
  )
```

```{r teinsert_allelefreq_by_subclass, fig.width = 6.2, fig.height = 7.4}
df_intersect_filter %>% 
  filter(repeat_subclass %in% unlist(strsplit(repeat_subclasses_to_plot, ','))) %>% 
  ggplot(aes(freq, fill = score <= 10)) +
  geom_bar() +
  xlab("Insertion frequency across samples (out of 8)") +
  scale_y_continuous(
    "Number of TE fragments in bubbles", expand = expansion(mult = c(.05, .2)) 
  ) +
  scale_fill_manual("sequence divergence <= 10%", values = c('darkorange', 'royalblue')) +
  facet_wrap(.~repeat_subclass, scales = 'free_y', ncol = 4) +
  theme(
    legend.position = 'top',
    legend.key.size=unit(0.02, "npc"),
    legend.title = element_text(size=10), 
    legend.text = element_text(size=9),
    axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1)
  )
```

```{r teinsert_allelefreq_by_subclass-KEYNOTE, fig.width = 9.8, fig.height = 4.6}
df_intersect_filter %>% 
  filter(repeat_subclass %in% unlist(strsplit(repeat_subclasses_to_plot, ','))) %>% 
  ggplot(aes(freq, fill = score <= 10)) +
  geom_bar() +
  xlab("Insertion frequency across samples (out of 8)") +
  scale_y_continuous(
    "Number of TE fragments in bubbles", expand = expansion(mult = c(.05, .2)) 
  ) +
  scale_fill_manual("sequence divergence <= 10%", values = c('darkorange', 'royalblue')) +
  facet_wrap(.~repeat_subclass, scales = 'free_y', nrow = 3) +
  theme(
    legend.position = 'top',
    legend.key.size=unit(0.02, "npc"),
    legend.title = element_text(size=10), 
    legend.text = element_text(size=9),
    axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1)
  )
```

Write TE intersect fragments into text file for analysis in separate document.
```{r}
df_intersect_filter %>% write_csv("./local/malawi_haplochromines-te-intersects.csv")
```
