small_RNA-seq.snakefile

"""
Snakefile to analyse small_RNA-seq data.

TODO: Some figures and summaries may be overridden when changing the mapper. The mapper name should be added to their path.

There are 2 types of counts. One resulting from small RNA annotations using a custom script, and one resulting from remapping and counting using featureCount.

The first one can contain "chimeric" annotations (possibly resulting from overlapping annotations?):
```
bli@naples:/extra2/Documents/Mhe/bli/small_RNA-seq_analyses$ head results_HRDE1_IP/bowtie2/mapped_C_elegans/annotation/hrde1_flag_HRDE1_IP_1_18-26_on_C_elegans/all_siu_counts.txt
gene	count
DNA?|DNA?|NDNAX1_CE:81	1
DNA?|DNA?|NDNAX2_CE:110_and_WBGene00005903	1
DNA?|DNA?|NDNAX2_CE:45	2
DNA?|DNA?|NDNAX3_CE:31	1
DNA|CMC-Chapaev|Chapaev-1_CE:45	1
DNA|CMC-Chapaev|Chapaev-1_CE:58_and_WBGene00005792	1
DNA|CMC-Chapaev|Chapaev-2_CE:40_and_WBGene00021886	1
DNA|CMC-Chapaev|Chapaev-2_CE:41_and_WBGene00021886	3
DNA|CMC-Chapaev|Chapaev-2_CE:42_and_WBGene00021886	5

Maybe it would be better to not use these counts.
```

"""
import sys
major, minor = sys.version_info[:2]
if major < 3 or (major == 3 and minor < 6):
    sys.exit("Need at least python 3.6\n")

# TODO: Try to find what the proportion of small reads map at unique position within repetitive elements.
# What is the distribution of unique reads among repetitive elements of a given family: is it evenly spread, or biased?
# Is it different if we use all small RNAs instead of just unique ones?
# For bowtie2, allow MAPQ >= 23 to get "not too ambiguously mapped reads"

# TODO: implement RPM folds from feature_counts results (and also for Ribo-seq pipeline)

# TODO: scatterplots log(IP_RPM) vs log(input_RPM) (or mutant vs WT), colouring with gene list
# TODO: heatmaps with rows (genes) ranked by expression in WT (mean reference level (wild types or inputs)), columns representing fold in mut vs WT
# Possibly remove the low expressed ones
# TODO: MA plot: fold vs "mean" count (to be defined)?

# TODO: extract most abundant piRNA reads (WT_48hph)
#bioawk -c fastx '{print $seq}' results_prg1_HS/bowtie2/mapped_C_elegans/reads/WT_RT_1_18-26_on_C_elegans/piRNA.fastq.gz | sort | uniq -c | sort -n | mawk '$1 >= 25 {print}' | wc -l
# -> 6857
# map them on the genome (without the piRNA loci), or on pseudo-genomes with specific features (histone genes...), tolerating a certain amount of mismatches (using bowtie1)
# Or: scan the genome (or "pseudo-genome") to find where there are matches with 0 to 3 mismatches at specific positions (13-21) (and perfect matches in the 5' zone)

# TODO: meta-profiles with IP and input on the same meta profile for Ortiz_oogenic and Ortiz_spermatogenic
# Make external script to generate meta-profiles on a custom list of libraries, and a given list of genes.

# TODO: boxplots with contrasts as series and a selection of gene lists
# Ex: 3 time points and oogenic and spermatogenic -> 6 boxes (custom program needed)
# Or simpler: one figure per gene list -> done for a full group of contrasts in make_gene_list_lfc_boxplots
# TODO: boxplots for all genes, for different small RNA categories (same as above but with all genes)

# Total number of "non-structural" (mapped - (fwd-(t,sn,sno,r-RNA))) to compute RPKM
# Quick-and-dirty: use gene span (TES - TSS) -> done for repeats
# Better: use length of union of exons for each gene -> done for genes
# Ideal: use transcriptome-based per-isoform computation
# Actually, we don't want a normalization by length. We deal with small RNAs, not transcripts
# TODO: make a bigwig of IP/input: is it possible?
# For RNA-seq?: use (oocyte, spermatogenic) intersected with "CSR-1-loaded" (=targets for different time point), csr1ADH_vs_WT_all_alltypes_up_genes_ids.txt (and intersected with non-spermatogenic) for meta-profiles


# Possibly filter out on RPKM
# Then, compute folds of RP(K)M IP/input (for a given experiment, i.-e. REP)
# and give list of genes sorted by such folds -> see /Gene_lists/csr1_prot_si_supertargets*

# Exploratory:
# Heatmap (fold)
# genes | experiment

# Then either:
# - take the genes (above log-fold threshold) common across replicates
# - look at irreproducible discovery rate (http://dx.doi.org/doi:10.1214%2F11-AOAS466)
# -> define CSR-1-loaded

# For metagene: see --metagene option of computeMatrix
# Retrieve gtf info after filtering out interfering genes based on merged bed

# TODO
# Locate potential targets of piRNA (using perfect complementarity), in a given annot_biotypes category for a given set of piRNA (the whole list, or for instance, up or down-regulated ones...).
# To do this: map all perfect matches, make a bed file, intersect with the annotation file for the given biotype, with antisense (note that this will automatically filter out self matches) -> done but only for all piRNAs
# Plot the count of endo-siRNA around those sites (from 100 upstream up to 100 downstream, using meta-TSS-like analysis) -> done


import os
OPJ = os.path.join
from glob import glob
from re import sub
from pickle import load
from fileinput import input as finput
from sys import stderr
from subprocess import Popen, PIPE, CalledProcessError
# Garbage-collect (to be used before Popen)
# see comment on https://stackoverflow.com/a/13329386/1878788
from gc import collect

# Useful for functional style
from itertools import chain, combinations, product, repeat, starmap
from functools import partial, reduce
from operator import or_ as union
from cytoolz import concat, flip, merge_with, take_nth, valmap


def concat_lists(lists):
    return list(concat(lists))


def dont_merge(*values):
    """This function can be passed to *merge_with* so that
    merge fails when dicts share keys."""
    try:
        ((val,),) = values
    except ValueError:
        raise ValueError("The dictionaries are not supposed to share keys.")
    return val


# Useful data structures
from collections import OrderedDict as od
from collections import defaultdict, Counter


import warnings


def formatwarning(message, category, filename, lineno, line):
    """Used to format warning messages."""
    return "%s:%s: %s: %s\n" % (filename, lineno, category.__name__, message)


warnings.formatwarning = formatwarning


from snakemake.io import apply_wildcards
# from gatb import Bank
from mappy import fastx_read
# To parse SAM format
import pyBigWig

# To compute correlation coefficient
from scipy.stats.stats import pearsonr
# To catch errors when plotting KDE
from scipy.linalg import LinAlgError
# For data processing and displaying
from sklearn import preprocessing
from sklearn.decomposition import PCA
import matplotlib as mpl
# To be able to run the script without a defined $DISPLAY
# https://github.com/mwaskom/seaborn/issues/1262
#mpl.use("agg")
mpl.use("PDF")
#mpl.rcParams["figure.figsize"] = 2, 4
mpl.rcParams["font.sans-serif"] = [
    "Arial", "Liberation Sans", "Bitstream Vera Sans"]
