---
output:
  html_document:
    toc: true
    self_contained: yes
version: 0.1.1
author: Tamas Ryszard Sztanka-Toth, Nikolaos Karaiskos
email: tamasryszard.sztanka-toth@mdc-berlin.de, nikolaos.karaiskos@mdc.berlin.de
license: GPL
---

```{r knitr_options, include=FALSE, cache=FALSE}
knitr::opts_chunk$set(
  autodep = TRUE,
  message = FALSE,
  warning = FALSE,
  comment = NA
)
```

```{r libraries, include = F, cache=F}
library(tidyverse)
library(yaml)
library(knitr)
library(magrittr)
library(kableExtra)
library(cowplot)

theme_set(theme_cowplot(18))

cpalette <- list('orange' = '#D55E00', 'blue' = '#0072B2', 'green' = '#009E73', 'black' = '#000000', 'yellow' = '#F0E442', 
				 'grey' = '#999999', 'light_orange' = "#E69F00", 'light_blue' = "#56B4E9")

clrs = c('umis'=cpalette$light_orange, 'genes' = cpalette$light_blue,
         'reads'=cpalette$green, 'pcr'=cpalette$pink, 'pct_counts_mt'= 'black')

source(snakemake@params$r_shared_scripts)
```

```{r read_data, echo = F}
obs_df <- read_csv(snakemake@input$obs_df) %>%
    mutate(pct_counts_mt = as.double(pct_counts_mt))
var_df <- read_csv(snakemake@input$var_df)

n_beads = nrow(obs_df)
data_empty <- n_beads == 0

if (!data_empty){
    data_empty <- sum(obs_df$total_counts) < 1
}
```

### Overview

```{r create_overview_table, echo = F}
puck_width_um <- snakemake@params$puck_variables$width_um
parameter_stats <- rbind(c('UMI filter', snakemake@wildcards$umi_cutoff),
          c('# genes in data', nrow(var_df)),
          c('# of spots in data', obs_df %>% select(cell_bc) %>% unique %>% nrow),
          c('median UMI', median(obs_df$total_counts)),
          c('median genes', median(obs_df$n_genes_by_counts)),
          c('puck width (um)', puck_width_um))

parameter_stats %>%
    kbl(col.names=NULL) %>%
    kable_classic_2(full_width=F) %>%
    #pack_rows('Sample info', 1, 7) %>%
    pack_rows('Data info', 1, 5)
```

### Histogram of metrics over beads

```{r data_empty, echo = F, eval = data_empty, results='asis'}
cat('This dataset has 0 beads passing the filters')
```

```{r plot_histogram_of_metrics, echo = F, fig.width=10, eval=!data_empty}
pl1 <- obs_df %>%
    select(cell_bc, total_counts, n_genes_by_counts) %>%
    dplyr::rename(umis=total_counts,
                  genes=n_genes_by_counts) %>%
    gather('obs', 'value', -cell_bc) %>%
    ggplot(aes(value, fill=obs)) +
        geom_histogram(bins=100) +
        scale_fill_manual(values=clrs) +
        facet_wrap(~obs, ncol=2, scales='free')

if (length(unique(obs_df$total_counts)) > 50){
    pl1 <- pl1 +
        scale_x_log10() +
        annotation_logticks(sides='b') +
        labs(x='log(value)', y='count')
} else {
    pl1 <- pl1 +
        labs(x ='value', y='count')
}

pl2 <- obs_df %>%
    select(cell_bc, pct_counts_mt) %>%
    gather('obs', 'value', -cell_bc) %>%
    ggplot(aes(value, fill=obs)) +
        geom_histogram(bins=100) +
        scale_fill_manual(values=clrs) +
        facet_wrap(~obs, ncol=2, scales='free')+
        labs(x='value', y='count')

plot_grid(pl1,
          plot_grid(pl2, NULL, ncol=2, rel_widths=c(1.5,1)),
          ncol=1, rel_heights=c(1,1), labels="")
```


```{r check_completeness, echo = F, fig.size=10, fig.width=10}
obs_df <- obs_df %>%
    gather('res', 'cluster', starts_with('leiden'))

data_complete = F

if ('cluster' %in% colnames(obs_df)){
    data_complete = T
}

is_spatial <- snakemake@params$is_spatial & data_complete

fig_height=6
fig_width=7
```

```{r incompleteness_warning, eval=!data_complete, echo=F, results='asis'}
cat('### WARNING: data incomplete\n\n')
cat(paste0('This dataset has ', n_beads, ' cells passing the filter of ', snakemake@wildcards$umi_cutoff, ' and ', nrow(var_df), ' genes.\n\n'))
cat(paste0('This dataset is too small, so it couldn\'t be properly clustered and analysed automatically'))
```


