shell:
    """umi_tools whitelist\
    --stdin {input}\
    --bc-pattern='(?P<cell_1>.{{{params.cell_barcode_length}}})(?P<umi_1>.{{{params.umi_barcode_length}}})'\
    --extract-method=regex\
    --set-cell-number={params.num_cells}\
    --log2stderr > {output}"""

wrapper:
    '0.36.0/bio/fastqc'

wrapper:
    '0.36.0/bio/multiqc'

debug_flag <- FALSE
# check if DEBUG flag is set
if (snakemake@config$DEBUG) {
  debug_flag <- TRUE
  message("In debug mode: saving R objects to inspect later")
  path_debug <- file.path(snakemake@config$LOCAL$results, "debug")
  dir.create(path_debug, showWarnings = FALSE)
  save(snakemake, file = file.path(path_debug, "create_summary_stats_snakemake.rdata"))
}
#### /debug

library(dplyr) # Dataframe manipulation
library(Matrix) # Sparse matrices
library(stringr)
library(RColorBrewer)
library(devtools)
library(Seurat)
library(plotly)

samples <- snakemake@params$sample_names
batches <- snakemake@params$batches


# importing Seurat object

seuratobj <- readRDS(file  = file.path(snakemake@input$R_objects))
meta.data <- seuratobj@meta.data

# subset only highest stamps as set in config.yaml
# this is necessary since there are more stamps selected as a safty margin which now have to be taken out again to calculate stats.

meta.data.sub <- meta.data %>%
  group_by(orig.ident) %>%
  arrange(desc(nCounts)) %>%
  slice(1:expected_cells[1]) %>% # makes sure only # expected cell are kept
  as.data.frame()

gini_index <- function (x, weights = rep(1, length = length(x))) {
    ox      <- order(x)
    x       <- x[ox]
    weights <- weights[ox] / sum(weights)
    p       <- cumsum(weights)
    nu      <- cumsum(weights * x)
    n       <- length(nu)
    nu      <- nu/nu[n]
    sum(nu[-1] * p[-n]) - sum(nu[-n] * p[-1])
}

#------------------------------------ post-filter-stats
# stats only based after keeping most abundant `expected-cells`
# taken out from seurat object generated in violine_plot rule.

# median calculator

stats_post <- meta.data.sub %>%
  group_by(orig.ident) %>%
  summarise(
    Total_nb_reads                 = sum(nCounts),
    Nb_STAMPS                      = mean(expected_cells), # should be all the same anyway..
    Median_reads_per_STAMP         = round(median(nCounts), 2),
    Mean_reads_per_STAMP           = round(mean(nCounts), 2),
    Total_nb_UMIs                  = sum(nUMI),
    Median_UMIs_per_STAMP          = round(median(nUMI), 2),
    Mean_UMIs_per_STAMP            = round(mean(nUMI), 2),
    Mean_UMIs_per_Gene             = round(mean(umi.per.gene), 2),
    Median_number_genes_per_STAMP  = round(median(nGene), 2),
    Mean_number_genes_per_STAMP    = round(mean(nGene), 2),
    Mean_Ribo_pct                  = round(100 * mean(pct.Ribo), 2),
    Mean_Mito_pct                  = round(100 * mean(pct.mito), 2),
    Mean_Count_per_UMI             = round(sum(nCounts) / sum(nUMI), 2),
    Read_length                    = mean(read_length), # should be all the same anyway..
    Number_barcodes_used_for_debug = n()
  ) %>%
  as.data.frame()

row.names(stats_post) <- stats_post$orig.ident



# highest, lowest count/UMI Stamp
# pre STAMP stats

# hist out goes into knee plots
# 'results/logs/{sample}_hist_out_cell.txt'
# """export _JAVA_OPTIONS=-Djava.io.tmpdir={params.temp_directory} && BAMTagHistogram -m {params.memory}\
# TAG=XC\
# https://hpc.nih.gov/apps/dropseq.html
# there is not hint in documentation on any filtering (only read quality)

#------------------------------------ pre-filter-stats
# calculating statistics based on barcodes before thresholding it (i.e. keeping the most abudant barcodes based on `expected cells`)
# This is based on 'logs/{sample}_hist_out_cell.txt' generated by `BAMTagHistogram` from dropseq-tools
# TOOO: The documentation only mentions duplicate and quality filter. But it seems do more filtering since most barcodes are expected to have only one read assinged but there are usually more. Find out.
# https://hpc.nih.gov/apps/dropseq.html

