shell:
    """
    cellranger count \
    {params.feature_ref} \
    {params.n_cells} \
    {params.extra_p}  \
    {params.introns} \
    --id {wildcards.sample} \
    --transcriptome {params.genome} \
    --libraries {input.libraries} \
    --localcores {threads} \
    --localmem {params.mem_gb} \
    &> {log} && \
    # a folder in results/counts/{wildcards.sample} is automatically created due to the output declared, which 
    # is a problem to move the cellranger output files. The workaround of deleting that folder fixes that.
    rm -r results/01_counts/{wildcards.sample} && \
    mv {wildcards.sample} results/01_counts/{wildcards.sample}
    """

shell:
    """
    cellranger-atac count \
    {params.extra_p} \
    --id {wildcards.sample} \
    --reference {params.genome} \
    --fastqs data/clean \
    --localcores {threads} \
    --localmem {params.mem_gb} \
    &> {log} && \
    # a folder in results/counts/{wildcards.sample} is automatically created due to the output declared, which 
    # is a problem to move the cellranger output files. The workaround of deleting that folder fixes that.
    rm -r results/01_counts/{wildcards.sample} && \
    mv {wildcards.sample} results/01_counts/{wildcards.sample}
"""

shell:
    """
    cellranger-arc count \
    {params.introns} \
    {params.extra_p} \
    --id {wildcards.sample} \
    --reference {params.genome} \
    --libraries {input.libraries} \
    --localcores {threads} \
    --localmem {params.mem_gb} \
    &> {log} && \
    # a folder in results/counts/{wildcards.sample} is automatically created due to the output declared, which 
    # is a problem to move the cellranger output files. The workaround of deleting that folder fixes that.
    rm -r results/01_arc/{wildcards.sample} && \
    mv {wildcards.sample} results/01_arc/{wildcards.sample}
    """

shell:
    """
    cellranger count \
    {params.feature_ref} \
    {params.n_cells} \
    {params.extra_p}  \
    {params.introns} \
    --id {wildcards.sample} \
    --chemistry=ARC-v1 \
    --transcriptome {params.genome} \
    --libraries {input.libraries} \
    --localcores {threads} \
    --localmem {params.mem_gb} \
    &> {log} && \
    # a folder in results/counts/{wildcards.sample} is automatically created due to the output declared, which 
    # is a problem to move the cellranger output files. The workaround of deleting that folder fixes that.
    rm -r results/01_counts/{wildcards.sample} && \
    mv {wildcards.sample} results/01_counts/{wildcards.sample}
    """

script:
    "../scripts/python/extract_barcodes.py"

shell:
    """
    java -Xmx200G -Xss1G -jar /UMICollapse/umicollapse.jar fastq -k {wildcards.hd} --tag -i {input} -o {output} 2> {log}
    """

script:
    "../scripts/python/correct_barcodes.py"

shell:
    """
    cat {input} > {output}
    """

script:
    "../scripts/R/generate_feature_ref_larry.R"

script:
    "../scripts/R/create_seurat.R"

script:
    "../scripts/R/barcode_summary.Rmd"

script:
    "../scripts/R/barcode_filtering.R"

script:
    "../scripts/R/merge_seurat.R"

script:
    "../scripts/R/RNA_exploration.Rmd"

shell:
    """
    ln -s {input[0]} {output.fw}
    ln -s {input[1]} {output.rv}
    """

shell:
    """
    ln -s {input} {output}
    """

shell:
    """
    cat {input.fw} > {output.fw} 2> {log}
    cat {input.rv} > {output.rv} 2>> {log}
    """

shell:
    """
    cat {input} > {output} 2> {log}
    """

shell:
    """
    mv {input.fw} {output.fw}
    mv {input.rv} {output.rv}
    """

shell:
    """
    mv {input.fw} {output.fw}
    mv {input.rv} {output.rv}
    """

shell:
    """
    mv {input.fw} {output.fw}
    mv {input.rv} {output.rv}
    mv {input.r3} {output.r3}
    """

script:
    "../scripts/python/create_library_csv.py"

script:
    "../scripts/python/generate_cellhashing_ref.py"

script:
    "../scripts/R/generate_feature_ref.R"