```{r umap_text, eval=data_complete, echo =F, results='asis'}
cat('### UMAP plots using different resolutions {.tabset}\n\n')
```


```{r plot_umi_umap, eval=data_complete, echo=F, fig.height=fig_height, fig.width=fig_width}
obs_df <- obs_df %>%
    mutate(cluster = factor(cluster)) %>%
    mutate(log1p_total_counts = log2(1+total_counts))

n_cells <- obs_df %>%
    select(cell_bc) %>%
    unique %>%
    nrow()

def_plot_bead_size <- ifelse(n_cells > 5000, 0.4, 0.75)
def_plot_bead_size <- ifelse(n_cells > 10000, 0.1, def_plot_bead_size)
def_plot_bead_size <- ifelse(n_cells > 25000, 0.05, def_plot_bead_size)

obs_colnames <- colnames(obs_df)

# barcode file attached at the python level
is_hexagonal <- snakemake@params$run_mode_variables$mesh_type == 'hexagon' & snakemake@params$run_mode_variables$mesh_data

filter_quant <- quantile(obs_df$total_counts, 0.9, na.rm=T)
limits <- c(0, quantile(obs_df$total_counts, 0.91, na.rm=T))

if (is_spatial){

    # calculate breaks and limits for the puck
    puck_file <- read_csv(snakemake@input$puck_file)
    x_limits <- puck_file$x_pos %>% {c(min(.), max(.))}
    y_limits <- puck_file$y_pos %>% {c(min(.), max(.))}

    ratio <- (x_limits[2] - x_limits[1] ) / (y_limits[2] - y_limits[1])

    scale_factor <- ifelse(puck_width_um < 3000, 2, 3)
    mm_dist <- max(10^scale_factor, round(puck_width_um/3, digits =-scale_factor))
    mm_diff <- mm_dist / 1000

    x_mm_breaks <- seq(0, puck_width_um, mm_dist)
    x_mm_breaks <- paste0(x_mm_breaks * mm_diff / mm_dist, 'mm')
    y_mm_breaks <- seq(0, puck_width_um / ratio, mm_dist)
    y_mm_breaks <- paste0(y_mm_breaks * mm_diff / mm_dist, 'mm')

    px_to_um <- (x_limits[2] - x_limits[1]) / snakemake@params$puck_variables$width_um

    x_breaks <- seq(x_limits[1], x_limits[2], px_to_um * mm_dist)
    y_breaks <- seq(y_limits[1], y_limits[2], px_to_um * mm_dist)

    puck_bead_size <- min(def_plot_bead_size, ifelse(snakemake@params$run_mode_variables$mesh_data,
        snakemake@params$run_mode_variables$mesh_spot_diameter_um / 40,
        snakemake@params$puck_variables$spot_diameter_um / 40))
    
    res_colnames <- obs_df$res %>%
        unique

    umi_pl <- obs_df %>%
        filter(res == res_colnames[1]) %>%
        arrange(total_counts) %>%
        mutate(total_counts = ifelse(total_counts > filter_quant, filter_quant, total_counts))

    if(is_hexagonal){
        umi_pl <- umi_pl %>%
            ggplot(aes(x=x_pos, y=y_pos, fill = total_counts, group=1)) +
                geom_hex(stat='identity', color='gray95') +
                coord_fixed()+
                scale_x_continuous(labels = x_mm_breaks, breaks = x_breaks, limits=x_limits) +
                scale_y_continuous(labels = y_mm_breaks, breaks = y_breaks, limits=y_limits) +
                scale_fill_viridis_c(option =  "magma", limits = limits) +
                guides(fill = guide_colorbar(barheight = 15)) + 
                labs(fill='UMI count', x='', y='')

    }
    else{
        umi_pl <- umi_pl %>%
            ggplot(aes(x_pos, y_pos, color = total_counts)) +
                geom_point(size=puck_bead_size) + 
                coord_fixed()+
                scale_x_continuous(labels = x_mm_breaks, breaks = x_breaks, limits=x_limits) +
                scale_y_continuous(labels = y_mm_breaks, breaks = y_breaks, limits=y_limits) +
                scale_color_viridis_c(option =  "magma", limits = limits) +
                guides(color = guide_colorbar(barheight = 15)) + 
                labs(color='# of UMIs\nper spatial unit', x='', y='')
    }
    umi_pl <- umi_pl + 
        theme(panel.background = element_rect(fill = 'gray95'), 
             legend.spacing = unit(0.1, 'cm'),
            axis.line = element_line(color = 'black'),
            text = element_text(color='black', size=18))
    umi_pl
}

```

