---
title: "Linking SVs to Transposable Elements (pseudogenome)"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(
  echo = TRUE, message = FALSE, warning = FALSE,
  cache = TRUE, cache.path = '.cache_knitr/pseudo_TE_stats/',
  fig.width = 7, fig.height = 4, dpi = 300,
  fig.path='figure/3-calculate_TE_pseudo/',
  dev = c('png', 'pdf')
)
```
 
In this document, we investigate the biological context of the SVs, and their association with transposable elements. The main focus is the plot the differences in TE composition at the bubbles across the different assemblies. Due to time constraints, my approach for this is to intersect the flexible nodes with the transposon annotations on the pseudogenomes built on the backbone. The scripts in the main folder of `script/D_pseudoreference/` need to be run first.

I've also written this script to be flexible in case we need to rerun this with a different backbone.

At the moment, we also consider the TE composition across all flexible nodes regardless of the success rate in allele calling. This shouldn't pose a problem because we will calculate this as a percentage of total flexible node length by species. 

In case we have to filter for only flexible nodes in ubiquitous bubbles, we need to modify `1-get_paths_w_coverage_from_gfa.sh` to include the bubble ID in the name of the flexible nodes BED files. This might come in useful later when we want to filter for bubbles explained by polymorphic TE in at least one of the samples (for the genome tracks), might refer to `7-extract_TE_in_longest_alleles.R` for this.

Please make sure `2a-analyse_sv_properties.Rmd` has already been run. We need its output files as input for this document.

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

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/get_chrom_sizes.R")
source("modules/generic_granges_functions.R")
source("modules/pseudogenome_functions.R")

##########
subdir <- "malawi_haplochromines"
bb_name <- "astCal1.2_ensembl"
bb_name_sqlite <- "astCal_v1.2"
bb_name_abrv <- "astCal"
```

## Data import
### Import data from pseudogenomes
```{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
```

### Intersect SV with TE
```{r}
gr_intersect_list <- lapply(color_schema_species$species_abrv, intersect_flexnode_with_te)
names(gr_intersect_list) <- color_schema_species$species_abrv
```

### Calculate total sizes of pseudogenomes
```{r}
df_pseudo_sizes <- tibble(
    species = color_schema_species$species_abrv, 
    genome_size = sapply(gr_flexnode_list, function(x) sum(seqlengths(x)) ),
    flex_size = sapply(gr_flexnode_list, function(x) sum(width(x)) ) 
  ) %>% 
  mutate(species = factor(species, levels = color_schema_species$species_abrv))

head(df_pseudo_sizes)
```

## Rainbow plot (for calliptera pseudogenome)
Should be quite similar to the original, so I think it's fine to use it interchangably.

Genomewide:
```{r rainbow_calliptera-genomewide, fig.height=5, fig.width=8}
gr_te_list$astCal %>% 
  as.data.frame() %>% 
  filter(repmask_id != "rnd-1_family-37") %>% 
  left_join(df_pseudo_sizes, by="species") %>% 
  ggplot(aes(score+0.01, width / genome_size * 100, fill=repeat_class)) +
  stat_summary_bin(geom='bar', fun='sum', position='stack', breaks = seq(0, 49, 0.5)) +
  scale_fill_manual('Transposon Class', values = color_schema_te$color, breaks = color_schema_te$te_family) +
  scale_x_continuous(sec.axis = sec_axis(
    ~ . / 100 * generation_time / mu / 1e6,
    breaks = c(0.9, 9.2, 20, 62.1), # c(0.9, 9.2, 20, 62.1, 87.3, 150),
    labels = c(
      "Onset of Lake Malawi radiation (0.9 mya)",
      "Haplochromine-Lamprologni lineage split (9.2 mya)",
      "Split from Oreochromis (20.0 mya)",
      "East Africa cichlids split from Cichlidae (62.1 mya)") )) +
      #"Cichlidae split from Percomorphaceae (87.3 mya)",
      #"Gondwana supercontinent begins to breakup (150 mya)") )) +
  geom_vline(xintercept = c(0.11, 1.07, 2.33, 7.245), linetype="dotted") +
  #geom_vline(xintercept = c(0.11, 1.07, 2.33, 7.245, 10.185, 17.5), linetype="dotted") +
  xlab("Divergence from ancestral consensus (%)") + 
  ylab("Percentage sequence of genome") +
  coord_cartesian(ylim=c(0, NA), xlim = c(0, 40)) +
  theme(axis.text.x.top = element_text(angle = 45, hjust = 0),
        legend.position = c(0.9, 0.5))
```