mpl.rcParams["font.family"] = "sans-serif"
#mpl.rcParams["figure.figsize"] = [16, 30]

from matplotlib import cm
from matplotlib.colors import Normalize

from matplotlib import numpy as np
from math import ceil, floor
from matplotlib.backends.backend_pdf import PdfPages
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
# import seaborn.apionly as sns
# import husl
# predefined seaborn graphical parameters
sns.set_context("talk")

from libsmallrna import (
    PI_MIN, PI_MAX, SI_MIN, SI_MAX,
    SI_SUFFIXES, SI_PREFIXES,
    SMALL_TYPES_TREE,
    type2RNA,
    types_under,
    SMALL_TYPES,
    SI_TYPES,
    SIU_TYPES,
    SISIU_TYPES,
    RMSK_SISIU_TYPES,
    JOINED_SMALL_TYPES)
# Do this outside the workflow
#from libhts import gtf_2_genes_exon_lengths, repeat_bed_2_lengths
from libdeseq import do_deseq2
from libhts import aligner2min_mapq
from libhts import status_setter, make_empty_bigwig
from libhts import median_ratio_to_pseudo_ref_size_factors, size_factor_correlations
from libhts import plot_paired_scatters, plot_norm_correlations, plot_counts_distribution, plot_boxplots, plot_histo
from libworkflows import texscape, wc_applied, ensure_relative, cleanup_and_backup
from libworkflows import get_chrom_sizes, column_converter, make_id_list_getter
from libworkflows import read_int_from_file, strip_split, file_len, last_lines, save_plot, SHELL_FUNCTIONS
from libworkflows import filter_combinator, read_feature_counts, sum_by_family, sum_feature_counts, sum_htseq_counts, warn_context
from smincludes import rules as irules
from smwrappers import wrappers_dir
strip = str.strip

alignment_settings = {"bowtie2": "-L 6 -i S,1,0.8 -N 0"}

# Positions in small RNA sequences for which to analyse nucleotide distribution
#POSITIONS = ["first", "last"]
POSITIONS = config["positions"]
# Orientations of small RNAs with respect to an annotated feature orientation.
# "fwd" and "rev" restrict feature quantification to sense or antisense reads.
ORIENTATIONS = config["orientations"]
# Codes for types of small RNAs
# Codes ending with "u" are for uridinylated endo siRNAs
# (recognized as not having mapped without first trimming a poly-T tail.
# "prot", "te" and "pseu" prefixes indicate the type of mRNA template that is inferred to have been used for endo siRNA synthesis by RdRP
# See small_RNA_seq_annotate.py and libsmallrna.pyx for more details


# small RNA types on which to run DESeq2
DE_TYPES = config["de_types"]
assert set(DE_TYPES) <= set(SMALL_TYPES + JOINED_SMALL_TYPES), "%s\n%s" % (", ".join(DE_TYPES), ", ".join(set(SMALL_TYPES + JOINED_SMALL_TYPES)))
# small RNA types on which to compute IP/input RPM folds
# prot_si does not include polyU siRNA, so the counts here may be lower than in pisimi
# siu includes SIU_TYPES, but not SI_TYPES, so the counts here may be lower than in pisimi
# pisimi includes all (SI_TYPES, SIU_TYPES, pi and mi (pi and mi will not be on the same genes as the siRNAs))
#IP_TYPES = ["pisimi", "siu", "prot_si"]
# TODO: update cross_HTS with pimi22G
# TODO: what kind of pimi22G ? -> pi, mi and si_22G, but not siu_22G
IP_TYPES = [f"pimi{SI_MIN}G", f"siu_{SI_MIN}G", f"prot_si_{SI_MIN}G",
            f"siu_{SI_MAX}G", f"prot_si_{SI_MAX}G"]
assert set(IP_TYPES) <= set(
    SMALL_TYPES + [f"all_si_{SI_MIN}G", f"all_si_{SI_MAX}G"] + JOINED_SMALL_TYPES), ", ".join(IP_TYPES)
#IP_TYPES = config["ip_types"]
# Cutoffs in log fold change
LFC_CUTOFFS = [0.5, 1, 2]
UP_STATUSES = [f"up{cutoff}" for cutoff in LFC_CUTOFFS]
DOWN_STATUSES = [f"down{cutoff}" for cutoff in LFC_CUTOFFS]
#STANDARDS = ["zscore", "robust", "minmax", "unit"]
#STANDARDS = ["zscore", "robust", "minmax"]
# hexbin jointplot for principal components crashes on MemoryError for PCA without standardization
#STANDARDS = ["robust", "identity"]
STANDARDS = ["robust"]

COMPL = {"A" : "T", "C" : "G", "G" : "C", "T" : "A", "N" : "N"}

# Possible feature ID conversions
ID_TYPES = ["name", "cosmid"]

#########################
# Reading configuration #
#########################
# key: library name
# value: 3' adapter sequence
lib2adapt = config["lib2adapt"]
trim5 = config["trim5"]
trim3 = config["trim3"]
# key: library name
# value: path to raw data
lib2raw = config["lib2raw"]
LIBS = list(lib2raw.keys())
#REF=config["WT"]
#MUT=config["mutant"]
# Used to associate colours to libraries
# key: series name
# value: list of libraries
colour_series_dict = config["colour_series"]
genotype_series = colour_series_dict.get("genotype_series", [])
# Groups of libraries to put together on a same metagene profile
lib_groups = config["lib_groups"]
GROUP_TYPES = list(lib_groups.keys())
merged_groups = merge_with(concat_lists, *lib_groups.values())
ALL_LIB_GROUPS = list(merged_groups.keys())
#all_libs_in_group = concat_lists(merge_with(concat_lists, *lib_groups.values()).values())
REPS = config["replicates"]
#TREATS = config["treatments"]
# Conditions to compare using DESeq2
# Note: we don't want contrasts of the following type:
# - ({geno}, {geno}) where various treatments are taken into account
# - ({treat}, {treat}) where various genotypes are taken into account
# We also don't want contrasts where more than one factor vary (ex: prg1_RT vs WT_HS30)
#ALL_COND_PAIRS = [(MUT, REF)] + list(combinations(reversed(TREATS), 2)) + [("%s_%s" % (lib2, treat2), "%s_%s" % (lib1, treat1)) for ((lib1, treat1), (lib2, treat2)) in combinations(product(LIBS, TREATS), 2)]
#COND_PAIRS = [(lib2, lib1) for (lib1, lib2) in [
#    ["%s_%s" % (geno, treat) for (geno, treat) in product(genos, treats)] for (genos, treats) in chain(
#        product(combinations(LIBS, 1), combinations(TREATS, 2)),
#        product(combinations(LIBS, 2), combinations(TREATS, 1)))]]
#CONTRASTS = ["%s_vs_%s" % cond_pair for cond_pair in COND_PAIRS]
#CONTRAST2PAIR = dict(zip(CONTRASTS, COND_PAIRS))
DE_COND_PAIRS = config["de_cond_pairs"]
msg = "\n".join([
    "Some contrats do not use known library names.",
    "Contrasts:"
    ", ".join([f"({cond}, {ref})" for (cond, ref) in DE_COND_PAIRS])])
