---
output:
  html_document:
    toc: true
    self_contained: yes
    toc_float: true
    toc_depth: 4
    theme: flatly
    highlight: tango
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('grey'= "#999999", 'light_orange'="#E69F00",
         'light_blue'="#56B4E9", 'green' = "#009E73",
         'yellow' = "#F0E442", 'blue'= "#0072B2", 'orange'="#D55E00",
         'pink'="#CC79A7")

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 spatial units

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

```{r plot_histogram_of_metrics, echo = F, fig.height=8, fig.width=10, eval=!data_empty}
library(scales)

# name metric lookup
metric_names <- list('umis' = 'UMIs', 'genes' = 'genes', 'pcr' = 'reads / UMIs',
                     'reads'='reads', 'pct_counts_mt' = 'mitochondrial counts')

y_labels <- paste0(c('# of', '# of', '# of', '# of','% of'), ' ', metric_names)
names(y_labels) <- names(metric_names)

to_plot <- obs_df %>%
    select(cell_bc, n_reads, total_counts, pct_counts_mt, n_genes_by_counts) %>%
    mutate(pcr = n_reads / total_counts) %>%
    dplyr::rename(umis=total_counts,
                  reads = n_reads,
                  genes=n_genes_by_counts) %>%
    gather('obs', 'val', -cell_bc)

metric_plot <- function(metric, legend_pos='none'){
    y_label <- y_labels[[metric]]
    to_plot <- to_plot %>%
        filter(obs == metric)

    # calculate minimum difference for a given observation
    min_difference <- to_plot %>%
        filter(!is.na(val)) %>%
        summarise(min_val = min(val,na.rm=T ),
                  max_val = max(val,na.rm=T),
                  difference = max_val - min_val) %>%
        summarise(min_difference = min(difference)) %$%
        min_difference

    pl <-  to_plot %>%
       ggplot(aes(x = val, fill=obs)) +
            geom_histogram(bins=100) +
            scale_fill_manual(values=clrs) +
            labs(fill='', x=y_label,
                 y=paste0('# of\n',
                          'spatial units')) +
            theme(legend.position = legend_pos, strip.background=element_rect(fill='white'),
                  text = element_text(size =18 ), plot.margin = margin(10, 10, 10, 10))

    if(min_difference >= 100){
        pl <- pl + 
            scale_x_log10(breaks = c(1, 1e1, 1e2, 1e3, 1e4, 1e5, 1e6),
                          labels=trans_format('log10',math_format(10^.x))) +
            annotation_logticks(sides='b')
    }

    return(pl)
}

# plot the metrics
pl1 = metric_plot('reads')
pl2 = metric_plot('pcr')
pl3 = metric_plot('genes')
pl4 = metric_plot('umis')
pl5 = metric_plot('pct_counts_mt')

print(plot_grid(pl1, pl2, pl3,
        pl4, pl5, align='vh', hjust=-1, labels="", ncol=2))
```


```{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 plot_umi_2d_title, results = 'asis', echo = F}
if(is_spatial){
    cat('\n\n### Distribution of UMIs in 2D{.tabset}\n\n')
    cat(paste0('Original data = original UMI counts\n\n',
               'Scaled data = the top 10% of UMIs (ordered high to low) are set to the ',
               'minimum UMI of the first 10% UMIs, for visualisation purposes.'))
}
```

```{r plot_umi_2d_function, 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
is_hexagonal <- FALSE
# commented out for now until ggplot corrects the geom_hex or until
# we migrate to python-only reports.

plot_umi_2d <- function(top_cutoff = FALSE){
    if(top_cutoff){
        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))
    } else {
        limits <- c(0, max(obs_df$total_counts + 1))
    }


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

    px_by_um <- snakemake@params$pbf_metrics$px_by_um
    puck_width_um <- (x_limits[2] - x_limits[1]) / px_by_um

    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')

    x_breaks <- seq(x_limits[1], x_limits[2], px_by_um * mm_dist)
    y_breaks <- seq(y_limits[1], y_limits[2], px_by_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

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

    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_umi_2d, echo = F, results = 'asis'}
if(is_spatial){
    cat('\n\n#### Original data\n\n')
    print(plot_umi_2d(top_cutoff = FALSE))
    cat('\n\n#### Scaled data\n\n')
    print(plot_umi_2d(top_cutoff = TRUE))

}
```


```{r umap_text, eval=data_complete, echo =F, results='asis'}
cat('\n\n### Automated analysis with different clustering resolutions {.tabset}\n\n')
```

```{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())

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

if(is_spatial){
    nhood_dat <- tryCatch({
        read_csv(snakemake@input$nhood_enrichment)},
        error = function(e) {
            message("File not found or could not be read: ", e$message)
            return(NULL)  # Return NULL or an empty tibble/dataframe if needed
            }
    )
}

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) {
        # calculate breaks and limits for the puck
        x_limits <- obs_df$x_pos %>% {c(min(.), max(.))}
        y_limits <- obs_df$y_pos %>% {c(min(.), max(.))}

        px_by_um <- snakemake@params$pbf_metrics$px_by_um
        puck_width_um <- (x_limits[2] - x_limits[1]) / px_by_um

        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')

        x_breaks <- seq(x_limits[1], x_limits[2], px_by_um * mm_dist)
        y_breaks <- seq(y_limits[1], y_limits[2], px_by_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))

        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))

        if (!is.null(nhood_dat)) {
            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)

        nhood_plot <- 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(umap_plot)
        print(physical_plot)
        
        tryCatch({
            print(nhood_dat)},
            error = function(e) {
                pass}
        )
                

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