Flexible regions / bubbles:
```{r rainbow_calliptera-flexible, fig.height=5, fig.width=8}
gr_intersect_list$astCal %>% 
  as.data.frame() %>% 
  filter(repmask_id != "rnd-1_family-37") %>% 
  left_join(df_pseudo_sizes, by="species") %>% 
  ggplot(aes(te_score+0.01, width / flex_size * 100, fill=te_class)) +
  stat_summary_bin(geom='bar', fun='sum', position='stack', breaks = seq(0, 49, 0.5)) +
  scale_fill_manual('Transposon Class', values = color_schema_te$color, breaks = color_schema_te$te_family) +
  scale_x_continuous(sec.axis = sec_axis(
    ~ . / 100 * generation_time / mu / 1e6,
    breaks = c(0.9, 9.2, 20, 62.1), # c(0.9, 9.2, 20, 62.1, 87.3, 150),
    labels = c(
      "Onset of Lake Malawi radiation (0.9 mya)",
      "Haplochromine-Lamprologni lineage split (9.2 mya)",
      "Split from Oreochromis (20.0 mya)",
      "East Africa cichlids split from Cichlidae (62.1 mya)") )) +
      #"Cichlidae split from Percomorphaceae (87.3 mya)",
      #"Gondwana supercontinent begins to breakup (150 mya)") )) +
  geom_vline(xintercept = c(0.11, 1.07, 2.33, 7.245), linetype="dotted") +
  #geom_vline(xintercept = c(0.11, 1.07, 2.33, 7.245, 10.185, 17.5), linetype="dotted") +
  xlab("Divergence from ancestral consensus (%)") + 
  ylab("Percentage sequence of flexible region") +
  coord_cartesian(ylim=c(0, NA), xlim = c(0, 40)) +
  theme(axis.text.x.top = element_text(angle = 45, hjust = 0),
        legend.position = c(0.9, 0.5))
```

## TE divergence at bubbles by samples
Compile TE percentages in flexible regions by TE class and across species.
```{r}
df_intersect_compiled <- c(
  # all TEs
  lapply(gr_intersect_list, function(x) as.data.frame(x) %>% mutate(te_class = "All Transposons")),
  
  # all TEs except Unknown fragments
  lapply(gr_intersect_list, function(x) as.data.frame(x) %>% filter(te_class != "Unknown") %>% mutate(te_class = "All Transposons excluding Unknown")),
  
  # the TE classes separately
  lapply(gr_intersect_list, as.data.frame)  ) %>% 
  bind_rows() %>% 
  mutate(te_class = factor(te_class,
    levels = c('All Transposons', 'All Transposons excluding Unknown', 'DNA', 'LINE', 'LTR', 'SINE', 'Retroposon', 'Helitron', 'Unknown'))
  ) %>% 
  left_join(df_pseudo_sizes, by="species")
```