assert all([cond in LIBS and ref in LIBS for (cond, ref) in DE_COND_PAIRS]), msg
IP_COND_PAIRS = config["ip_cond_pairs"]
msg = "\n".join([
    "Some contrats do not use known library names.",
    "Contrasts:"
    ", ".join([f"({cond}, {ref})" for (cond, ref) in IP_COND_PAIRS])])
assert all([cond in LIBS and ref in LIBS for (cond, ref) in IP_COND_PAIRS]), ""
COND_PAIRS = DE_COND_PAIRS + IP_COND_PAIRS
DE_CONTRASTS = [f"{cond1}_vs_{cond2}" for [cond1, cond2] in DE_COND_PAIRS]
IP_CONTRASTS = [f"{cond1}_vs_{cond2}" for [cond1, cond2] in IP_COND_PAIRS]
contrasts_dict = {"de" : DE_CONTRASTS, "ip" : IP_CONTRASTS}
CONTRASTS = DE_CONTRASTS + IP_CONTRASTS
CONTRAST2PAIR = dict(zip(CONTRASTS, COND_PAIRS))
MIN_LEN = config["min_len"]
MAX_LEN = config["max_len"]
size_selected = "%s-%s" % (MIN_LEN, MAX_LEN)
# bli@naples:/pasteur/entites/Mhe/Genomes$ cat C_elegans/Caenorhabditis_elegans/Ensembl/WBcel235/Annotation/Genes/miRNA.bed | mawk '{hist[$3-$2]++} END {for (l in hist) print l"\t"hist[l]}' | sort -n
# 17	1
# 18	2
# 19	4
# 20	18
# 21	57
# 22	172
# 23	138
# 24	45
# 25	11
# 26	3
# This should ideally come from genome configuration:
MI_MAX = 26
read_type_max_len = {
    size_selected: int(MAX_LEN),
    "pi": min(PI_MAX, int(MAX_LEN)),
    "si": min(SI_MAX, int(MAX_LEN)),
    "mi": min(MI_MAX, int(MAX_LEN))}
READ_TYPES_FOR_COMPOSITION = [
    size_selected, "nomap_siRNA", "all_siRNA",
    f"all_si_{SI_MIN}GRNA", f"all_si_{SI_MAX}GRNA",
    *type2RNA(SMALL_TYPES)]
READ_TYPES_FOR_MAPPING = [
    *READ_TYPES_FOR_COMPOSITION,
    "all_siuRNA",
    f"si_{SI_MIN}GRNA", f"si_{SI_MAX}GRNA",
    f"siu_{SI_MIN}GRNA", f"siu_{SI_MAX}GRNA"]
# Types of annotation features, as defined in the "gene_biotype"
# GTF attribute sections of the annotation files.
COUNT_BIOTYPES = config["count_biotypes"]
BOXPLOT_BIOTYPES = config.get("boxplot_biotypes", [])
ANNOT_BIOTYPES = config["annot_biotypes"]
assert "protein_coding_3UTR" not in set(ANNOT_BIOTYPES), "It seems having 3' or 5' UTR annotations makes 22G RNA disappear from those regions. Only use the more general protein_coding_UTR in annot_biotypes"

RMSK_BIOTYPES = [
    "DNA_transposons_rmsk",
    "RNA_transposons_rmsk",
    "satellites_rmsk",
    "simple_repeats_rmsk"]
RMSK_FAMILIES_BIOTYPES = [
    "DNA_transposons_rmsk_families",
    "RNA_transposons_rmsk_families",
    "satellites_rmsk_families",
    "simple_repeats_rmsk_families"]

BIOTYPES_TO_JOIN = {
    "all_rmsk": [biotype for biotype in COUNT_BIOTYPES + BOXPLOT_BIOTYPES if biotype in RMSK_BIOTYPES],
    "all_rmsk_families": [biotype for biotype in COUNT_BIOTYPES + BOXPLOT_BIOTYPES if biotype in RMSK_FAMILIES_BIOTYPES],
    # We only count "protein_coding", not "protein_codin_{5UTR,CDS,3UTR}"
    "alltypes": [biotype for biotype in COUNT_BIOTYPES + BOXPLOT_BIOTYPES if not biotype.startswith("protein_coding_")]}
# Filter out those with empty values
JOINED_BIOTYPES = [joined_biotype for (joined_biotype, biotypes) in BIOTYPES_TO_JOIN.items() if biotypes]

# cf Gu et al. (2009), supplement:
# -----
# To compare 22G-RNAs derived from a gene, transposon, or pseudogene between two
# samples, each sample was normalized using the total number of reads less structural RNAs, i.e.
# sense small RNA reads likely derived from degraded ncRNAs, tRNAs, snoRNAs, rRNAs,
# snRNAs, and scRNAs. Degradation products of structural RNAs map to the sense strand, with a
# poorly defined size profile and 1 st nucleotide distribution. At least 25 22G-RNA reads per million,
# nonstructural reads in one of the two samples was arbitrarily chosen as a cutoff for comparison
# analyses. A change of 2-fold or more between samples was chosen as an enrichment threshold.
# Because some 21U-RNAs or miRNAs overlap with protein coding genes, reads derived from
# miRNA loci within a window of ± 4 nt and all the known 21U-RNAs were filtered out prior to
# comparison analysis.
# -----
# And Germano Cecere, about scRNA:
# -----
# Its an old nomenclature and in anycase there is only one of this annotated scRNAs
# (small cytoplasmic RNA genes).
# https://www.ncbi.nlm.nih.gov/books/NBK19701/table/genestructure_table2/?report=objectonly
# Don't even pay attention to this
# -----
STRUCTURAL_BIOTYPES = ["tRNA", "snRNA", "snoRNA", "rRNA", "ncRNA"]
GENE_LISTS = config["gene_lists"]
BOXPLOT_GENE_LISTS = config["boxplot_gene_lists"]
#BOXPLOT_GENE_LISTS = [
#    "all_genes",
#    "replication_dependent_octamer_histone",
#    "piRNA_dependent_prot_si_down4",
#    "csr1_prot_si_supertargets_48hph",
#    "spermatogenic_Ortiz_2014", "oogenic_Ortiz_2014"]
aligner = config["aligner"]
########################
# Genome configuration #
########################
genome_dict = config["genome_dict"]
genome = genome_dict["name"]
chrom_sizes = get_chrom_sizes(genome_dict["size"])
chrom_sizes.update(valmap(int, genome_dict.get("extra_chromosomes", {})))
genomelen = sum(chrom_sizes.values())
genome_db = genome_dict["db"][aligner]
# bed file binning the genome in 10nt bins
genome_binned = genome_dict["binned"]
annot_dir = genome_dict["annot_dir"]
exon_lengths_file = OPJ(annot_dir, "union_exon_lengths.txt"),
# TODO: figure out the difference between OPJ(convert_dir, "wormid2name.pickle") and genome_dict["converter"]
convert_dir = genome_dict["convert_dir"]
gene_lists_dir = genome_dict["gene_lists_dir"]
avail_id_lists = set(glob(OPJ(gene_lists_dir, "*_ids.txt")))
index = genome_db