import os

"""Create the content of library.csv for the feature barcoding and multiome pipeline
from cellranger.

Based on what is written in the column lib_type from the units file,
write the following: "fastq folder,fastq name,library type".
"""

abs_path        = os.getcwd()
sample          = snakemake.wildcards["sample"]
is_feature_bc   = snakemake.params["is_feature_bc"]
is_cell_hashing = snakemake.params["is_cell_hashing"]
lib_type        = snakemake.params["library_type"]

lib_csv = dict()


########################################################################################################
# Create library dict
########################################################################################################
if is_feature_bc:
    lib_csv['FB'] = f'{abs_path}/data/clean,{sample}_FB,Custom'

if is_cell_hashing:
    lib_csv['CH'] = f'{abs_path}/data/clean,{sample}_CH,Antibody Capture'

elif lib_type == 'GEX':
    lib_csv['GEX'] = f'{abs_path}/data/clean,{sample}_GEX,Gene Expression'

elif lib_type == 'ATAC':
    lib_csv['ATAC'] = f'{abs_path}/data/clean,{sample}_ATAC,Chromatin Accessibility'

elif lib_type == 'ARC' and not is_feature_bc:
    lib_csv['ATAC'] = f'{abs_path}/data/clean,{sample}_ATAC,Chromatin Accessibility'
    lib_csv['GEX']  = f'{abs_path}/data/clean,{sample}_GEX,Gene Expression'

elif lib_type == 'ARC' and is_feature_bc:
    lib_csv['GEX']  = f'{abs_path}/data/clean,{sample}_GEX,Gene Expression'

# Save library dict to csv
out_csv = '\n'.join(lib_csv.values())
with open(snakemake.output["library"], 'wt') as output_handle:
    output_handle.write("fastqs,sample,library_type\n")
    output_handle.write(out_csv)


########################################################################################################
# Create library dict for arc
########################################################################################################
# If the library is multiome and there is feature barcode data, the pipeline must be run 2 times
# One for rna + larry and the other for rna+atac. Due to this, we need 2 different library csv files.
if lib_type == 'ARC' and is_feature_bc:

    lib_csv2 = dict()
    lib_csv2['ATAC'] = f'{abs_path}/data/clean,{sample}_ATAC,Chromatin Accessibility'
    lib_csv2['GEX']  = f'{abs_path}/data/clean,{sample}_GEX,Gene Expression'

    # Save library (arc) dict to csv
    out_csv2 = '\n'.join(lib_csv2.values())
    with open(snakemake.output["library_arc"], 'wt') as output_handle:
        output_handle.write("fastqs,sample,library_type\n")
        output_handle.write(out_csv2)

else:
    lib_csv2 = "If library type is not ARC, ignore this file."
    with open(snakemake.output["library_arc"], 'wt') as output_handle:
        output_handle.write(lib_csv2)

import gzip
import re
from Bio.SeqIO.QualityIO import FastqGeneralIterator

#------------------------------------------------------------------------------------------------------------------------------------
# Declare input/output files
#------------------------------------------------------------------------------------------------------------------------------------
input_fb  = snakemake.input["fb"]
input_cb  = snakemake.input["cb"]
output_fb = snakemake.output["filt_fb"]
output_cb = snakemake.output["filt_cb"]
barcode_patterns = snakemake.params["barcode_dict"]


#------------------------------------------------------------------------------------------------------------------------------------
# Declare functions
#------------------------------------------------------------------------------------------------------------------------------------
def extract_barcodes(input_file, barcode_patterns):
    """
    This function will take a fastq file from larry barcode enrichment and a dictionary in which the
    keys are the barcode patterns and the corresponding larry library name/color  (i.e: {...TG...AG... : GFP}).

    It will parse the fastq file, look for reads that contain the barcode and extract the barcode sequence
    from the read. All the fastq records containing barcodes will be stored in a dictionary, as a list. The
    keys of the dictionary are the larry colors, and the values are a list of the fastq records containing each
    specific barcode color.
    """

    patterns      = [re.compile(r'{}'.format(barcode)) for barcode in barcode_patterns.keys()]
    patterns_dict = {key: value for key, value in zip(patterns, barcode_patterns.values())}
    barcode_reads = {key: [] for key in barcode_patterns.values()}

    with gzip.open(input_file, 'rt') as input_handle:
        # Iterate over each record in the fastq file
        for title, seq, qual in FastqGeneralIterator(input_handle):
            # Look for barcode patterns
            for pattern in patterns:
                match = pattern.search(seq)
                if match:
                    # Update record to contain just matched sequence 
                    seq  = seq[match.start() : match.end()]
                    qual = qual[match.start() : match.end()]
                    # Save updated records to dictionary. Every key is a different barcode color, the content are all the records
                    # corresponding to that color.
                    barcode_reads[patterns_dict[pattern]] += [title, seq, qual],
                    break

    return(barcode_reads)