Visualise it:
```{r te_composition_divergence, fig.height=8, fig.width=6.5}
df_intersect_compiled %>% 
  filter(repmask_id != "rnd-1_family-37") %>% 
  ggplot(aes(te_score+0.01, width / flex_size * 100, color=species)) +
  stat_summary_bin(geom='line', fun='sum', position='identity', breaks = seq(0, 49, 0.5)) +
  scale_fill_manual('Transposon Class', values = color_schema_te$color, breaks = color_schema_te$te_family) +
  facet_wrap(.~te_class, scales='free_y', ncol = 2) +
  geom_vline(xintercept = c(0.11, 1.07, 2.33, 7.245), linetype="dotted") +
  #geom_vline(xintercept = c(0.11, 1.07, 2.33, 7.245, 10.185, 17.5), linetype="dotted") +
  scale_color_manual('Species', values = color_schema_species$color, breaks = color_schema_species$species_abrv) +
  xlab("Divergence from ancestral consensus (%)") +
  ylab("Percentage sequence of flexible region") +
  scale_x_continuous(breaks = seq(0, 12, 2)) +
  coord_cartesian(xlim=c(0,11)) +
  theme(
    axis.text.x.top = element_text(angle = 45, hjust = 0),
    panel.grid.major.y = element_line(colour = "grey80"),
    legend.position = 'top',
    legend.title = element_text(size = rel(0.8)),
    legend.text = element_text(size = rel(0.7))
  )
```
## Finer subclass calculations
Aggregate data first.
```{r}
df_intersect_list_agg_subclass <- df_intersect_compiled %>% 
  filter(!te_class %in% c('All Transposons', 'All Transposons excluding Unknown')) %>% 
  group_by(species, te_class, te_name, te_subclass, flex_size) %>% 
  summarise(width_sum = sum(width) ) %>% 
  ungroup() %>% 
  mutate(perc_sum = round(width_sum / flex_size * 100, 3))
```

For the selected subset of species.
```{r subclass_perc_seq_subset, fig.width=5.5, fig.height=8}
species_subset <- c("astCal", "troMau", "copChr")

df_intersect_list_agg_subclass %>% 
  filter(species %in% species_subset, te_class != 'Unknown') %>% 
  ggplot(aes(x = perc_sum, y = te_name, fill = te_class )) +
  geom_col(color='black') +
  geom_text(aes(label = paste0(format(round(perc_sum, 1), n_small=1), "%") ), 
            hjust=-0.1, vjust=0.5, size=2.4) +
  facet_wrap(.~factor(species, levels=species_subset), nrow=1) +
  scale_x_continuous(
    'Percentage sequence of SVs annotated as TE subclass',
    expand = expansion(mult = c(.05, .35))
  ) +
  scale_y_discrete(NULL, limits=rev) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  theme_bw() +
  theme(legend.position = 'none')
```

Just for calliptera.
```{r subclass_perc_seq_subset_astCal, fig.width=3.2, fig.height=8}
df_intersect_list_agg_subclass %>% 
  filter(species == 'astCal', te_class != 'Unknown') %>% 
  ggplot(aes(x = perc_sum, y = te_name, fill = te_class )) +
  geom_col(color='black') +
  geom_text(aes(label = paste0(format(round(perc_sum, 1), n_small=1), "%") ), 
            hjust=-0.1, vjust=0.5, size=2.4) +
  facet_wrap(.~factor(species, levels=species_subset), nrow=1) +
  scale_x_continuous(
    'Percentage sequence (SVs)',
    expand = expansion(mult = c(.05, .35))
  ) +
  scale_y_discrete(NULL, limits=rev) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  theme_bw() +
  theme(legend.position = 'none')
```


## Calculate TE composition at bubbles
Now we want to make bar charts and tables of the TE composition. We need to perform a `GenomicsRanges::reduce` function to avoid overestimation of the size of these regions.
```{r}
calculate_item_percentage <- function(species, 
    gr_i_list = gr_intersect_list, gr_f_list = gr_flexnode_list, 
    te_class = NULL, te_age = NULL, te_len = NULL) {
  
  gr_i <- gr_i_list[[species]] # feature
  gr_f <- gr_f_list[[species]] # genomewide / flex node
  
  if (!is.null(te_class)) {
    if (length(te_class) == 1) {
      gr_i <- gr_i[gr_i$te_class == te_class]
    } else {
      gr_i <- gr_i[gr_i$te_class %in% te_class]
    }
  }
  
  if (!is.null(te_age)) {
    gr_i <- gr_i[gr_i$te_score < te_age]
  }
  
  if (!is.null(te_len)) {
    gr_i <- gr_i[gr_i$te_len > te_len]
  }
  
  ans <- calculate_reduced_length(gr_i) / calculate_reduced_length(gr_f)
  return(ans)
}
```