#output_dir = config["output_dir"]
#workdir: config["output_dir"]
output_dir = os.path.abspath(".")
local_annot_dir = config.get("local_annot_dir", OPJ("annotations"))
log_dir = config.get("log_dir", OPJ("logs"))
data_dir = config.get("data_dir", OPJ("data"))
# To put the results of small_RNA_seq_annotate
mapping_dir = OPJ(aligner, f"mapped_{genome}")
reads_dir = OPJ(mapping_dir, "reads")
annot_counts_dir = OPJ(mapping_dir, "annotation")
# The order after pi and mi must match: this is used in make_read_counts_summary
ANNOT_COUNTS_TYPES = [
    "pi", "mi", "all_si",
    f"all_si_{SI_MIN}G", f"all_si_{SI_MAX}G",
    *SI_TYPES]
ANNOT_COUNTS_TYPES_U = [
    "all_siu",
    f"all_siu_{SI_MIN}G", f"all_siu_{SI_MAX}G",
    *SIU_TYPES]
feature_counts_dir = OPJ(mapping_dir, "feature_count")
READ_TYPES_FOR_COUNTING = [f"{size_selected}_and_nomap_siRNA", f"prot_si_{SI_MIN}GRNA_and_prot_siu_{SI_MIN}GRNA", *type2RNA([f"prot_si_{SI_MIN}G", f"prot_si_{SI_MAX}G"])]
# Reads remapped and counted using featureCounts
REMAPPED_COUNTED = [
    f"{small_type}_{mapping_type}_{biotype}_{orientation}_transcript" for (
        small_type, mapping_type, biotype, orientation) in product(
            READ_TYPES_FOR_COUNTING,
            [f"on_{genome}", f"unique_on_{genome}"],
            #[f"on_{genome}"],
            COUNT_BIOTYPES + JOINED_BIOTYPES,
            ORIENTATIONS)]
# Used to skip some genotype x treatment x replicate number combinations
# when some of them were not sequenced
forbidden = {frozenset(wc_comb.items()) for wc_comb in config["missing"]}
CONDITIONS = [{
    "lib" : lib,
    "rep" : rep} for rep in REPS for lib in LIBS]
# We use this for various things in order to have always the same library order:
COND_NAMES = ["_".join((
    cond["lib"],
    cond["rep"])) for cond in CONDITIONS]
COND_COLUMNS = pd.DataFrame(CONDITIONS).assign(
    cond_name=pd.Series(COND_NAMES).values).set_index("cond_name")
#SIZE_FACTORS = ["raw", "deduped", size_selected, "mapped", "siRNA", "miRNA"]
#SIZE_FACTORS = [size_selected, "mapped", "miRNA"]
# TESTED_SIZE_FACTORS = ["mapped", "non_structural", "siRNA", "miRNA", "median_ratio_to_pseudo_ref"]
TESTED_SIZE_FACTORS = ["mapped", "non_structural", "all_sisiuRNA", "miRNA", "median_ratio_to_pseudo_ref"]
#SIZE_FACTORS = ["mapped", "miRNA", "median_ratio_to_pseudo_ref"]
# "median_ratio_to_pseudo_ref" is a size factor adapted from
# the method described in the DESeq paper, but with addition
# and then substraction of a pseudocount, in order to deal with zero counts.
# This seems to perform well (see "test_size_factor" results).
DE_SIZE_FACTORS = ["non_structural", "median_ratio_to_pseudo_ref"]
SIZE_FACTORS = ["non_structural"]
#NORMALIZER = "median_ratio_to_pseudo_ref"

# For metagene analyses
#META_MARGIN = 300
META_MARGIN = 0
META_SCALE = 500
#UNSCALED_INSIDE = 500
UNSCALED_INSIDE = 0
#META_MIN_LEN = 1000
META_MIN_LEN = 2 * UNSCALED_INSIDE
MIN_DIST = 2 * META_MARGIN
READ_TYPES_FOR_METAPROFILES = [
    size_selected,
    f"all_si_{SI_MIN}GRNA", f"all_si_{SI_MAX}GRNA",
    f"si_{SI_MIN}GRNA", f"si_{SI_MAX}GRNA",
    f"siu_{SI_MIN}GRNA", f"siu_{SI_MAX}GRNA",
    "miRNA", "piRNA", "all_siRNA"]

# split = str.split
#
#
# def strip_split(text):
#     return split(strip(text), "\t")
#
#
# # http://stackoverflow.com/a/845069/1878788
# def file_len(fname):
#     p = Popen(
#         ['wc', '-l', fname],
#         stdout=PIPE,
#         stderr=PIPE)
#     (result, err) = p.communicate()
#     if p.returncode != 0:
#         raise IOError(err)
#     return int(result.strip().split()[0])


def add_dataframes(df1, df2):
    return df1.add(df2, fill_value=0)


def sum_dataframes(dfs):
    return reduce(add_dataframes, dfs)


def sum_counts(fname):
    collect()
    p = Popen(
        # ['awk', '$1 ~ /^piRNA$|^miRNA$|^pseudogene$|^satellites_rmsk$|^simple_repeats_rmsk$|^protein_coding_|^.NA_transposons_rmsk$/ {sum += $2} END {print sum}', fname],
        # slightly faster
        ['mawk', '$1 ~ /^piRNA$|^miRNA$|^pseudogene$|^satellites_rmsk$|^simple_repeats_rmsk$|^protein_coding_|^.NA_transposons_rmsk$/ {sum += $2} END {print sum}', fname],
        stdout=PIPE,
        stderr=PIPE)
    (result, err) = p.communicate()
    if p.returncode != 0:
        raise IOError(err)
    try:
        return int(result.strip().split()[0])
    except IndexError:
        warnings.warn(f"No counts in {fname}\n")
        return 0

def sum_te_counts(fname):
    collect()
    p = Popen(
        ['awk', '$1 !~ /WBGene/ {sum += $2} END {print sum}', fname],
        stdout=PIPE,
        stderr=PIPE)
    (result, err) = p.communicate()
    if p.returncode != 0:
        raise IOError(err)
    return int(result.strip().split()[0])


import matplotlib.patheffects


class Scale(matplotlib.patheffects.RendererBase):
    def __init__(self, sx, sy=None):
        self._sx = sx
        self._sy = sy

    def draw_path(self, renderer, gc, tpath, affine, rgbFace):
        affine = affine.identity().scale(self._sx, self._sy) + affine
        renderer.draw_path(gc, tpath, affine, rgbFace)