# TODO: Nr UMI for pre filter

stats_pre <- data.frame(matrix(nrow = length(samples), ncol = 10))
colnames(stats_pre) <- c(
  "Sample",
  "Batch",
  "Total_raw_reads",
  "Nr_barcodes_total",
  "Nr_barcodes_more_than_1_reads",
  "Nr_barcodes_more_than_10_reads",
  "percentile99",
  "percentile95",
  "percentile50",
  "Gini-index"
)

stats_pre[, "Sample"] <- samples
stats_pre[, "Batch"]  <- batches


for (i in 1:length(samples)) {
  # importing 'logs/{sample}_hist_out_cell.txt'
  hist_out <- read.table(
    file = snakemake@input$hist_cell[i],
    header = FALSE, stringsAsFactors = FALSE
  )
  mysample <- samples[i]
  reads <- hist_out$V1
  barcodes <- hist_out$V2
  # calculations on reads
  # total reads are not sum(reads)! Needs to be taken from
  # results/logs/cutadapt/sample1_R1.qc.txt
  #
  # read in full text file "sample_R2.qc.txt"
  filedump <- readLines(snakemake@input$R2qc[i])
  # subset line matching a pattern
  total_reads <- filedump[grep("Total reads processed:", filedump)] %>% #extract line
    str_extract("[0-9,]+") %>% # extract number from line
    str_replace_all(",", "") %>% # delete comma for subsequent numeric casting
    as.numeric()
  reads_cumsum      <- cumsum(reads)
  reads_cumsum_perc <- (reads_cumsum / sum(reads))
  # reporting stats
  stats_pre[i, "Total_raw_reads"]                <- total_reads
  stats_pre[i, "Reads_assigned_to_expected_STAMPs"] <- sum(reads[1:stats_post$Nb_STAMPS[i]])
  stats_pre[i, "Nr_barcodes_total"]              <- length(barcodes)
  stats_pre[i, "percentile99"]                   <- which.min(reads_cumsum_perc < 0.99)
  stats_pre[i, "percentile95"]                   <- which.min(reads_cumsum_perc < 0.95)
  stats_pre[i, "percentile50"]                   <- which.min(reads_cumsum_perc < 0.50)
  stats_pre[i, "Nr_barcodes_more_than_1_reads"]  <- sum(reads > 1)
  stats_pre[i, "Nr_barcodes_more_than_10_reads"] <- sum(reads > 10)
  stats_pre[i, "Gini-index"]                     <- round(gini_index(reads), 2)
  expected_cells <- as.numeric(filter(stats_post, orig.ident==mysample) %>% select(Nb_STAMPS))
  # % of reads left after applying expected_cells cuttoff
  stats_post[mysample, "Pct_reads_after_filter_expected_cells"] <-
    round(100 * (
                 reads_cumsum[expected_cells] / sum(reads)
    ), 2)
  # % of reads left after applying all filters including mapping etc.
  # Thats the effecive usable reads of the sequencing run
  stats_post[mysample, "Pct_reads_after_filter_everything"] <-
    round(100 * (
                 filter(stats_post, orig.ident==mysample) %>% select(Total_nb_reads) /
                 stats_pre [i, "Total_raw_reads"]), 2
         )
}

stats_pre <- stats_pre %>%
  arrange(Sample)

stats_post <- stats_post %>%
  arrange((orig.ident))

# output
write.csv(stats_pre,  file.path(snakemake@output$stats_pre))
write.csv(stats_post, file.path(snakemake@output$stats_post)) # writes table for excel

if (debug_flag) {
  save.image(file = file.path(path_debug, "create_summary_stats_workspace.rdata"))
}

from Bio import SeqIO
import gzip
from collections import Counter

fastq_parser = SeqIO.parse(gzip.open(snakemake.input.R1, "rt"), "fastq")
sequences = []
n=0
for fastq_R1 in fastq_parser:
		sequences.append(str(fastq_R1.seq))
		n+=1
		if(n==10000000):
			break