Sweep it across different TE names sets. Filter out TEs smaller than 50bp and older than 25% divergence.
```{r}
df_te_compiled <- bind_rows(
    lapply(color_schema_te$te_family,
      function(x) {
        result <- sapply(
          color_schema_species$species_abrv, calculate_item_percentage,
          gr_i_list = gr_intersect_list, gr_f_list = gr_flexnode_list,
          te_class = x, te_age = 25, te_len = 50
        )
        return(tibble(te_class = x, species=names(result), te_perc=result))
      }
    )
  ) %>% 
  mutate(
    species = factor(species, levels = color_schema_species$species_abrv),
    te_class = factor(te_class, levels = color_schema_te$te_family)
  )
```

```{r class_perc_seq-flexible, fig.width=12, fig.height=3.5}
df_te_compiled %>% 
  ggplot(aes(x=te_perc*100, y=te_class, fill=te_class)) +
  geom_col(color='black') +
  geom_text(aes(label = paste0(format(round(te_perc*100, 2), n_small=2), "%")), 
            hjust=-0.1, vjust=0.5, size=2.4) +
  facet_wrap(. ~species, nrow=1) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  scale_x_continuous('Percentage sequence of A. calliptera flexible regions annotated as TE class', 
                     expand = expansion(mult = c(.05, .4))) +
  scale_y_discrete(NULL, limits=rev) +
  theme_bw() +
  theme(legend.position = 'none')
```

Show the flexible node percentage for all species. The genomewide for pseudogenome is not reliable, so we don't do this here. Check `0-compare_ownVScalip_RepModel.Rmd` for the genomewide percentages.

All TEs and repetitive elements.
```{r}
bind_rows(
    lapply(c(50, 25, 10, 5, 1),
      function(x) {
        result <- sapply(color_schema_species$species_abrv, calculate_item_percentage, 
                         te_class = NULL, te_age = x, te_len = 50)
        return(tibble(age_filter = x, species=names(result), te_perc=round(result*100, 2)) )
      }
    )
  ) %>% 
  pivot_wider(id_cols="age_filter", names_from=species, values_from=te_perc) %>% 
  reactable()
```

All TEs excluding Unknowns.
```{r}
bind_rows(
    lapply(c(50, 25, 10, 5, 1),
      function(x) {
        result <- sapply(color_schema_species$species_abrv, calculate_item_percentage, 
                         te_class = c('DNA', 'LINE', 'LTR', 'SINE', 'Retroposon', 'Helitron'), te_age = x, te_len = 50)
        return(tibble(age_filter = x, species=names(result), te_perc=round(result*100, 2)) )
      }
    )
  ) %>% 
  pivot_wider(id_cols="age_filter", names_from=species, values_from=te_perc) %>% 
  reactable()
```

Table: percentages of different TE classes in flexible region at different age thresholds.
```{r}
tmp_df <- tibble()
for (age_filter in c(50, 25, 10, 5, 1)) {
  tmp_df_curr <- bind_rows(
      lapply(color_schema_te$te_family,
        function(x) {
          result <- sapply(
            color_schema_species$species_abrv, calculate_item_percentage,
            gr_i_list = gr_intersect_list, gr_f_list = gr_flexnode_list,
            te_class = x, te_age = age_filter, te_len = 50
          )
          return(tibble(age = age_filter, te_class = x, species=names(result), te_perc=round(result*100, 2) ))
        }
      )
    ) %>% 
    pivot_wider(id_cols=c('age', 'te_class'), names_from=species, values_from=te_perc) 
  tmp_df <- bind_rows(tmp_df, tmp_df_curr)
}
reactable(tmp_df); rm(tmp_df)
```