def print_wc(wildcards):
    print(wildcards)
    return []


def check_wc(wildcards):
    #if hasattr(wildcards, "read_type"):
    #    if wildcards.read_type.endswith("_on_C"):
    #        print(wildcards)
    if any(attr.endswith("_on_C") for attr in vars(wildcards).values() if type(attr) == "str"):
        print(wildcards)
    return []


filtered_product = filter_combinator(product, forbidden)

# Limit risks of ambiguity by imposing replicates to be numbers
# and restricting possible forms of some other wildcards
wildcard_constraints:
    lib="|".join(LIBS),
    #treat="|".join(TREATS),
    rep="\d+",
    min_dist="\d+",
    min_len="\d+",
    #max_len="\d+",
    biotype="|".join(set(COUNT_BIOTYPES + ANNOT_BIOTYPES + JOINED_BIOTYPES)),
    id_list="|".join(GENE_LISTS),
    type_set="|".join(["all", "protein_coding", "protein_coding_TE"]),
    mapping_type="|".join([f"on_{genome}", f"unique_on_{genome}"]),
    #small_type="si|siu|sisiu|all_si|all_siu|all_sisiu|%s" % "|".join(SMALL_TYPES + JOINED_SMALL_TYPES),
    small_type="si|siu|sisiu|all_si|all_siu|all_sisiu|all_si_%sG|all_si_%sG|%s" % (SI_MIN, SI_MAX, "|".join(SMALL_TYPES + JOINED_SMALL_TYPES)),
    #small_type="si|siu|sisiu|all_si|all_siu|all_sisiu|%s" % "|".join(
    #    [f"all_si_{SI_MIN}G", "all_si_{SI_MAX}G"] + SMALL_TYPES + JOINED_SMALL_TYPES),
    mapped_type="|".join([size_selected, "nomap_siRNA"]),
    read_type="|".join([*REMAPPED_COUNTED, *READ_TYPES_FOR_MAPPING, *READ_TYPES_FOR_COUNTING, "trimmed", "nomap"]),
    infix="si|siu",
    suffix="|".join(SI_SUFFIXES),
    standard="zscore|robust|minmax|unit|identity",
    orientation="all|fwd|rev",
    contrast="|".join(CONTRASTS),
    norm="|".join(TESTED_SIZE_FACTORS),
    lib_group="|".join(ALL_LIB_GROUPS),
    group_type="|".join(GROUP_TYPES),
    fold_type="|".join(["mean_log2_RPM_fold", "log2FoldChange", "lfcMLE"]),

#ruleorder: map_on_genome > sam2indexedbam > compute_coverage > remap_on_genome > resam2indexedbam > recompute_coverage

# Define functions to be used in shell portions
shell.prefix(SHELL_FUNCTIONS)


###################################
# Preparing the input of rule all #
###################################

bigwig_files = [
    # individual libraries
    expand(
        expand(
            OPJ(mapping_dir, "{lib}_{rep}",
                "{lib}_{rep}_{{read_type}}_on_%s_by_{{norm}}_{{orientation}}.bw" % genome),
            filtered_product, lib=LIBS, rep=REPS),
        read_type=READ_TYPES_FOR_MAPPING, norm=SIZE_FACTORS, orientation=["all"]),
    # means of replicates
    expand(
        expand(
            OPJ(mapping_dir, "{lib}_mean",
                "{lib}_mean_{{read_type}}_on_%s_by_{{norm}}_{{orientation}}.bw" % genome),
            filtered_product, lib=LIBS),
        read_type=READ_TYPES_FOR_MAPPING, norm=SIZE_FACTORS, orientation=["all"]),
    ]

meta_profiles = [
        #expand(OPJ(local_annot_dir, "transcripts_{type_set}", "merged_isolated_{min_dist}.bed"), type_set=["all", "protein_coding", "protein_coding_TE"], min_dist="0 5 10 25 50 100 250 500 1000 2500 5000 10000".split()),
        #expand(OPJ(local_annot_dir, "transcripts_{type_set}", "merged_isolated_{min_dist}_{biotype}_min_{min_len}.bed"), type_set=["all", "protein_coding", "protein_coding_TE"], min_dist="0 5 10 25 50 100 250 500 1000 2500 5000 10000".split(), biotype=["protein_coding"], min_len=[str(META_MIN_LEN)]),
        [expand(
            OPJ("figures", "mean_meta_profiles_meta_scale_{meta_scale}",
                "{read_type}_by_{norm}_{orientation}_on_{type_set}_merged_isolated_{min_dist}_{biotype}_min_{min_len}_{group_type}_{lib_group}_meta_profile.pdf"),
            meta_scale=[str(META_SCALE)],
            read_type=READ_TYPES_FOR_METAPROFILES,
            norm=SIZE_FACTORS, orientation=["all"], type_set=["protein_coding_TE"], min_dist=[str(MIN_DIST)],
            biotype=["protein_coding", "DNA_transposons_rmsk", "RNA_transposons_rmsk"], min_len=[str(META_MIN_LEN)],
            group_type=[group_type], lib_group=list(lib_group.keys())) for (group_type, lib_group) in lib_groups.items()],
        # Same as above ?
        # [expand(
        #     OPJ("figures", "mean_meta_profiles_meta_scale_{meta_scale}",
        #         "{read_type}_by_{norm}_{orientation}_on_{type_set}_merged_isolated_{min_dist}_{biotype}_min_{min_len}_{group_type}_{lib_group}_meta_profile.pdf"),
        #     meta_scale=[str(META_SCALE)], read_type=[size_selected, "siRNA", "siuRNA", "miRNA", "piRNA", "all_siRNA"],
        #     norm=SIZE_FACTORS, orientation=["all"], type_set=["protein_coding_TE"], min_dist=[str(MIN_DIST)],
        #     biotype=["protein_coding", "DNA_transposons_rmsk", "RNA_transposons_rmsk"], min_len=[str(META_MIN_LEN)],
        #     group_type=[group_type], lib_group=list(lib_group.keys())) for (group_type, lib_group) in lib_groups.items()],
        [expand(
            OPJ("figures", "mean_meta_profiles_meta_scale_{meta_scale}",
                "{read_type}_by_{norm}_{orientation}_on_{type_set}_merged_isolated_{min_dist}_{biotype}_min_{min_len}_{group_type}_{lib_group}_meta_profile.pdf"),
            meta_scale=[str(META_SCALE)],
            read_type=READ_TYPES_FOR_METAPROFILES,
            norm=SIZE_FACTORS, orientation=["all"], type_set=["protein_coding"], min_dist=[str(MIN_DIST)],
            biotype=["protein_coding"], min_len=[str(META_MIN_LEN)],
            group_type=[group_type], lib_group=list(lib_group.keys())) for (group_type, lib_group) in lib_groups.items()],
        [expand(
            OPJ("figures", "mean_meta_profiles_meta_scale_{meta_scale}",
                "{read_type}_by_{norm}_{orientation}_on_{type_set}_merged_isolated_{min_dist}_{id_list}_{group_type}_{lib_group}_meta_profile.pdf"),
            meta_scale=[str(META_SCALE)],
            read_type=READ_TYPES_FOR_METAPROFILES,
            norm=SIZE_FACTORS, orientation=["all"], type_set=["protein_coding_TE"], min_dist=["0"], id_list=GENE_LISTS,
            group_type=[group_type], lib_group=list(lib_group.keys())) for (group_type, lib_group) in lib_groups.items()],
        ## TODO: Resolve issue with bedtools
        # expand(
        #     OPJ("figures", "{lib}_{rep}",
        #         "{read_type}_by_{norm}_{orientation}_pi_targets_in_{biotype}_profile.pdf"),
        #     lib=LIBS, rep=REPS, read_type=READ_TYPES_FOR_METAPROFILES,
        #     norm=SIZE_FACTORS, orientation=["all"], biotype=["protein_coding"]),
    ]