def extracted_bc_to_fq(output_fastqs, barcode_reads):
    """
    Take the dictionary of larry colors and fastq records containing barcodes and save it to multiple fastq
    files, one for each larry color.
    """
    # Order of output files is the same as the keys of the dictionary since both of them are taken from the same config variable
    for i, key in enumerate(barcode_reads.keys()):
        with gzip.open(output_fastqs[i], 'wt') as output_handle:
            for title, seq, qual in barcode_reads[key]:
                _ = output_handle.write(f"@{title}\n{seq}\n+\n{qual}\n")


def subset_cb_fastq(input_fastq, output_fastq, barcode_reads):
    """
    Subset the fastq file that contains the cellular barcodes (instead of larry barcodes) to contain the same read ids
    of the larry barcode fastq, which has been filtered and now does not contain all the initial fastq entries.

    Also, since the larry bc fq is split in different colors (if there is sequential barcoding), the order of fastq
    reads will be different. For this, the cellular barcode fastq must be filtered and then sorted to match the larry
    fastq.
    """
    read_primer      = re.compile(r"\s.*$")
    filtered_ids     = [read_primer.sub('', fastq_record[0]) for larry_color in barcode_reads.values() for fastq_record in larry_color]
    filtered_ids_set = set(filtered_ids)
    fastq_records    = {}

    with gzip.open(input_fastq, 'rt') as input_handle, gzip.open(output_fastq, 'wt') as output_handle:
        # Parse fastq and store records in dictionary using read_id as key
        for title, seq, qual in FastqGeneralIterator(input_handle):
            title_no_read_id = read_primer.sub('', title)
            if title_no_read_id in filtered_ids_set:
                fastq_records[title_no_read_id] = [title, seq, qual]

        # Sort the keys (read_id) of the dict to match the order of read_ids from the larry fastq
        index_map = {v: i for i, v in enumerate(filtered_ids)}
        fastq_records = dict( sorted(fastq_records.items(), key=lambda pair: index_map[pair[0]]) )

        # Save to fastq
        for title, seq, qual in fastq_records.values():
                _ = output_handle.write(f"@{title}\n{seq}\n+\n{qual}\n")


#------------------------------------------------------------------------------------------------------------------------------------
# Main
#------------------------------------------------------------------------------------------------------------------------------------
# Filter reads containing barcodes from feature barcoding read and write to new fastq file
barcode_reads = extract_barcodes(input_fb, barcode_patterns) 
extracted_bc_to_fq(output_fb, barcode_reads)

# Subset the other fastq (the one containing cellular barcodes) based on the read ids filtered before
subset_cb_fastq(input_cb, output_cb, barcode_reads) 

def generate_cellhash_ref_csv(cell_hashing_dict):
    """Generate cmo-set csv file indicating the totalseq names and barcode sequences + positions"""
    cmo_csv = "id,name,read,pattern,sequence,feature_type\n"
    for key, value in cell_hashing_dict.items():
        cmo_csv += f"{key},{key},{value[0]},{value[1]},{value[2]},Antibody Capture\n"

    return cmo_csv

cell_hashing_Abs  = snakemake.params["cellhash_ab_names"]
cell_hashing_dict = snakemake.params["cell_hashing"]
cell_hashing_dict = dict((Ab, cell_hashing_dict[Ab]) for Ab in cell_hashing_Abs) # Use just cellhash Abs corresponding to this sample