## Fold enrichment
This has to be done only with the backbone assembly. We want to calculate this for the major TE classes, as well as try a range of age divergence filters.

Import TE coordinates on native backbone.
```{r}
source("modules/import_transposons_into_granges.R")
gr_te_bb <- import_transposons_into_granges(bb_name_abrv)
gr_te_bb$te_len <- width(gr_te_bb)

# get intersect
gr_intersect_bb <- gr_intersect_list[[bb_name_abrv]]
```

Import SV coordinates.
```{r}
df_sv <- read_csv(
  glue("./local/{subdir}-summary-preprocessed.csv"),
  col_types = paste0('cciiiiicciiicilliidc', paste0(rep('ciicic', num_of_genomes), collapse=''), collapse=''),
  na = c("", "NA", ".")
)

chrom_sizes <- read_delim(
  glue(Sys.getenv("HOME"), "/igv/genomes/{bb_name}/chrom.sizes"),
  col_names = c('chr', 'len')
)
chrom_sizes <- Seqinfo(chrom_sizes$chr, chrom_sizes$len)
bb_size <- sum(seqlengths(chrom_sizes))

gr_sv_bb <- GRanges(
  seqnames = df_sv$chr, 
  ranges = IRanges(start = df_sv$start, end = df_sv$end), id = df_sv$id, 
  seqinfo = chrom_sizes
)
```

Calculate fold enrichment of TEs in SVs relative to the wider genome.
```{r}
calculate_fold_enrichment <- function(te_class, te_age = NULL, te_len = NULL) {
  if (!is.null(te_class)) {
    if (length(te_class) == 1) {
      gr_te_bb <- gr_te_bb[gr_te_bb$repeat_class == te_class]
      gr_intersect_bb <- gr_intersect_bb[gr_intersect_bb$te_class == te_class]
    } else {
      gr_te_bb <- gr_te_bb[gr_te_bb$repeat_class %in% te_class]
      gr_intersect_bb <- gr_intersect_bb[gr_intersect_bb$te_class %in% te_class]
    }
  }
  
  if (!is.null(te_age)) {
    gr_te_bb <- gr_te_bb[gr_te_bb$score < te_age]
    gr_intersect_bb <- gr_intersect_bb[gr_intersect_bb$te_score < te_age]
  }
  
  if (!is.null(te_len)) {
    gr_te_bb <- gr_te_bb[gr_te_bb$te_len > te_len]
    gr_intersect_bb <- gr_intersect_bb[gr_intersect_bb$te_len > te_len]
  }
  
  te_perc_genome <- calculate_reduced_length(gr_te_bb) / bb_size
  te_perc_flex <- calculate_reduced_length(gr_intersect_bb) / calculate_reduced_length(gr_sv_bb) 
  
  return(
    tibble(te_genome = te_perc_genome, te_flex = te_perc_flex)
  )
}

tmp_df <- tibble()
for(age_filter in c(50, 25, 10, 5, 1) ) {
  tmp_df_curr <- c(
      list(
        calculate_fold_enrichment(NULL, te_len = 0, te_age = age_filter) %>% 
          add_column(te = "All TEs")
      ),
      list(
        calculate_fold_enrichment(c('DNA', 'LINE', 'LTR', 'SINE', 'Retroposon', 'Helitron'), te_len = 0, te_age = age_filter) %>% 
          add_column(te = "All TEs excluding Unknown") 
      ),
      lapply(c('DNA', 'LINE', 'LTR', 'SINE', 'Retroposon', 'Helitron', 'Unknown'), 
        function(x) {
          calculate_fold_enrichment(x, te_len = 0, te_age = age_filter) %>% add_column(te = x)
        }
      )
    ) %>%
    bind_rows() %>% 
    add_column(age_filter = age_filter)
  
  tmp_df <- bind_rows(tmp_df, tmp_df_curr)
  rm(tmp_df_curr)
}

df_fe_bb <- tmp_df %>% 
  mutate(fe = te_flex / te_genome)
rm(tmp_df)
```