read_graphs = [
    expand(
        expand(
            OPJ("figures", "{{lib}}_{{rep}}", "{read_type}_base_composition_from_{position}.pdf"),
            read_type=READ_TYPES_FOR_COMPOSITION, position=["start", "end"]),
        filtered_product, lib=LIBS, rep=REPS),
    expand(
        expand(
            OPJ("figures", "{{lib}}_{{rep}}", "{read_type}_base_logo_from_{position}.pdf"),
            read_type=READ_TYPES_FOR_COMPOSITION, position=["start", "end"]),
        filtered_product, lib=LIBS, rep=REPS),
    expand(
        expand(
            OPJ("figures", "{{lib}}_{{rep}}", "{read_type}_{position}_base_composition.pdf"),
            read_type=READ_TYPES_FOR_COMPOSITION, position=POSITIONS),
        filtered_product, lib=LIBS, rep=REPS),
    expand(
        expand(
            OPJ("figures", "{{lib}}_{{rep}}", "{read_type}_{position}_base_logo.pdf"),
            read_type=READ_TYPES_FOR_COMPOSITION, position=POSITIONS),
        filtered_product, lib=LIBS, rep=REPS),
    expand(
        expand(
            OPJ("figures", "{{lib}}_{{rep}}", "{read_type}_size_distribution.pdf"),
            read_type=["trimmed", "nomap"]),
        filtered_product, lib=LIBS, rep=REPS),
    expand(
        OPJ("figures", "{lib}_{rep}", f"{size_selected}_smallRNA_barchart.pdf"),
        filtered_product, lib=LIBS, rep=REPS),
    expand(
        OPJ("figures", "{lib}_{rep}", "nb_reads.pdf"),
        filtered_product, lib=LIBS, rep=REPS),
    ]

ip_fold_heatmaps = []
de_fold_heatmaps = []
ip_fold_boxplots = []
# Temporary, until used for boxplots:
remapped_folds = []
remapped_fold_boxplots = []
if contrasts_dict["ip"]:
    if BOXPLOT_BIOTYPES:
        remapped_fold_boxplots = expand(
            OPJ("figures", "{contrast}", "by_biotypes",
                "{contrast}_{read_type}_{mapping_type}_{biotypes}_{orientation}_transcript_{fold_type}_{gene_list}_boxplots.pdf"),
            contrast=IP_CONTRASTS,
            read_type=READ_TYPES_FOR_COUNTING,
            mapping_type=[f"on_{genome}", f"unique_on_{genome}"],
            #mapping_type=[f"on_{genome}"],
            # TODO: Read this from config file
            biotypes=["_and_".join(BOXPLOT_BIOTYPES)],
            orientation=ORIENTATIONS,
            fold_type=["mean_log2_RPM_fold"],
            gene_list=BOXPLOT_GENE_LISTS)
    remapped_folds = expand(
        OPJ(mapping_dir, f"RPM_folds_{size_selected}", "all", "remapped", "{counted_type}_mean_log2_RPM_fold.txt"),
        counted_type=REMAPPED_COUNTED)
    # Should be generated as a dependency of the "all" contrast:
    ## TODO: check that this is generated
    # remapped_folds += expand(
    #     OPJ(mapping_dir, f"RPM_folds_{size_selected}", "{contrast}", "remapped",
    #         "{contrast}_{read_type}_{mapping_type}_{biotype}_{orientation}_transcript_RPM_folds.txt"),
    #     contrast=IP_CONTRASTS,
    #     read_type=READ_TYPES_FOR_COUNTING,
    #     mapping_type=[f"on_{genome}", f"unique_on_{genome}"],
    #     biotype=["alltypes"],
    #     orientation=ORIENTATIONS)
    ##
    # snakemake -n OK
    # expand(
    #     OPJ(mapping_dir, f"RPM_folds_{size_selected}", "all",
    #         "{read_type}_{mapping_type}_{biotype}_{orientation}_transcript_mean_log2_RPM_fold.txt"),
    #     read_type=READ_TYPES_FOR_COUNTING, mapping_type=[f"on_{genome}"], biotype=COUNT_BIOTYPES, orientation=ORIENTATIONS),
    # snakemake -n OK
    # expand(
    #     OPJ(mapping_dir, f"RPM_folds_{size_selected}", "{contrast}",
    #         "{contrast}_{small_type}_{mapping_type}_{biotype}_{orientation}_transcript_RPM_folds.txt"),
    #     contrast=IP_CONTRASTS, small_type=READ_TYPES_FOR_COUNTING, mapping_type=[f"on_{genome}"], biotype=COUNT_BIOTYPES, orientation=ORIENTATIONS),
    ip_fold_heatmaps = expand(
        OPJ("figures", "fold_heatmaps", "{small_type}_{fold_type}_heatmap.pdf"),
        small_type=IP_TYPES, fold_type=["mean_log2_RPM_fold"])
    ip_fold_boxplots = expand(
        OPJ("figures", "all_{contrast_type}",
            "{contrast_type}_{small_type}_{fold_type}_{gene_list}_boxplots.pdf"),
        contrast_type=["ip"], small_type=IP_TYPES, fold_type=["mean_log2_RPM_fold"],
        gene_list=BOXPLOT_GENE_LISTS)
    # Subdivided by biotype (to enable distinction between 5'UTR, CDS and 3'UTR)
    ## TODO: activate this?
    # ip_fold_boxplots = expand(
    #     OPJ("figures", "all_{contrast_type}",
    #         "{contrast_type}_{small_type}_{fold_type}_{gene_list}_{biotype}_boxplots.pdf"),
    #     contrast_type=["ip"], small_type=IP_TYPES, fold_type=["mean_log2_RPM_fold"],
    #     gene_list=BOXPLOT_GENE_LISTS, biotype=["protein_coding_5UTR", "protein_coding_CDS", "protein_coding_3UTR"])
    ##
if contrasts_dict["de"]:
    de_fold_heatmaps = expand(
        OPJ("figures", "fold_heatmaps", "{small_type}_{fold_type}_heatmap.pdf"),
        small_type=DE_TYPES, fold_type=["log2FoldChange"])