cellhash_csv = generate_cellhash_ref_csv(cell_hashing_dict)

with open(snakemake.output[0], 'wt') as output_handle:
    output_handle.write(cellhash_csv)

knitr::opts_chunk$set(
    echo           = TRUE,
    error          = FALSE,
    fig.align      = "center",
    message        = FALSE,
    warning        = FALSE,
    fig.width      = 10,
    fig.height     = 8
)

log <- file(snakemake@log[[1]], open = "wt")
sink(log)
sink(log, type = "message")

library(Seurat)
library(tidyverse)
library(DropletUtils)
source("workflow/scripts/R/functions.R")

get_larry_molinfo <- function(molinfo) {

  larry_bc_pos <- which(molinfo$feature.type == "Custom")

  larry_molinfo <- molinfo$data %>% 
    data.frame() %>% 
    dplyr::filter(gene %in% larry_bc_pos) %>% 
    mutate(larry_bc = molinfo$genes[gene]) %>% # Create col with larry bc name instead of idx from molinfo$genes
    mutate(larry_color = str_replace(larry_bc, "_.*$", "")) %>% # Add color information to split barcodes
    dplyr::select(-gene)

    return(larry_molinfo)

}


filter_umi_reads <- function(df_molinfo, read_threshold) {

  df_molinfo <- df_molinfo %>% 
    dplyr::filter(reads >= read_threshold)

  return(df_molinfo)

}


filter_bc_umis <- function(df_molinfo, umi_threshold) {

  df_molinfo <- df_molinfo %>% 
    dplyr::group_by(cell, larry_bc, larry_color) %>% 
    dplyr::count() %>% 
    dplyr::rename(n_umi = n) %>% 
    dplyr::filter(n_umi >= umi_threshold) %>% 
    dplyr::ungroup()

  return(df_molinfo)

}


efficiency_larry <- function(df_molinfo, umi_threshold, n_cells) {

  efficiency <- df_molinfo %>% 
    filter_bc_umis(umi_threshold) %>% 
    dplyr::group_by(larry_color) %>% 
    dplyr::distinct(cell) %>% 
    dplyr::count() %>% 
    dplyr::mutate(efficiency    = n/n_cells,
                  umi_threshold = umi_threshold
                  ) %>% 
    dplyr::ungroup()

  return(efficiency)

}


plot_larry_efficiency <- function(molinfo_larry, read_threshold, n_cells, umi_thresholds) {

  larry_molinfo_read_filt <- molinfo_larry %>% 
    filter_umi_reads(read_threshold)

  larry_eficiencies <- purrr::map_df(umi_thresholds,
             \(umi_threshold) efficiency_larry(larry_molinfo_read_filt, umi_threshold, n_cells)
             )

  p <- larry_eficiencies %>%
    ggplot2::ggplot(aes( x = umi_threshold, y = efficiency)) +
    geom_line() +
    facet_wrap(~ larry_color) +
    theme_bw() +
    ggtitle(paste0(read_threshold, " read threshold"))

  return(p)

}

prop_integrations_larry <- function(df_molinfo, umi_threshold) {

  integrtions <- df_molinfo %>% 
    filter_bc_umis(umi_threshold) %>% 
    dplyr::group_by(larry_color) %>% 
    dplyr::count(cell) %>% 
    dplyr::summarise(
      p_cells_multiple_integrations = sum(n > 1)/n()
    ) %>% 
    dplyr::mutate(umi_threshold = umi_threshold) %>% 
    dplyr::ungroup()

  return(integrtions)

}

plot_larry_mult_int <- function(molinfo_larry, read_threshold, umi_thresholds) {

  larry_molinfo_read_filt <- molinfo_larry %>% 
    filter_umi_reads(read_threshold)

  larry_eficiencies <- purrr::map_df(umi_thresholds,
             \(umi_threshold) prop_integrations_larry(larry_molinfo_read_filt, umi_threshold)
             )

  p <- larry_eficiencies %>%
    ggplot2::ggplot(aes( x = umi_threshold, y = p_cells_multiple_integrations)) +
    geom_line() +
    facet_wrap(~ larry_color) +
    theme_bw() +
    ggtitle(paste0(read_threshold, " read threshold")) +
    scale_x_continuous(breaks = umi_thresholds)

  return(p)

}