def parse_barcodes(fastq_parser):
	counts={}
	ranges = range(5,len(sequences[0]))

	for cell_bc_length in ranges:
		counts[cell_bc_length] = list()
		for fastq_R1 in sequences:
			counts[cell_bc_length].append(fastq_R1[0:cell_bc_length])
	return(counts)



counts = parse_barcodes(fastq_parser)

with open(snakemake.output[0], "w") as outfile:
	outfile.write('bc_length,first_counts\n')
	for cell_bc_length in counts:
		outfile.write('{},{}\n'.format(cell_bc_length,str(Counter(counts[cell_bc_length]).most_common(100)[0][1])))

import pysam
import re
import csv
from Bio import SeqIO
import gzip
from collections import defaultdict
import sys

#This function fills in a dict with readids
#and their corresponding cell and umi barcodes until it finds
#a specific read id

discard_secondary_alignements = snakemake.params['discard_secondary_alignements']

barcodes_struct = {
	'BC_start':snakemake.params['BC_start'],
	'BC_end':snakemake.params['BC_end'],
	'UMI_start':snakemake.params['UMI_start'],
	'UMI_end':snakemake.params['UMI_end']
	}

def parse_barcodes(fastq_parser, query_name, read_barcodes, barcodes_struct):
	for fastq_R1 in fastq_parser:
		# Some sequencers give a /1 and /2 to R1 and R2 read ids respectively. This attempts to solve the issue #69.
		if '/' in fastq_R1.id:
			R1_id = fastq_R1.id[:fastq_R1.id.find("/")]
		else:
			R1_id = fastq_R1.id
		read_barcodes[R1_id]['XC'] = str(fastq_R1.seq)[barcodes_struct['BC_start']:barcodes_struct['BC_end']]
		read_barcodes[R1_id]['XM'] = str(fastq_R1.seq)[barcodes_struct['UMI_start']:barcodes_struct['UMI_end']]
		if(read_barcodes[R1_id]['XM']==''):
			sys.SystemExit('UMI empty for read {}.\n The barcode is: {}.\nWhole entry is:{}'.format(R1_id, fastq_R1.seq,fastq_R1))
		if (R1_id == query_name):
			return(fastq_parser,read_barcodes)
	return(fastq_parser,read_barcodes)

infile_bam = pysam.AlignmentFile(snakemake.input[0], "rb")

fastq_parser = SeqIO.parse(gzip.open(snakemake.input[1], "rt"), "fastq")

outfile = pysam.AlignmentFile(snakemake.output[0], "wb", template=infile_bam)

read_barcodes = defaultdict(lambda :{'XC':'','XM':''})

for bam_read in infile_bam:
	if(discard_secondary_alignements & bam_read.is_secondary):
		continue
	if (bam_read.query_name) in read_barcodes:
		current_barcodes = read_barcodes.pop(bam_read.query_name)
		tags = bam_read.get_tags()
		tags.extend([
			('XC', current_barcodes['XC'],'Z'),
			('XM', current_barcodes['XM'],'Z')])
		bam_read.set_tags(tags)
	else:
		fastq_parser,read_barcodes = parse_barcodes(fastq_parser, bam_read.query_name, read_barcodes, barcodes_struct)
		if (bam_read.query_name) not in read_barcodes:
			raise SystemExit('Read {} from mapped file is missing in reference fastq file!'.format(bam_read.query_name))
			os.remove(snakemake.output[0])
		current_barcodes = read_barcodes.pop(bam_read.query_name)
		tags = bam_read.get_tags()
		tags.extend([
			('XC', current_barcodes['XC'],'Z'),
			('XM', current_barcodes['XM'],'Z')])
		bam_read.set_tags(tags)
	outfile.write(bam_read)

debug_flag <- FALSE
# check if DEBUG flag is set
if (snakemake@config$DEBUG) {
  debug_flag <- TRUE
  message("In debug mode: saving R objects to inspect later")
  path_debug <- file.path(snakemake@config$LOCAL$results, "debug")
  dir.create(path_debug, showWarnings = FALSE)
  save(snakemake, file = file.path(path_debug, "plot_adapter_content_snakemake.rdata"))
}

#------------------------------------ debugging

library(ggplot2)
library(dplyr)
library(viridis)

samples <- snakemake@params$sample_names
batches <- snakemake@params$batches