Visualise results.
```{r te_fold_enrichment, fig.width = 4.8, fig.height = 3}
df_fe_bb %>% 
  filter(!te %in% c('SINE', 'Retroposon', 'Helitron') ) %>% 
  ggplot(aes(x = factor(age_filter), y = fe, fill=te) ) +
  geom_col(color = 'black', width=0.7, position=position_dodge()) +
  theme(panel.grid.major.y = element_line(colour = "grey80"), legend.position = 'top') +
  scale_fill_manual(NULL, values = c('gray60', 'gray20', color_schema_te$color[c(1:3, 7)])) +
  scale_y_continuous('Fold enrichment in bubbles', breaks = seq(0, 5)) +
  scale_x_discrete('Divergence from ancestral consensus',
    labels = c('<1%', '<5%', '<10%', '<25%', 'all')
  ) +
  theme(
    legend.title = element_text(size = rel(0.9)),
    legend.text = element_text(size = rel(0.9)),
    legend.key.size=unit(0.03, "npc")
  )
```

```{r}
df_fe_bb %>% 
  dplyr::relocate(te, age_filter, fe) %>%
  mutate(across(where(is.double), ~ round(., 4)) ) %>% 
  filter(age_filter %in% c(50, 10)) %>% 
  arrange(age_filter) %>%
  reactable()
```

Genomewide vs SV, side by side:
```{r te_genome_v_sv, fig.width = 5.2, fig.height = 2.9}
library(ggpattern)
df_fe_bb %>% 
  filter(age_filter==50, !te %in% c("All TEs excluding Unknown", "All TEs")) %>% 
  dplyr::select(te, te_genome, te_flex) %>% 
  pivot_longer(te_genome:te_flex, names_pattern='te_(.*)', values_to="perc_seq") %>% 
  mutate(
    perc_seq = perc_seq*100, 
    name = if_else(name=="genome", "Genome-wide", "Structural Variants")
  ) %>% 
  ggplot(aes(x=factor(te, levels=color_schema_te$te_family), y=perc_seq, fill=te, pattern=name)) +
  geom_col_pattern(color='black', position='dodge', pattern_fill='black') +
  geom_text(aes(label = paste0(format(round(perc_seq, 2), n_small=2), "%")), hjust=0.5, vjust=-0.7, size=2.5) +
  facet_wrap(.~name) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  scale_y_continuous(expand = expansion(mult = c(.05, .1))) +
  xlab(NULL) +
  ylab('Percentage sequence') +
  #theme_linedraw() +
  theme(
    axis.text.x = element_text(angle = 35, hjust = 1), 
    legend.position = 'none'
  )
```

Young TEs classified as 10% divergence and below.
```{r te_genome_v_sv-young, fig.width = 5.2, fig.height = 2.9}
df_fe_bb %>% 
  filter(age_filter==10, !te %in% c("All TEs excluding Unknown", "All TEs")) %>% 
  dplyr::select(te, te_genome, te_flex) %>% 
  pivot_longer(te_genome:te_flex, names_pattern='te_(.*)', values_to="perc_seq") %>% 
  mutate(
    perc_seq = perc_seq*100, 
    name = if_else(name=="genome", "Genomewide", "Structural Variants")
  ) %>% 
  ggplot(aes(x=factor(te, levels=color_schema_te$te_family), y=perc_seq, fill=te, pattern=name)) +
  geom_col_pattern(color='black', position='dodge', pattern_fill='black') +
  geom_text(aes(label = paste0(format(round(perc_seq, 2), n_small=2), "%")), hjust=0.5, vjust=-0.7, size=2.5) +
  facet_wrap(.~name) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  scale_y_continuous(expand = expansion(mult = c(.05, .1))) +
  xlab(NULL) +
  ylab('Percentage sequence') +
  theme_linedraw() +
  theme(
    axis.text.x = element_text(angle = 35, hjust = 1), 
    legend.position = 'none'
  )
```