seurat_rds    <- snakemake@input[[1]]
molec_info_h5 <- snakemake@params[["molecule_info"]]

seurat      <- readRDS(seurat_rds)
molinfo     <- read10xMolInfo(molec_info_h5,
                              get.cell    = TRUE,
                              get.umi     = TRUE,
                              get.gem     = FALSE,
                              get.gene    = TRUE,
                              get.reads   = TRUE,
                              get.library = FALSE
                              )

# Modify cell names form molinfo file to match those from seurat (basically add whatever is)
# before and after the cellular barcode (which consists in random 16 nucleotides, by now)
cell_prefix <- Cells(seurat) %>% str_replace("[AGTC]{16}-.*", "")
cell_suffix <- Cells(seurat) %>% str_replace(".*_[AGTC]{16}", "")

molinfo$data <- molinfo$data %>% 
  data.frame() %>%
  dplyr::mutate(cell = paste0(cell_prefix, cell, cell_suffix)) %>% 
  filter(cell %in% Cells(seurat)) %>% 
  DataFrame()

molinfo_larry <- get_larry_molinfo(molinfo)

summary_larry_thresholds <- map(1:10,
    \(read_threshold) plot_larry_efficiency(
                      molinfo_larry  = molinfo_larry,
                      read_threshold = read_threshold,
                      n_cells        = length(Cells(seurat)),
                      umi_thresholds = 1:10
                      )
    )

names(summary_larry_thresholds) <- paste0(1:10, " reads")

in_tabs(summary_larry_thresholds, level = 1L)

summary_larry_integrations <- map(1:10,
    \(read_threshold) plot_larry_mult_int(
                      molinfo_larry  = molinfo_larry,
                      read_threshold = read_threshold,
                      umi_thresholds = 1:10
                      )
    )

names(summary_larry_integrations) <- paste0(1:10, " reads")

in_tabs(summary_larry_integrations, level = 1L)

log <- file(snakemake@log[[1]], open = "wt")
sink(log)
sink(log, type = "message")

library(tidyverse)

#------------------------------------------------------------------------------------------
# Load feature ref from every sample and merge it into 1
#------------------------------------------------------------------------------------------
feature_ref <- snakemake@input %>%
  purrr::map(read_csv) %>% 
  bind_rows() %>% 
  mutate(
    id    = str_replace(id, "_.*$", ""),
    name = id
  ) %>% 
  distinct() %>% 
  group_by(id) %>% 
  mutate(
    id   = paste0(id, "_", 1:n()),
    name = id
  ) %>%
  select(id, name, read, pattern, sequence, feature_type)

write_csv(feature_ref, snakemake@output[[1]])

log <- file(snakemake@log[[1]], open = "wt")
sink(log)
sink(log, type = "message")

library(tidyverse)

#------------------------------------------------------------------------------------------
# Load feature ref from larry/cellhashing
#------------------------------------------------------------------------------------------
if (length(snakemake@input[[1]]) == 1) {

	# If just cellhash or larry are set, return the same csv
	feature_ref <- read_csv(snakemake@input[[1]])
	write_csv(feature_ref, snakemake@output[[1]])

} else { # Otherwise combine both larry and cellhash references into 1

	larry_ref    <- read_csv(snakemake@input[["larry_ref"]])
	cellhash_ref <- read_csv(snakemake@input[["cell_hash_ref"]])

	combined_ref <- bind_rows(larry_ref, cellhash_ref)
	write_csv(combined_ref, snakemake@output[[1]])

}

log <- file(snakemake@log[[1]], open = "wt")
sink(log)
sink(log, type = "message")

### Libraries
library(Seurat)
library(tidyverse)

### Read and merge seurat objects if there are multiple samples
if (length(snakemake@input)  == 1) {
    seurat <- readRDS(snakemake@input[[1]])
} else {
    seurat_objects <- map(snakemake@input, \(x) readRDS(x))
    seurat         <- merge(seurat_objects[[1]], seurat_objects[2:length(seurat_objects)] %>% unlist())
}

saveRDS(seurat, snakemake@output[["seurat"]])