#Read files into a list
cutadapt_clean_list <- list()
for (i in seq_along(samples)){
  cutadapt_clean           <- read.csv(snakemake@input[[i]][1], header = TRUE)
  cutadapt_clean$Sample    <- samples[i]
  cutadapt_clean$Batch     <- batches[i]
  cutadapt_clean_list[[i]] <- cutadapt_clean
}

# combining adaptors accross samples
cutadapt_counts <- Reduce(rbind, cutadapt_clean_list, NULL)

#Transform it into percentages
cutadapt_counts <- group_by(cutadapt_counts, Sample, Pair) %>%
  mutate(Percentages=Count/sum(Count))
#      Adapter                           Sequence Pair Count  Sample  Batch
# 1 PrefixNX/1                AGATGTGTATAAGAGACAG   R1     7 sample1 Batch1
# ...
# 6  Trans2_rc CTGTCTCTTATACACATCTCCGAGCCCACGAGAC   R2     5 sample2 Batch2

p1 <- ggplot(cutadapt_counts, aes(x=Sample, y = Percentages, fill = Adapter))  +
  geom_bar(stat = "identity") +
  facet_grid(Pair ~ Batch, scales = "free") +
  theme_minimal() +
  ggtitle("Comparison accross samples of adapter content") +
  scale_x_discrete(label=abbreviate) +
  scale_y_continuous(labels = scales::percent) +
  theme(axis.text.x=element_text(angle = 90, hjust = 0)) +
  scale_fill_viridis(discrete=TRUE)

ggsave(plot=p1, filename=snakemake@output$pdf)

if (debug_flag) {
  save.image(file = file.path(path_debug, "plot_adapter_content_workspace.rdata"))
}

debug_flag <- FALSE
if (snakemake@config$DEBUG) {
  debug_flag <- TRUE
  message("In debug mode: saving R objects to inspect later")
  path_debug <- file.path(snakemake@config$LOCAL$results, "debug")
  dir.create(path_debug, showWarnings = FALSE)
  save(snakemake, file = file.path(path_debug,
     paste0("plot_knee_plot_snakemake_",
            attr(snakemake, "wildcard")$sample, ".rdata"))
          )
}
#### /debug

library(ggplot2)
library(plyr)
# Create the cumulative plot
data=read.table(file = snakemake@input[[1]][1], header=FALSE, stringsAsFactors=FALSE)
barcodes = data$V2
total_reads = sum(data$V1)
y_raw=cumsum(data$V1)
y=(y_raw/total_reads)
x = 1:length(y)
plot_data = data.frame(rank = x,cum_sum=y, Barcode=data$V2)
x_scale = snakemake@params$cells * 4

knee_plot <- ggplot(plot_data, aes(x=rank, y=cum_sum)) +
 geom_point(size = 0.1) +
 xlim(0,x_scale) +
 geom_vline(xintercept=snakemake@params$cells, linetype="dashed", color = "red") +
 ggtitle(paste0(snakemake@wildcards$sample, '\nTotal reads: ', prettyNum(total_reads))) +
 theme(plot.title = element_text(size=10)) +
 labs(x='STAMPS', y='Cumulative fraction of reads') +
 scale_y_continuous(labels = scales::percent) +
 theme_classic()

if(!is.null(snakemake@input$barcodes))
{
  selected_cells <- read.csv(snakemake@input$barcodes, header=FALSE, stringsAsFactors=FALSE)
  knee_plot <- knee_plot +
    geom_point(data = plot_data[plot_data$Barcode %in% selected_cells$V1,],
               aes(x=rank, y=cum_sum, color='Selected'), size=0.1) +
    scale_color_manual(values=c('Selected'='green'))
}
ggsave(knee_plot, file=snakemake@output$pdf, width = 4, height = 3)