exploratory_graphs = [
    ip_fold_heatmaps,
    de_fold_heatmaps,
    # Large figures, not very readable
    # expand(
    #     OPJ(mapping_dir, f"deseq2_{size_selected}", "{contrast}", "{contrast}_{small_type}_by_{norm}_pairplots.pdf"),
    #     contrast=DE_CONTRASTS, small_type=DE_TYPES, norm=DE_SIZE_FACTORS),
    ## TODO: debug PCA
    #expand(
    #    OPJ("figures", "{small_type}_{standard}_PCA.pdf"),
    #    small_type=["pisimi"], standard=STANDARDS),
    #expand(
    #    OPJ("figures", "{small_type}_{standard}_PC1_PC2_distrib.pdf"),
    #    small_type=["pisimi"], standard=STANDARDS),
    ##
    #expand(
    #    OPJ("figures", "{small_type}_clustermap.pdf"),
    #    small_type=SMALL_TYPES),
    #expand(
    #    OPJ("figures", "{small_type}_zscore_clustermap.pdf"),
    #    small_type=SMALL_TYPES),
    #expand(
    #    OPJ("figures", "{contrast}", "{small_type}_zscore_clustermap.pdf"),
    #    contrast=CONTRASTS, small_type=DE_TYPES),
    ]

de_fold_boxplots = expand(
    OPJ("figures", "{contrast}",
        "{contrast}_{small_type}_{fold_type}_{gene_list}_boxplots.pdf"),
    contrast=DE_CONTRASTS, small_type=DE_TYPES, fold_type=["log2FoldChange", "mean_log2_RPM_fold"],
    gene_list=["all_gene_lists"])
ip_fold_boxplots_by_contrast = expand(
    OPJ("figures", "{contrast}",
        "{contrast}_{small_type}_{fold_type}_{gene_list}_boxplots.pdf"),
    contrast=IP_CONTRASTS, small_type=IP_TYPES, fold_type=["mean_log2_RPM_fold"],
    gene_list=["all_gene_lists"])

fold_boxplots = [de_fold_boxplots, ip_fold_boxplots_by_contrast, ip_fold_boxplots]

rule all:
    input:
        expand(
            OPJ(mapping_dir, "{lib}_{rep}", "{read_type}_on_%s_coverage.txt" % genome),
            filtered_product, lib=LIBS, rep=REPS, read_type=[size_selected]),
        bigwig_files,
        meta_profiles,
        # snakemake -n OK
        # expand(
        #     expand(
        #         OPJ(feature_counts_dir, "summaries",
        #             "{{lib}}_{{rep}}_{read_type}_on_%s_{orientation}_transcript_counts.txt" % (genome)),
        #         read_type=READ_TYPES_FOR_MAPPING, orientation=ORIENTATIONS),
        #     filtered_product, lib=LIBS, rep=REPS),
        # TODO
        # snakemake -n OK
        # expand(
        #     expand(
        #         OPJ(feature_counts_dir,
        #             "{{lib}}_{{rep}}_{small_type}_counts.txt"),
        #         small_type=REMAPPED_COUNTED),
        #     filtered_product, lib=LIBS, rep=REPS),
        # snakemake -n not OK: summaries contains all biotypes
        # expand(
        #     expand(
        #         OPJ(feature_counts_dir, "summaries",
        #             "{{lib}}_{{rep}}_{small_type}_counts.txt"),
        #         small_type=REMAPPED_COUNTED),
        #     filtered_product, lib=LIBS, rep=REPS),
        # REMAPPED_COUNTED = [f"{small_type}_on_{genome}_{biotype}_{orientation}_transcript" for (small_type, biotype, orientation) in product(READ_TYPES_FOR_MAPPING, set(COUNT_BIOTYPES + ANNOT_BIOTYPES), ORIENTATIONS)]
        expand(
            expand(
                OPJ(feature_counts_dir, "summaries",
                    "{{lib}}_{{rep}}_%s_and_nomap_siRNA_on_%s_{orientation}_transcript_counts.txt" % (size_selected, genome)),
                orientation=ORIENTATIONS),
            filtered_product, lib=LIBS, rep=REPS),
        expand(
            expand(
                OPJ(feature_counts_dir, "summaries",
                    "{{lib}}_{{rep}}_prot_si_%sGRNA_and_prot_siu_%sGRNA_on_%s_{orientation}_transcript_counts.txt" % (SI_MIN, SI_MIN, genome)),
                orientation=ORIENTATIONS),
            filtered_product, lib=LIBS, rep=REPS),
        read_graphs,
        #expand(OPJ(feature_counts_dir, "summaries", "{lib}_{rep}_nb_non_structural.txt"), filtered_product, lib=LIBS, rep=REPS),
        #OPJ(mapping_dir, f"RPM_folds_{size_selected}", "all", "pisimi_mean_log2_RPM_fold.txt"),
        OPJ(mapping_dir, f"RPM_folds_{size_selected}", "all", f"pimi{SI_MIN}G_mean_log2_RPM_fold.txt"),
        # Not looking ad deseq2 results any more
        #expand(OPJ(mapping_dir, f"deseq2_{size_selected}", "all", "pisimi_{fold_type}.txt"), fold_type=["log2FoldChange"]),
        #expand(OPJ(mapping_dir, "RPM_folds_%s" % size_selected, "{contrast}", "{contrast}_{small_type}_RPM_folds.txt"), contrast=IP_CONTRASTS, small_type=DE_TYPES),
        exploratory_graphs,
        remapped_folds,
        fold_boxplots,
        remapped_fold_boxplots,
        #expand(OPJ("figures", "{small_type}_unit_clustermap.pdf"), small_type=SMALL_TYPES),
        #expand(OPJ(
        #    feature_counts_dir,
        #    "all_{read_type}_{mapping_type}_{biotype}_{orientation}_transcript_counts.txt"), read_type=READ_TYPES_FOR_COUNTING, mapping_type=[f"on_{genome}"], biotype=ANNOT_BIOTYPES, orientation=ORIENTATIONS),
        # expand(OPJ(
        #     feature_counts_dir,
        #     "all_{small_type}_{mapping_type}_{biotype}_{orientation}_transcript_counts.txt"),
        #     small_type=READ_TYPES_FOR_MAPPING, mapping_type=[f"on_{genome}"], biotype=set(COUNT_BIOTYPES + ANNOT_BIOTYPES), orientation=ORIENTATIONS),
        expand(
            OPJ(annot_counts_dir, f"all_{size_selected}_on_{genome}", "{small_type}_RPM.txt"),
            small_type=["mi", *SI_TYPES, *SIU_TYPES,  f"pimi{SI_MIN}G"]),
        # piRNA and satel_siu raise ValueError: `dataset` input should have multiple elements when plotting
        # simrep_siu raise TypeError: Empty 'DataFrame': no numeric data to plot
        expand(
            OPJ("figures", "{small_type}_norm_correlations.pdf"),
            small_type=["mi", *SI_TYPES, *SIU_TYPES,  f"pimi{SI_MIN}G"]),
        expand(
            OPJ("figures", "{small_type}_norm_counts_distrib.pdf"),
            small_type=["mi", *SI_TYPES, *SIU_TYPES,  f"pimi{SI_MIN}G"]),