```{r plot_clusters_umap_puck, echo =F, fig.height=fig_height, fig.width=fig_width, eval=data_complete, results='asis'}
library(pals)

cluster_clrs <- unname(glasbey())

#top10_marker_table <- read_table2(snakemake@input$cluster_markers)

for (i in obs_df %$% res %>% unique){
    res <- as.double(strsplit(i, '_')[[1]][2])
    cat(paste0('\n\n#### ', res, ' resolution {.tabset}\n\n'))
    dat <- obs_df %>%
        filter(res == i) %>%
        dplyr::select(cell_bc, umap_0, umap_1, cluster)

    umap_plot <- dat %>%
        ggplot(aes(umap_0, umap_1, color = cluster)) +
            geom_point(size=def_plot_bead_size) +
            guides(colour = guide_legend(override.aes = list(size=3)))+
            coord_fixed() +
            theme(axis.text = element_blank(),
                  legend.position = 'none',
                axis.ticks = element_blank(), axis.line = element_line(color='white'))

    n_clusters <- length(unique(dat$cluster))
    if(n_clusters< length(cluster_clrs)){
        umap_plot <- umap_plot + scale_color_manual(values=cluster_clrs)
    }

    if (is_spatial) {
        physical_plot <- obs_df %>%
            filter(res == i) %>%
            dplyr::select(cell_bc, x_pos, y_pos, cluster)

        if(is_hexagonal){
            physical_plot <- physical_plot %>%
            ggplot(aes(x_pos, y_pos, fill = cluster)) +
                geom_hex(stat='identity', color='gray95') +
                guides(fill = guide_legend(override.aes = list(size=3), ncol=2))+
                coord_fixed()+
                labs(x='', y='') +
                scale_x_continuous(labels = x_mm_breaks, breaks = x_breaks, limits=x_limits) +
                scale_y_continuous(labels = y_mm_breaks, breaks = y_breaks, limits=y_limits)
            if(n_clusters< length(cluster_clrs)){
                physical_plot <- physical_plot + scale_fill_manual(values=cluster_clrs)
            }
        } else{
            physical_plot <- physical_plot %>%
            ggplot(aes(x_pos, y_pos, color = cluster)) +
                geom_point(size=puck_bead_size) +
                guides(colour = guide_legend(override.aes = list(size=3), ncol=2))+
                coord_fixed()+
                labs(x='', y='') +
                scale_x_continuous(labels = x_mm_breaks, breaks = x_breaks, limits=x_limits) +
                scale_y_continuous(labels = y_mm_breaks, breaks = y_breaks, limits=y_limits)
            if(n_clusters< length(cluster_clrs)){
                physical_plot <- physical_plot + scale_color_manual(values=cluster_clrs)
            }

        }
        physical_plot <- physical_plot +
            theme(panel.background = element_rect(fill = 'gray95'),
                legend.spacing = unit(0.1, 'cm'),
                axis.line = element_line(color = 'black'),
                plot.subtitle=element_text(size=18),
                text = element_text(color='black')) +
            ggtitle('', subtitle=paste0('resolution = ', res))

        print(umap_plot)
        print(physical_plot)
    } else{
        print(umap_plot)
    }
}
```

```{r nhood_enrichment_title, echo =F, eval=is_spatial, results='asis'}
cat('\n\n### Neighborhood enrichment\n\n')
```

```{r plot_nhood_enrichment, echo =F, eval=is_spatial}
nhood_dat <- read_csv(snakemake@input$nhood_enrichment)

for (i in obs_df %$% res %>% unique){
    res <- as.double(strsplit(i, '_')[[1]][2])
    dat <- nhood_dat %>%
        filter(resolution == res) %>%
        mutate(a = as.character(1000 + cluster_a),
               b = as.character(1000 + cluster_b)) %>%
        mutate(zscore = ifelse(zscore > 100, 100, zscore),
               zscore = ifelse(zscore < -50, -50, zscore))

    labs <- dat %>%
        select(cluster_a, cluster_b, a, b) %>%
        filter(cluster_a %% 5 == 0, cluster_b %% 5 == 0)

    pl <- dat %>%
        ggplot(aes(a, b, fill=zscore)) +
            geom_tile(color='white') + 
            scale_fill_viridis_c(option='inferno', limits = c(-51, 101)) +
        guides(fill = guide_colorbar(barheight = 12)) + 
        scale_x_discrete(labels = labs$cluster_a, breaks=labs$a) +
        scale_y_discrete(labels = labs$cluster_b, breaks=labs$b) +
        coord_fixed() +
        labs(fill='neighborhood\nenrichment\nscore', x='cluster identity', y='cluster identity') +
        ggtitle('', subtitle=paste0('resolution = ', res))+
        coord_equal()+
        theme(plot.subtitle=element_text(size=18))
    print(pl)
}
```