if (debug_flag) {
  library(gridExtra)
  library(grid)
  # potential usefull to change lag of diff calculation
  mylag <- 1
  # data <- read.table(file = file, header=FALSE, stringsAsFactors=FALSE)
  # head(data,3)
  #        V1           V2
  # 1 1145137 CCCTTCGTCTGC
  # 2 1039974 ATAGTTTTTTAA
  # 3  912199 GCATGAAACTTC

  # borrowed from https://stackoverflow.com/questions/6836409/finding-local-maxima-and-minima
  localMaxima <- function(x) {
    # Use -Inf instead if x is numeric (non-integer)
    y <- diff(c(-.Machine$integer.max, x)) > 0L
    rle(y)$lengths
    y <- cumsum(rle(y)$lengths)
    y <- y[seq.int(1L, length(y), 2L)]
    if (x[[1]] == x[[2]]) {
      y <- y[-1]
    }
    y
  }

  reads        <- data$V1
  barcodes     <- data$V2
  total_reads  <- sum(reads)
  reads_cumsum <- cumsum(reads)
  # 1st dervivative (diff) needs also to be padded to keep same vector length
  reads_diff   <- c(diff(reads,lag=mylag,differences = 1),rep(0,mylag))

  # 2nd derivative: twice as much padding:
  reads_diff_diff   <- c(diff(reads,lag=mylag,differences = 2),rep(0,mylag*2))
  reads_cumsum_perc <- (reads_cumsum/total_reads)
  x                 <- 1:length(reads_cumsum_perc)
  plot_data         <- data.frame(rank       = x,
                                  cum_sum    = reads_cumsum_perc,
                                  read_count = reads,
                                  Barcode    = data$V2,
                                  diff       = reads_diff,
                                  diffdiff   = reads_diff_diff)

  # head(plot_data,6)
  #   rank     cum_sum read_count      Barcode    diff diffdiff
  # 1    1 0.007192916    1145137 CCCTTCGTCTGC -105163   -22612
  # 2    2 0.013725274    1039974 ATAGTTTTTTAA -127775    14025
  # 3    3 0.019455043     912199 GCATGAAACTTC -113750    61486
  # 4    4 0.024470318     798449 GTGTGGGTCTCT  -52264    34985
  # 5    5 0.029157308     746185 CGTACTGACTAC  -17279   -60441
  # 6    6 0.033735764     728906 GTTCGTCCCGCC  -77720    69104
  mystats <- paste0("| Nr barcodes total: ", length(barcodes), ' \n ',
                    "Nr barcodes for 50% reads: ", which.min(reads_cumsum_perc<0.50)," | ",
                    "Nr barcodes for 95% reads: ", which.min(reads_cumsum_perc<0.95)," | ",
                    "Nr barcodes for 99% reads: ", which.min(reads_cumsum_perc<0.99)
                    )

  # x_scale        <- which(reads_cumsum_perc>0.99)[1]
  x_scale        <- snakemake@params$cells * 4
  plot_data_head <- head(plot_data, x_scale)
  # plot_data_head$reads_diff_smooth <- predict(loess(diff~rank,data=plot_data_head))
  plot_data_head$reads_diffdiff_smooth <- predict(loess(diffdiff~rank,span=0.2,data=plot_data_head))

  ## Finding knee in knee-plot:
  # Best approach so far is to calclate the 2nd derivative of the read counts per STAMP (reads_diff_diff), smooth it (loess) and find it's maxima. Since ther can be several maxima, they are just ploted and it's up to the user to visually asses and decide.
  # finding local maxima in 2nd derivative:
  #https://stackoverflow.com/questions/6836409/finding-local-maxima-and-minima
  loc.max <- localMaxima(plot_data_head$reads_diffdiff_smooth)
  # local maxima are then colored green as lines in plots
  plot_data_head_sub <- plot_data_head[loc.max,]
  # which.peaks(plot_data_head$reads_diff_smooth2)

  knee_plot_ext <- ggplot(plot_data_head, aes(x=rank, y=cum_sum)) +
    xlim(0,x_scale) +
    ylim(0,1) +
    geom_text(data=plot_data_head_sub,aes(label = round(cum_sum,2)),nudge_y=-0.05, vjust = "inward", hjust = "inward") +
    geom_vline(xintercept=snakemake@params$cells, linetype="dashed", color = "red") +
    #   geom_vline(xintercept=100, linetype="dashed", color = "red") +
    geom_vline(xintercept=loc.max, col="lightgreen") +
    geom_point(size = 0.1)  +
    ggtitle(paste0(snakemake@wildcards$sample, '\nTotal reads: ', prettyNum(total_reads), mystats)) +
    theme(plot.title = element_text(size=10)) +
    labs(x='STAMPS', y='Cumulative fraction of reads')
  read_count_plot <- ggplot(plot_data_head, aes(x=rank, y=read_count)) +
    geom_text(data=plot_data_head_sub,aes(label = read_count),nudge_y=-0.05, vjust = "inward", hjust = "inward", check_overlap = TRUE) +
    #   geom_smooth() +
    xlim(0,x_scale) +
    #   ylim(0,1000) +
    geom_vline(xintercept=snakemake@params$cells, linetype="dashed", color = "red") +
    #   geom_vline(xintercept=100, linetype="dashed", color = "red") +
    geom_vline(xintercept=loc.max, col="lightgreen") +
    geom_point(size = 0.1)  +
    theme(plot.title = element_text(size=10)) +
    labs(x='STAMPS', y='Read counts per STAMP')

  # knee_plot_ext = knee_plot_ext + scale_y_continuous(labels = scales::percent)
  diff_plot <- ggplot(plot_data_head, aes(x=rank, y=diff)) +
    geom_text(data=plot_data_head_sub,aes(label = diff),nudge_y=-0.05, vjust = "inward", hjust = "inward", check_overlap = TRUE) +
    #   geom_smooth() +
    xlim(0,x_scale) +
    #   ylim(0,1000) +
    geom_vline(xintercept=snakemake@params$cells, linetype="dashed", color = "red") +
    #   geom_vline(xintercept=100, linetype="dashed", color = "red") +
    geom_vline(xintercept=loc.max, col="lightgreen") +
    geom_point(size = 0.1)  +
    theme(plot.title = element_text(size=10)) +
    #1st derivative read count diff to next STAMP")
    labs(x='STAMPS', y="1st derivative of read counts")
  diff_diff_plot <- ggplot(plot_data_head,
                           aes(x=rank, y=reads_diffdiff_smooth)) +
    geom_text(data=plot_data_head_sub,
              aes(label = rank),nudge_y=-1,
              vjust = "inward",
              hjust = "inward",
              check_overlap = TRUE) +
    xlim(0,x_scale) +
    ylim(-30,30) +
    geom_vline(xintercept=snakemake@params$cells,
               linetype="dashed", color = "red") +
    #   geom_vline(xintercept=100, linetype="dashed", color = "red") +
    geom_vline(xintercept=loc.max, col="lightgreen") +
    geom_point(size = 0.1)  +
    theme(plot.title = element_text(size=10)) +
    labs(x='STAMPS', y='2nd derivative of read counts')

  if(!is.null(snakemake@input$barcodes))
  {
    selected_cells <- read.csv(snakemake@input$barcodes, header=FALSE, stringsAsFactors=FALSE)
    knee_plot_ext <- knee_plot_ext +
      geom_point(data = plot_data_head[plot_data_head$Barcode %in% selected_cells$V1,],
                 aes(x=rank, y=cum_sum, color='Selected'), size=0.1) +
      scale_color_manual(values=c('Selected'='green')) +
      theme(legend.position="none")
    diff_diff_plot <- diff_diff_plot +
      geom_point(data = plot_data_head[plot_data_head$Barcode %in% selected_cells$V1,],
                 aes(x=rank, y=reads_diffdiff_smooth, color='Selected'), size=0.1) +
      scale_color_manual(values=c('Selected'='green')) +
      theme(legend.position="none")
  }
  #   scale_y_continuous(position = "right")
  gp1 <- ggplotGrob(knee_plot_ext)
  gp2 <- ggplotGrob(read_count_plot)
  gp3 <- ggplotGrob(diff_plot)
  gp4 <- ggplotGrob(diff_diff_plot)
  # grid::grid.newpage()
  # gg <- grid::grid.draw(rbind(gp1, gp2, gp3, gp4, size = "last"))
  gg <- gridExtra::arrangeGrob(rbind(gp1, gp2, gp3, gp4, size = "last"))
  # gg <- gridExtra::arrangeGrob(gp1, gp2, gp3, gp4, ncol = 1)

  # if barcode.csv is present in base directory, only use barcodes in there (rule plot_knee_plot_whitelist in map.smk)
  ggsave(gg, file=paste0(snakemake@output$pdf, "_extended.pdf"), width = 9, height = 11)
}

if (debug_flag) {
  save.image(file = file.path(path_debug,
       paste0("plot_knee_plot_workspace_",
              attr(snakemake, "wildcard")$sample, ".rdata"))
            )
}