library(readr)
library(dplyr)
library(tidyr)
library(ggplot2)
library(glue)

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

df_raw <- read_csv(snakemake@input[[1]]) %>% 
  mutate(species1 = forcats::fct_relevel(species1, color_schema_species$species_abrv),
         species2 = forcats::fct_relevel(species2, color_schema_species$species_abrv),
         bb_sv_perc = 100 - bb_detect_perc
  )

diagonal_df <- data.frame(
  species1 = color_schema_species$species_abrv, 
  species2 = color_schema_species$species_abrv, 
  bb_detect_perc = 1,
  bb_sv_perc = 0
)

# loop over several variant sizes
for (i in seq(1, length(snakemake@config$min_variant_sizes)) ) {
  L <- snakemake@config$min_variant_sizes[i]
  conserv_df <- df_raw %>% 
    filter(min_variant_size == L) %>% 
    select(species1, species2, bb_detect_perc, bb_sv_perc)
  
  p <- ggplot(conserv_df, aes(species1, species2, label = format(round(bb_sv_perc, 2), nsmall=2), fill = bb_sv_perc)) +
    geom_tile(color='black') +
    geom_text(color='black') +
    scale_fill_distiller('flexible\npercentage\non backbone\nvs query', palette='RdBu') +
    geom_tile(aes(x = species1, y = species2), diagonal_df, fill = 'black') +
    labs(x='backbone', y = 'query') +
    theme_classic() +
    theme(
      aspect.ratio = 1,
      axis.line = element_blank(), 
      axis.ticks = element_blank(),
      axis.text = element_text(size=10)
    )
  ggsave(snakemake@output$heatmap[i], width = 6.2, height = 4.8)
}

p <- df_raw %>% 
  ggplot(aes(x = min_variant_size, y = bb_sv_perc, color = species2)) +
  geom_point(size=1.2) + 
  geom_line(size=.5) +
  facet_wrap(.~species1, nrow=2) +
  scale_color_manual('query genome', values=color_schema_species$color, breaks=color_schema_species$species_abrv) +
  theme_classic() +
  xlab('minimum variant size (-L)') +
  ylab('flexible percentage on backbone vs query') +
  theme(legend.position='top', 
        panel.grid.major.y = element_line(colour = "grey90", size=0.5),
        legend.title = element_text(size = rel(0.75)),
        legend.text = element_text(size = rel(0.65)),
        legend.key.size=unit(0.04, "npc")
  )
ggsave(snakemake@output$robustness, width = 6.4, height = 4)