#absolute = "/pasteur/homes/bli/src/bioinfo_utils/snakemake_wrappers/includes/link_raw_data.rules"
#relative_include_path = "../snakemake_wrappers/includes/link_raw_data.snakefile"
#absolute_include_path = os.path.join(workflow.basedir, relative_include_path)
#assert os.path.exists(absolute_include_path)
#include: relative_include_path
include: ensure_relative(irules["link_raw_data"], workflow.basedir)


rule trim_and_dedup:
    input:
        rules.link_raw_data.output,
        #OPJ(data_dir, "{lib}_{rep}.fastq.gz"),
    params:
        adapter = lambda wildcards: lib2adapt[wildcards.lib],
        trim5 = trim5,
        trim3 = trim3,
    output:
        trimmed = OPJ(data_dir, "trimmed", "{lib}_{rep}_trimmed.fastq.gz"),
        nb_raw =  OPJ(data_dir, "trimmed", "{lib}_{rep}_nb_raw.txt"),
        nb_trimmed =  OPJ(data_dir, "trimmed", "{lib}_{rep}_nb_trimmed.txt"),
        nb_deduped =  OPJ(data_dir, "trimmed", "{lib}_{rep}_nb_deduped.txt"),
    threads: 2
    resources:
        mem_mb=1049300
    message:
        "Trimming adaptor from raw data, deduplicating reads, removing random 5' {trim5}-mers and 3' {trim3}-mers for {wildcards.lib}_{wildcards.rep}."
    benchmark:
        OPJ(log_dir, "trim_and_dedup", "{lib}_{rep}_benchmark.txt")
    log:
        cutadapt = OPJ(log_dir, "cutadapt", "{lib}_{rep}.log"),
        trim_and_dedup = OPJ(log_dir, "trim_and_dedup", "{lib}_{rep}.log"),
    shell:
        """
        zcat {input} \\
            | tee >(count_fastq_reads {output.nb_raw}) \\
            | cutadapt -a {params.adapter} --discard-untrimmed - 2> {log.cutadapt} \\
            | tee >(count_fastq_reads {output.nb_trimmed}) \\
            | dedup \\
            | tee >(count_fastq_reads {output.nb_deduped}) \\
            | trim_random_nt {params.trim5} {params.trim3}  2>> {log.cutadapt} \\
            | gzip > {output.trimmed} \\
            2> {log.trim_and_dedup}
        """


def awk_size_filter(wildcards):
    """Returns the bioawk filter to select reads of size from MIN_LEN to MAX_LEN."""
    return "%s <= length($seq) && length($seq) <= %s" % (MIN_LEN, MAX_LEN)


rule select_size_range:
    """Select (and count) reads in the correct size range."""
    input:
        rules.trim_and_dedup.output.trimmed
    output:
        selected = OPJ(data_dir, "trimmed", "{lib}_{rep}_%s.fastq.gz" % size_selected),
        nb_selected = OPJ(data_dir, "trimmed", "{lib}_{rep}_nb_%s.txt" % size_selected),
    params:
        awk_filter = awk_size_filter,
    message:
        "Selecting reads size %s for {wildcards.lib}_{wildcards.rep}." % size_selected
    shell:
        """
        bioawk -c fastx '{params.awk_filter} {{print "@"$name" "$4"\\n"$seq"\\n+\\n"$qual}}' {input} \\
            | tee >(count_fastq_reads {output.nb_selected}) \\
            | gzip > {output.selected}
        """


# TODO: update this
@wc_applied
def source_fastq(wildcards):
    """Determine the fastq file corresponding to a given read type."""
    read_type = wildcards.read_type
    if read_type == "raw":
        return rules.link_raw_data.output
    elif read_type == "trimmed":
        return rules.trim_and_dedup.output.trimmed
    elif read_type == size_selected:
        return rules.select_size_range.output.selected
    elif read_type == "nomap":
        return rules.map_on_genome.output.nomap_fastq
    elif read_type == "nomap_siRNA":
        return rules.extract_nomap_siRNAs.output.nomap_si
    elif read_type[:-3] in annotate_read_output:
        return annotate_read_output[read_type[:-3]]
    elif read_type == f"si_{SI_MIN}GRNA":
        return OPJ(
            reads_dir, "{lib}_{rep}_%s_on_%s" % (size_selected, genome),
            f"si_{SI_MIN}GRNA.fastq.gz"),
    elif read_type == f"si_{SI_MAX}GRNA":
        return OPJ(
            reads_dir, "{lib}_{rep}_%s_on_%s" % (size_selected, genome),
            f"si_{SI_MAX}GRNA.fastq.gz"),
    elif read_type == f"siu_{SI_MIN}GRNA":
        return OPJ(
            reads_dir, "{lib}_{rep}_%s_on_%s" % ("nomap_siRNA", genome),
            f"siu_{SI_MIN}GRNA.fastq.gz"),
    elif read_type == f"siu_{SI_MAX}GRNA":
        return OPJ(
            reads_dir, "{lib}_{rep}_%s_on_%s" % ("nomap_siRNA", genome),
            f"siu_{SI_MAX}GRNA.fastq.gz"),
    elif read_type[:-3] in annotate_read_output_U:
        return annotate_read_output_U[read_type[:-3]]
    elif read_type == f"sisiu_{SI_MIN}GRNA":
        return OPJ(reads_dir, "{lib}_{rep}_%s_on_%s" % (size_selected, genome), f"sisiu_{SI_MIN}GRNA.fastq.gz"),
    elif read_type == f"sisiu_{SI_MAX}GRNA":
        return OPJ(reads_dir, "{lib}_{rep}_%s_on_%s" % (size_selected, genome), f"sisiu_{SI_MAX}GRNA.fastq.gz"),
    elif read_type == f"all_sisiu_{SI_MIN}GRNA":
        return OPJ(reads_dir, "{lib}_{rep}_%s_on_%s" % (size_selected, genome), f"all_sisiu_{SI_MIN}GRNA.fastq.gz"),
    elif read_type == f"all_sisiu_{SI_MAX}GRNA":
        return OPJ(reads_dir, "{lib}_{rep}_%s_on_%s" % (size_selected, genome), f"all_sisiu_{SI_MAX}GRNA.fastq.gz"),
    else:
        raise NotImplementedError("Unknown read type: %s" % read_type)




rule map_on_genome:
    input:
        #rules.select_size_range.output.selected,
        #fastq = source_fastq,
        fastq = rules.select_size_range.output.selected,
    output:
        sam = temp(OPJ(mapping_dir, "{lib}_{rep}", "%s_on_%s.sam" % (size_selected, genome))),
        nomap_fastq = OPJ(aligner, f"not_mapped_{genome}", "{lib}_{rep}_%s_unmapped_on_%s.fastq.gz" % (size_selected, genome)),
    params: