---
title: "Preprocessing TEs"
output: html_document
---
  
```{r setup, include=FALSE}
knitr::opts_chunk$set(
  echo = TRUE, message = FALSE, warning = FALSE,
  cache = TRUE, cache.path = '.cache_knitr/preprocess_te/',
  fig.width = 8, fig.height = 5, dpi = 350,
  fig.path='figure/1-preprocess_te/',
  dev = c('png', 'pdf')
)
```

In this document, we are utilising various metrics to understand the properties of the TE annotations produced by RepeatModeler. Here we focus on the astCal library and see what kind of statistics it gives for the other species.

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

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/compare_annotation_overlap.R") # function to computes the overlap between annotations

# get genome sizes of the species as a dataframe
df_genome_sizes <- tibble(
  species = color_schema_species$species_abrv, 
  genome_size = sapply(chrom_sizes_list, function(x) sum(x$length)) ) %>% 
  mutate(species = factor(species, levels = color_schema_species$species_abrv))
```

## Data import and preprocessing
```{r}
source("modules/import_transposons_into_granges.R")

species_subset <- c("astCal", "troMau", "copChr")

gr_te_list <- lapply(color_schema_species$species_abrv, 
                     import_transposons_into_granges, annoType="calliptera")
names(gr_te_list) <- color_schema_species$species_abrv

df_te_list_raw <- bind_rows(lapply(gr_te_list, as.data.frame)) %>% 
  mutate(species = factor(species, levels = color_schema_species$species_abrv))
```

Aggregate data at various levels.
```{r}
# summarise key statistics grouped by individual RepeatModeler families (e.g. ltr-1_family-10)
df_te_list_agg_family <- df_te_list_raw %>% 
  group_by(repmask_id, species, name, repeat_class) %>% 
  summarise(
    width_mean = mean(width),
    width_max = max(width),
    width_sd = sd(width),
    width_sum = sum(width),
    score_mean = mean(score),
    score_sd = sd(score),
    count = n() ) %>% 
  ungroup() %>% 
  left_join(df_genome_sizes) %>% 
  mutate(perc_sum = round(width_sum / genome_size * 100, 3))

# summarise key statistics grouped by TE subclass (e.g. DNA/TcMar2)
df_te_list_agg_subclass <- df_te_list_agg_family %>% 
  group_by(name, repeat_class, species) %>% 
  summarise(count = sum(count), width_sum = sum(width_sum), n_family = n()) %>% 
  left_join(df_genome_sizes) %>% 
  mutate(perc_sum = round(width_sum / genome_size * 100, 3))

# summarise key statistics grouped by TE class (e.g. DNA)
df_te_list_agg_class <- df_te_list_agg_subclass %>% 
  group_by(repeat_class, species) %>% 
  summarise(count = sum(count), width_sum = sum(width_sum), n_family = sum(n_family)) %>% 
  left_join(df_genome_sizes) %>% 
  mutate(perc_sum = round(width_sum / genome_size * 100, 3))
```

## Identifying small and potentially artefactual RepeatModeler families
As shown below, some of the RepeatModeler families are extremely small, some only containing one individual fragment and not adding up to much total sequence length. I suspect these are artifacts from the sequence clustering step. 
```{r count_v_totalwidth, fig.width=6, fig.height=3.6}
df_te_list_agg_family %>% 
  filter(species == "astCal") %>% 
  ggplot(aes(x = count, y = width_sum, color = repeat_class)) +
  geom_density2d() +
  geom_point(color='black', size=0.2, alpha=0.25) +
  facet_wrap(.~repeat_class, nrow = 2) +
  coord_cartesian(xlim=c(1,NA), ylim=c(1,NA)) +
  scale_x_log10("number of fragments in RepeatModeler family", labels = trans_format("log10", math_format(10^.x)) ) +
  scale_y_log10("total sequence length of fragments", labels = trans_format("log10", math_format(10^.x)) ) +
  scale_color_manual(values = color_schema_te$color_dark, breaks = color_schema_te$te_family) +
  theme_bw() +
  theme(legend.position = 'none')
```

We could use some criteria to filter out potentially spurious TE families. Below shows the number of families removed by TE class and by species.
```{r}
threshold_count <- 10
threshold_width_sum <- 1000

tmp_df_removed <- df_te_list_agg_family %>% 
  filter(count < threshold_count, width_sum < threshold_width_sum) %>% 
  filter(species == "astCal")

# write down repmask_ids into a text file to remove
tibble(repmask_id=tmp_df_removed$repmask_id) %>% arrange(repmask_id) %>% write_delim("./local/repmaskid_to_remove.txt", col_names=FALSE)

# count number of repeat classes that are removed
tmp_df_removed %>% group_by(repeat_class) %>% summarise(n())
```

This would be what remains.
```{r}
df_te_list_agg_family %>% 
  filter(count >= threshold_count | width_sum >= threshold_count) %>% 
  filter(species == "astCal") %>% 
  group_by(repeat_class) %>% 
  summarise(n())
```

The actual code to remove these false families. Create a Boolean column about whether it belongs to an artifact family.
```{r}
tmp_df <- df_te_list_agg_family %>%
  ungroup() %>% 
  filter(species == "astCal") %>% 
  mutate(artifact = if_else(repmask_id %in% tmp_df_removed$repmask_id, TRUE, FALSE) ) %>% 
  select(repmask_id, artifact)
df_te_list <- left_join(df_te_list_raw, tmp_df, by = 'repmask_id') 
rm(tmp_df)
```

Print new statistics.
```{r}
glue("Keeps {x} out of {y} entries (across all species).",
  x = df_te_list %>% filter(artifact == FALSE) %>% nrow(),
  y = df_te_list %>% nrow()
)

# within calliptera
glue("Keeps {x} out of {y} (within calliptera).",
  x = df_te_list %>% filter(artifact == FALSE, species == 'astCal') %>% nrow(),
  y = df_te_list %>% filter(species == 'astCal') %>% nrow()
)
```


## Percentage sequence --- subclass level
The overall distributions at the class level have already been generated in `compare_ownVScalip_RepModel.Rmd`, but here are some plots at the subclass level.
```{r subclass_perc_seq, fig.width=12, fig.height=8}
# get percentage sequence per class and subclass
df_te_list_agg_subclass %>% 
  filter(repeat_class != 'Unknown') %>% 
  ggplot(aes(x = perc_sum, y = name, fill = repeat_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(.~species, nrow=1) +
  scale_x_continuous(
    'Percentage sequence of genome 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')
```

```{r subclass_perc_seq_subset, fig.width=5.5, fig.height=8}
# get percentage sequence per class and subclass
df_te_list_agg_subclass %>% 
  filter(species %in% species_subset, repeat_class != 'Unknown') %>% 
  ggplot(aes(x = perc_sum, y = name, fill = repeat_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(.~species, nrow=1) +
  scale_x_continuous(
    'Percentage sequence of genome 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}
# get percentage sequence per class and subclass
df_te_list_agg_subclass %>% 
  filter(species == "astCal", repeat_class != 'Unknown') %>% 
  ggplot(aes(x = perc_sum, y = name, fill = repeat_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(.~species, nrow=1) +
  scale_x_continuous(
    'Percentage sequence',
    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')
```

## Fragment length distribution --- subclass level
Overly small fragments might not be very interesting. An arbitrary but reasonable cutoff might be 50, as that is the `-L` parameter in minigraph, below which we would discard said TE fragments. As illustrated in the boxplots below, the IQR range of almost all families across all the studied species would survive this cutoff.
```{r boxplot_te_length, fig.width=12, fig.height=8}
df_te_list %>%
  ggplot(aes(x=name, y=width, fill=repeat_class)) +
  #geom_hline(yintercept=50, linetype='dashed', color='gray60') +
  geom_boxplot() +
  facet_wrap(.~species, nrow=1) +
  scale_x_discrete(NULL, limits=rev) +
  scale_y_log10("length of TE fragment (bp)", breaks = c(10,100,1000,10000) ) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  coord_flip() +
  theme(legend.position='none')
```

```{r boxplot_te_length_subset, fig.width=5.5, fig.height=8}
df_te_list %>%
  filter(species %in% species_subset) %>% 
  ggplot(aes(x=name, y=width, fill=repeat_class)) +
  #geom_hline(yintercept=50, linetype='dashed', color='gray60') +
  geom_boxplot() +
  facet_wrap(.~species, nrow=1) +
  scale_x_discrete(NULL, limits=rev) +
  scale_y_log10("length of TE fragment (bp)", breaks = c(10,100,1000,10000) ) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  coord_flip() +
  theme(legend.position='none')
```

```{r boxplot_te_length_subset_astCal, fig.width=3, fig.height=8}
df_te_list %>%
  filter(species == "astCal") %>% 
  ggplot(aes(x=name, y=width, fill=repeat_class)) +
  #geom_hline(yintercept=50, linetype='dashed', color='gray60') +
  geom_boxplot() +
  facet_wrap(.~species, nrow=1) +
  scale_x_discrete(NULL, limits=rev) +
  scale_y_log10("length of TE fragment (bp)", breaks = c(10,100,1000,10000) ) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  coord_flip() +
  theme(legend.position='none')
```

As well, overly divergent fragments might be false positives. Based on the boxplots below, the IQR range of almost all families across all the studied species would survive a cutoff of 25 to 30%.
```{r boxplot_te_divergence, fig.width=12, fig.height=8}
df_te_list %>%
  ggplot(aes(x=name, y=score, fill=repeat_class)) +
  #geom_hline(yintercept=30, linetype='dashed', color='gray60') +
  geom_boxplot() +
  facet_wrap(.~species, nrow=1) +
  scale_x_discrete(NULL, limits=rev) +
  scale_y_continuous("Divergence from inferred ancestral consensus (%)") +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  coord_flip() +
  theme(legend.position='none')
```

```{r boxplot_te_divergence_subset, fig.width=5.5, fig.height=8}
df_te_list %>%
  filter(species %in% species_subset) %>% 
  ggplot(aes(x=name, y=score, fill=repeat_class)) +
  #geom_hline(yintercept=30, linetype='dashed', color='gray60') +
  geom_boxplot() +
  facet_wrap(.~species, nrow=1) +
  scale_x_discrete(NULL, limits=rev) +
  scale_y_continuous("Divergence from inferred ancestral consensus (%)") +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  coord_flip() +
  theme(legend.position='none')
```

```{r boxplot_te_divergence_subset_astCal, fig.width=3, fig.height=8}
df_te_list %>%
  filter(species == "astCal") %>% 
  ggplot(aes(x=name, y=score, fill=repeat_class)) +
  #geom_hline(yintercept=30, linetype='dashed', color='gray60') +
  geom_boxplot() +
  facet_wrap(.~species, nrow=1) +
  scale_x_discrete(NULL, limits=rev) +
  scale_y_continuous("Divergence (%)") +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  coord_flip() +
  theme(legend.position='none')
```

## Maximum TE fragment by family
```{r max_te_width, fig.width=5, fig.height=8}
df_te_list_agg_family %>% 
  filter(species == "astCal") %>% 
  ggplot(aes(x=width_max, y=name, color=repeat_class)) +
  geom_point(position=position_jitter(w=0, h=0.2), size=1, alpha=0.6) +
  geom_boxplot(color='black', alpha=0, outlier.shape=NA) +
  scale_y_discrete(limits=rev) +
  scale_color_manual(values = color_schema_te$color_dark, breaks = color_schema_te$te_family) +
  xlab("Maximum fragment length in TE family") +
  ylab(NULL) +
  theme_minimal() +
  theme(legend.position = 'none')
```


## Fragment length and divergence distribution
We visualise the transposon landscape in terms of the **fragment length** and **divergence** from their inferred ancestral consensus. Compared to the version with individual-specific libraries, there is much less noise and that most of the assemblies have very similar TE composition. This shows that comparing TE annotations across different libraries is very difficult.
```{r landscape_single_species}
for (j in species_subset) {
  p <- df_te_list %>% 
    filter(species == {{j}} ) %>% 
    ggplot(aes(width, fill = cut(score, breaks=c(0,10,20,30,50), include.lowest = TRUE))) +
    geom_histogram(position='stack', color='black', alpha = 0.65, binwidth=0.1) +
    facet_wrap(.~factor(repeat_class, levels=color_schema_te$te_family), ncol = 4, scales = 'free_y') + 
    coord_cartesian(xlim=c(10,10000)) +
    scale_fill_brewer('divergence', palette = 'RdBu',
      labels = c('0-10%', '10-20%', '20-30%', '>30%')) +
    scale_x_log10("length of TE fragment (bp)", breaks = c(10,100,1000,10000),
      labels = trans_format("log10", math_format(10^.x)) ) +
    ylab('Number of TE fragments') +
    ggtitle(glue(j)) +
    theme(legend.position='top', legend.key.size=unit(0.025, "npc"),
        legend.title = element_text(size=10), legend.text = element_text(size=9)) 
  print(p)
}
```

```{r landscape_te_length, fig.width=11, fig.height=3}
df_te_list %>% 
  left_join(df_genome_sizes) %>% 
  ggplot(aes(x=width, y=1/genome_size*10**6, color=species)) +
  stat_summary_bin(geom='line', fun='sum', position='identity', breaks=sapply(seq(1,5,0.05), function(x) 10**x), alpha=0.7) + 
  facet_wrap(.~factor(repeat_class, levels=color_schema_te$te_family), nrow = 1, scales = 'free_y') + 
  coord_cartesian(xlim=c(10,10000)) +
  scale_color_manual(values=color_schema_species$color, breaks=color_schema_species$species_abrv) +
  scale_x_log10("length of TE fragment (bp)", breaks = sapply(seq(1,5), function(x) 10**x),
    labels = trans_format("log10", math_format(10^.x)) ) +
  ylab('Number of TE fragments\nper 1Mbp of genome') +
  theme(legend.position='top', legend.key.size=unit(0.025, "npc"),
        legend.title = element_text(size=10), legend.text = element_text(size=9)) 
```

```{r landscape_te_age, fig.width=11, fig.height=3}
df_te_list %>% 
  left_join(df_genome_sizes) %>% 
  filter(repmask_id != "rnd-1_family-37") %>% 
  ggplot(aes(x=score, y=1/genome_size*10**6, color=species)) +
  stat_summary_bin(geom='line', fun='sum', position='identity', breaks = seq(0, 49, 0.5), alpha=0.7) +
  facet_wrap(.~factor(repeat_class, levels=color_schema_te$te_family), nrow = 1, scales = 'free_y') + 
  scale_color_manual(values=color_schema_species$color, breaks=color_schema_species$species_abrv) +
  scale_x_continuous("Divergence from inferred ancestral consensus (%)") +
  ylab('Number of TE fragments\nper 1Mbp of genome') +
  theme(legend.position='top', legend.key.size=unit(0.025, "npc"),
        legend.title = element_text(size=10), legend.text = element_text(size=9)) 
```

These plots suggest that at the genome wide level, the TE distribution is extremely similar, which is expected given that it is estimated that 90% about the genomes are similar. This should be compared to the parallel plot in `further_explore_ownLibraryTE.Rmd` where individual species are plotted with their own library.

In conclusion, in addition to removing artefactual TE families, it might be worth removing TE fragments shorter than 50bp and older than 25% divergence. This may help give a cleaner signal later on.

## Rainbow plot
Here is the rainbow plot for `astCal` only.
```{r rainbow_plot, fig.height=6, fig.width=8}
df_te_list %>% 
  filter(species == "astCal") %>% 
  filter(repmask_id != "rnd-1_family-37") %>% 
  left_join(df_genome_sizes) %>% 
  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, 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, 10.185, 17.5), linetype="dotted") +
  xlab("Divergence from ancestral consensus (%)") + 
  ylab("Percentage sequence") +
  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.6))
```


