RNA-seq.snakefile

"""
Snakefile to analyse RNA-seq data.
"""
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: Extract reads with poly-T tail (up to 5) from non-mapped, trim the tail, remap those (like siu RNA)

# TODO: plot spike-in vs spike-in between libraries to detect anormal spike-ins: should be a straight line

# TODO: Add rules to take into account spike-ins.
# The size factor should be determined within the dynamic range of the spike-ins
# Use 1 RPKM for LDD
# See /Genomes/spike-ins/ERCC_RNA_Spike-In_Control_Mixes.pdf pp. 13-14
# Then use size factor to compute normalized RPKM and then fold changes like for small RNA IP

#TODO: same boxplots as small_RNA-seq
#TODO: for each contrast,  scatterplots with highlight of differentially expressed genes (with p-value colour code)
# and same scale for different biotypes
# boxplots of categories of genes (only protein-coding)
# Prepare scripts to compare with small RNA (fold vs fold)
#TODO: try to use bedtools intersect to do the counting (only reads strictly included in feature)
# bli@naples:/extra2/Documents/Mhe/bli/RNA_Seq_analyses$ mkfifo /tmp/gene_id
# bli@naples:/extra2/Documents/Mhe/bli/RNA_Seq_analyses$ mkfifo /tmp/count
# bli@naples:/extra2/Documents/Mhe/bli/RNA_Seq_analyses$ head -6 /tmp/WT_RT_1_protein_coding_counts_F1S.txt | tee >(cut -f9 | gtf2attr gene_id > /tmp/gene_id) | cut -f7,10 > /tmp/count & paste /tmp/gene_id /tmp/count
# "WBGene00022277"	-	170
# "WBGene00022276"	+	551
# "WBGene00022276"	+	522
# "WBGene00022276"	+	550
# "WBGene00022276"	+	550
# "WBGene00022276"	+	550


import os
OPJ = os.path.join
from glob import glob
from pickle import load
from collections import defaultdict
from itertools import combinations
from gzip import open as gzopen

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

import matplotlib as mpl
# To be able to run the script without a defined $DISPLAY
mpl.use("PDF")
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
# https://stackoverflow.com/a/42768093/1878788
# from matplotlib.backends.backend_pgf import FigureCanvasPgf
# mpl.backend_bases.register_backend('pdf', FigureCanvasPgf)
# TEX_PARAMS = {
#     "text.usetex": True,            # use LaTeX to write all text
#     "pgf.rcfonts": False,           # Ignore Matplotlibrc
#     "pgf.texsystem": "lualatex",  # hoping to avoid memory issues
#     "pgf.preamble": [
#         r'\usepackage{color}'     # xcolor for colours
#     ]
# }
# mpl.rcParams.update(TEX_PARAMS)
mpl.rcParams["font.sans-serif"] = [
    "Arial", "Liberation Sans", "Bitstream Vera Sans"]
mpl.rcParams["font.family"] = "sans-serif"

import numpy as np
import pandas as pd
# To catch errors when plotting KDE
from scipy.linalg import LinAlgError
from scipy.stats import pearsonr, shapiro
from sklearn.linear_model import LinearRegression
import seaborn as sns
# To merge bigwig files
import pyBigWig
from mappy import fastx_read

from rpy2.robjects import Formula, StrVector
from libhts import do_deseq2, status_setter, plot_lfc_distribution, plot_MA
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)
from libworkflows import ensure_relative, cleanup_and_backup
from libworkflows import get_chrom_sizes, column_converter, make_id_list_getter
from libworkflows import strip_split, save_plot, test_na_file, SHELL_FUNCTIONS, warn_context
from libworkflows import feature_orientation2stranded
from libworkflows import read_htseq_counts, sum_htseq_counts
from libworkflows import read_intersect_counts, sum_intersect_counts
from libworkflows import read_feature_counts, sum_feature_counts
from smincludes import rules as irules

alignment_settings = {"bowtie2": "", "hisat2": "", "crac": "-k 20 --stranded --use-x-in-cigar"}

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

# TODO: make this part of the configuration
LFC_RANGE = {
    "spike_ins" : (-10, 10),
    "protein_coding" : (-10, 10),
    "DNA_transposons_rmsk" : (-10, 10),
    "RNA_transposons_rmsk" : (-10, 10),
    "satellites_rmsk" : (-10, 10),
    "simple_repeats_rmsk" : (-10, 10),
    "DNA_transposons_rmsk_families" : (-10, 10),
    "RNA_transposons_rmsk_families" : (-10, 10),
    "simple_repeats_rmsk_families" : (-10, 10),
    "pseudogene" : (-10, 10),
    "alltypes" : (-10, 10)}
# 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]
#status2colour = make_status2colour(DOWN_STATUSES, UP_STATUSES)
#STATUSES = list(reversed(DOWN_STATUSES)) + ["down", "NS", "up"] + UP_STATUSES
#STATUS2COLOUR = dict(zip(STATUSES, sns.color_palette("coolwarm", len(STATUSES))))

# Or use --configfile instead
#configfile:
#    "RNA-seq.config.json"


spikein_dict = config["spikein_dict"]
spikein_microliters = spikein_dict["microliters"]
spikein_dilution_factor = spikein_dict["dilution_factor"]
# To get the expected number of attomoles of a given spike-in,
# multiply this by the value found in column 4 of Genomes/spike-ins/ERCC92_concentrations.txt
spikein_microliter_equivalent = spikein_microliters / spikein_dilution_factor
if spikein_microliter_equivalent:
    spikein_concentrations = pd.read_table(
        spikein_dict["concentrations"],
        index_col="ERCC ID",
        usecols=["ERCC ID", "concentration in Mix %s (attomoles/ul)" % spikein_dict["mix"]])
    from scipy.constants import Avogadro
    # attomole: 10^{-18} mole
    molecules_per_attomole = Avogadro / 10**18
    spikein_expected_counts = spikein_concentrations * spikein_microliter_equivalent * molecules_per_attomole
    spikein_expected_counts.columns = ["expected_counts"]

# http://sailfish.readthedocs.io/en/master/library_type.html
LIB_TYPE = config["lib_type"]
ORIENTATIONS = config["orientations"]
# key: library name
# value: dictionary
#   key: replicate number
#   value: path to raw data
lib2raw = config["lib2raw"]
LIBS = list(lib2raw.keys())
REPS = config["replicates"]
#COND_PAIRS = config["cond_pairs"]
#CONTRASTS = [f"{cond1}_vs_{cond2}" for [cond1, cond2] in 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))
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")


BIOTYPES = config["biotypes"]
DE_BIOTYPES = [biotype for biotype in LFC_RANGE.keys() if biotype in BIOTYPES]
if "spike_ins" in BIOTYPES:
    PAIR_SCATTER_BIOTYPES = ["spike_ins"]
else:
    PAIR_SCATTER_BIOTYPES = []
# TODO: make this configurable
ID_LISTS = [
    "germline_specific", "histone",
    "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"])
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"]
# 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")))
#COUNTERS = config["counters"]
# htseq_count is very slow
# bedtools intersect eats all the RAM when there are a lot of rRNAs:
# https://github.com/arq5x/bedtools2/issues/565
COUNTERS = ["feature_count"]
output_dir = config["output_dir"]
transcriptomes_dir = OPJ(output_dir, "merged_transcriptomes")
log_dir = OPJ(output_dir, f"logs_{genome}")
data_dir = OPJ(output_dir, "data")

#NORM_TYPES = config["norm_types"]
if spikein_microliter_equivalent:
    SIZE_FACTORS = ["protein_coding", "miRNA", "median_ratio_to_pseudo_ref", "spike_ins"]
    NORM_TYPES = ["spike_ins", "protein_coding", "median_ratio_to_pseudo_ref"]
else:
    SIZE_FACTORS = ["protein_coding", "miRNA", "median_ratio_to_pseudo_ref"]
    NORM_TYPES = ["protein_coding", "median_ratio_to_pseudo_ref"]
    assert "spike_ins" not in BIOTYPES
assert set(SIZE_FACTORS).issubset(set(BIOTYPES) | {"median_ratio_to_pseudo_ref"})
assert set(NORM_TYPES).issubset(set(SIZE_FACTORS))

# PolyU tail min and max length
minU, maxU = 1, 5

# Seems to interfere with temp files deletion
wildcard_constraints:
    rep="\d+",
    orientation="|".join(ORIENTATIONS),
    biotype="|".join(BIOTYPES + ["alltypes"]),
    fold_type="|".join(["mean_log2_RPM_fold", "log2FoldChange", "lfcMLE"]),

#ruleorder: join_all_counts > gather_counts

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


def specific_input(aligner):
    if aligner == "hisat2":
        files = [
            expand(OPJ(
                output_dir, aligner, f"mapped_{genome}",
                "ballgown", "{lib}_{rep}_on_%s.gtf" % genome), lib=LIBS, rep=REPS),
            OPJ(transcriptomes_dir, aligner, f"mapped_{genome}", "stringtie", "compare.stats"),
            OPJ(transcriptomes_dir, aligner, f"mapped_{genome}", "stringtie", "merged.gtf"),]
    elif aligner == "bowtie2":
        files = []
    elif aligner == "crac":
        files = []
    else:
        raise NotImplementedError("%s is not yet handeled." % aligner)
    return files

counts_files = [
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            "{lib}_mean_on_%s" % genome, "{biotype}_{orientation}_TPM.txt"),
        lib=LIBS, counter=COUNTERS, biotype=["alltypes"], orientation=ORIENTATIONS),
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{biotype}_{orientation}_TPM.txt"),
        counter=COUNTERS, biotype=["alltypes"], orientation=ORIENTATIONS),
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{biotype}_{orientation}_RPK.txt"),
        counter=COUNTERS, biotype=BIOTYPES + ["alltypes"], orientation=ORIENTATIONS),
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{biotype}_{orientation}_counts.txt"),
        counter=COUNTERS, biotype=BIOTYPES, orientation=ORIENTATIONS),
    # expand(
    #     OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
    #         "summaries", f"{{lib}}_{{rep}}_on_{genome}_{{orientation}}_counts.txt"),
    #     counter=COUNTERS, lib=LIBS, rep=REPS, orientation=ORIENTATIONS),
    # expand(
    #     OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
    #         f"{{lib}}_{{rep}}_on_{genome}", "{biotype}_{orientation}_counts.txt"),
    #     counter=COUNTERS, lib=LIBS, rep=REPS, biotype=BIOTYPES, orientation=ORIENTATIONS),
    # expand(
    #     OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
    #         f"{{lib}}_{{rep}}_on_{genome}", "{biotype}_{orientation}_counts_gene_names.txt"),
    #     counter=COUNTERS, lib=LIBS, rep=REPS, biotype=BIOTYPES, orientation=ORIENTATIONS),]
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            "summaries", "{lib}_{rep}_{mapped_type}_{orientation}_counts.txt"),
        counter=COUNTERS, lib=LIBS, rep=REPS,
        mapped_type=[f"on_{genome}", f"polyU_on_{genome}"], orientation=ORIENTATIONS),
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            "{lib}_{rep}_{mapped_type}", "{biotype}_{orientation}_counts.txt"),
        counter=COUNTERS, lib=LIBS, rep=REPS,
        mapped_type=[f"on_{genome}", f"polyU_on_{genome}"], biotype=BIOTYPES, orientation=ORIENTATIONS),
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            "{lib}_{rep}_{mapped_type}", "{biotype}_{orientation}_counts_gene_names.txt"),
        counter=COUNTERS, lib=LIBS, rep=REPS,
        mapped_type=[f"on_{genome}", f"polyU_on_{genome}"], biotype=BIOTYPES, orientation=ORIENTATIONS),]

# Actually not necessarily IP data
# TODO: generalize contrast types
if contrasts_dict["ip"]:
    ip_fold_boxplots = [
        expand(
            OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}", "fold_boxplots",
                "{contrast_type}_{orientation}_{biotype}_{fold_type}_{id_list}_boxplots.pdf"),
            counter=COUNTERS, contrast_type=["ip"], orientation=ORIENTATIONS, biotype=DE_BIOTYPES,
            fold_type=["mean_log2_RPM_fold"], id_list=ID_LISTS),]
else:
    ip_fold_boxplots = []

if contrasts_dict["de"]:
    de_fold_boxplots = [
        expand(
            OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}", "fold_boxplots",
                "{contrast_type}_{orientation}_{biotype}_{fold_type}_{id_list}_boxplots.pdf"),
            counter=COUNTERS, contrast_type=["de"], orientation=ORIENTATIONS, biotype=DE_BIOTYPES,
            fold_type=["log2FoldChange"], id_list=ID_LISTS),]
else:
    de_fold_boxplots = []

de_files = [
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            "deseq2", "{contrast}", "{orientation}_{biotype}", "counts_and_res.txt"),
        counter=COUNTERS, contrast=DE_CONTRASTS, orientation=ORIENTATIONS, biotype=DE_BIOTYPES),
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            "deseq2", "{contrast}", "{orientation}_{biotype}", "{fold_type}_distribution.pdf"),
        counter=COUNTERS, contrast=DE_CONTRASTS, orientation=ORIENTATIONS, biotype=DE_BIOTYPES,
        fold_type=["log2FoldChange"]),
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            "deseq2", "{contrast}", "{orientation}_{biotype}", "MA_with_{id_list}.pdf"),
        counter=COUNTERS, contrast=DE_CONTRASTS, orientation=ORIENTATIONS, biotype=DE_BIOTYPES,
        id_list=ID_LISTS + ["lfc_statuses"]),
    de_fold_boxplots]

bigwig_files = [
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{lib}_{rep}",
            "{lib}_{rep}_{orientation}_{mapped_type}_by_{norm_type}.bw"),
        lib=LIBS, rep=REPS, orientation=ORIENTATIONS,
        mapped_type=[f"on_{genome}", f"polyU_on_{genome}"], norm_type=NORM_TYPES),
    expand(
        OPJ(output_dir, aligner, f"mapped_{genome}", "{lib}_mean",
            "{lib}_mean_{orientation}_{mapped_type}_by_{norm_type}.bw"),
        lib=LIBS, orientation=ORIENTATIONS,
        mapped_type=[f"on_{genome}", f"polyU_on_{genome}"], norm_type=NORM_TYPES),]

if spikein_microliter_equivalent:
    spikein_files = [
        expand(OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{lib}_{rep}_spike_ins_{orientation}_TPM_response.pdf"),
            counter=COUNTERS, orientation=ORIENTATIONS, lib=LIBS, rep=REPS),]
else:
    spikein_files = []


rule all:
    input:
        specific_input(aligner),
        expand(
            OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_samtools_stats.txt"),
            lib=LIBS, rep=REPS),
        expand(
            OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_coverage.txt"),
            lib=LIBS, rep=REPS),
        bigwig_files,
        spikein_files,
        # Uses too much memory when there are many libraries
        # expand(OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
        #     f"all_on_{genome}", "scatter_pairs/{biotype}_{orientation}_{quantif_type}_scatter_pairs.pdf"),
        #     counter=COUNTERS, biotype=["alltypes"], orientation=ORIENTATIONS, quantif_type=["TPM"]),
        # expand(OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
        #     f"all_on_{genome}", "scatter_pairs/{biotype}_{orientation}_{quantif_type}_scatter_pairs.pdf"),
        #     counter=COUNTERS, biotype=PAIR_SCATTER_BIOTYPES, orientation=ORIENTATIONS, quantif_type=["counts", "RPK"]),
        expand(
            OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
                "summaries", f"all_on_{genome}_{{orientation}}_{{biotype}}_norm_correlations.pdf"),
            counter=COUNTERS, orientation=ORIENTATIONS, biotype=BIOTYPES + ["alltypes"]),
        expand(
            OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
                "summaries", f"all_on_{genome}_{{orientation}}_{{biotype}}_norm_counts_distrib.pdf"),
            counter=COUNTERS, orientation=ORIENTATIONS, biotype=BIOTYPES + ["alltypes"]),
        de_files,
        ip_fold_boxplots,
        counts_files,
        expand(
            OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
                f"all_on_{genome}", "replicates_comparison",
                "{lib}", "{biotype}_{orientation}_{quantif_type}_correlations.txt"),
            counter=COUNTERS, lib=LIBS, biotype=["alltypes"],
            orientation=ORIENTATIONS, quantif_type=["TPM"]),
        expand(
            OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
                f"all_on_{genome}", "replicates_comparison",
                "{lib}", "{biotype}_{orientation}_{quantif_type}_correlations.txt"),
            counter=COUNTERS, lib=LIBS, biotype=BIOTYPES + ["alltypes"],
            orientation=ORIENTATIONS, quantif_type=["counts", "RPM"]),


include: ensure_relative(irules["link_raw_data"], workflow.basedir)
#include: "../snakemake_wrappers/includes/link_raw_data.rules"


# def mapping_command(aligner):
#     """This function returns the shell commands to run given the *aligner*."""
#     if aligner == "hisat2":
#         shell_commands = """
# cmd="hisat2 -p {threads} --dta --seed 123 -t --mm -x %s -U {input.fastq} --no-unal --un-gz {output.nomap_fastq} -S {output.sam}"
# echo ${{cmd}} 1> {log.log}
# eval ${{cmd}} 1>> {log.log} 2> {log.err}
# """ % genome_db
#     elif aligner == "bowtie2":
#         shell_commands = """
# cmd="bowtie2 -p {threads} --seed 123 -t --mm -x %s -U {input.fastq} --no-unal --un-gz {output.nomap_fastq} -S {output.sam}"
# echo ${{cmd}} 1> {log.log}
# eval ${{cmd}} 1>> {log.log} 2> {log.err}
# """ % genome_db
#     else:
#         raise NotImplementedError("%s is not yet handled." % aligner)
#     return shell_commands


rule map_on_genome:
    input:
        fastq = OPJ(data_dir, "{lib}_{rep}.fastq.gz"),
    output:
        # temp because it uses a lot of space
        sam = temp(OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}.sam")),
        nomap_fastq = OPJ(output_dir, aligner, f"not_mapped_{genome}", f"{{lib}}_{{rep}}_unmapped_on_{genome}.fastq.gz"),
    params:
        aligner = aligner,
        index = genome_db,
        settings = alignment_settings[aligner],
    message:
        f"Mapping {{wildcards.lib}}_{{wildcards.rep}} on {genome}."
    log:
        log = OPJ(log_dir, "map_on_genome_{lib}_{rep}.log"),
        err = OPJ(log_dir, "map_on_genome_{lib}_{rep}.err")
    resources:
        io=5
    threads:
        4
    #shell:
    #    mapping_command(aligner)
    wrapper:
        "file:///pasteur/homes/bli/src/bioinfo_utils/snakemake_wrappers/map_on_genome"


rule extract_nomap_polyU:
    f"""
    Extract from the non-mappers those that end with a polyU tail
    of length from {minU} to {maxU}."""
    input:
        nomap_fastq = OPJ(output_dir, aligner, f"not_mapped_{genome}", f"{{lib}}_{{rep}}_unmapped_on_{genome}.fastq.gz"),
    output:
        nomap_polyU = OPJ(output_dir, aligner, f"not_mapped_{genome}", f"{{lib}}_{{rep}}_unmapped_on_{genome}_polyU.fastq.gz"),
        nomap_polyU_stats = OPJ(output_dir, aligner, f"not_mapped_{genome}", f"{{lib}}_{{rep}}_unmapped_on_{genome}_polyU_stats.txt"),
    run:
        stats = defaultdict(int)
        with gzopen(output.nomap_polyU, "wb") as fq_out:
            for (name, seq, qual) in fastx_read(input.nomap_fastq):
                trimmed = seq.rstrip("T")
                nb_T = len(seq) - len(trimmed)
                stats[nb_T] += 1
                if minU <= nb_T <= maxU:
                    # Find number of ending Ts
                    fq_out.write(f"@{name}\n{trimmed}\n+\n{qual:.{len(trimmed)}}\n".encode("utf-8"))
        with open(output.nomap_polyU_stats, "w") as stats_file:
            tot = 0
            for (nb_T, stat) in sorted(stats.items()):
                stats_file.write(f"{nb_T}\t{stat}\n")
                tot += stat
            stats_file.write(f"tot\t{tot}\n")


rule remap_on_genome:
    input:
        fastq = rules.extract_nomap_polyU.output.nomap_polyU,
    output:
        # temp because it might use a lot of space
        sam = temp(OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_polyU_on_{genome}.sam")),
        nomap_fastq = OPJ(output_dir, aligner, f"not_mapped_{genome}", f"{{lib}}_{{rep}}_polyU_unmapped_on_{genome}.fastq.gz"),
    params:
        aligner = aligner,
        index = genome_db,
        settings = alignment_settings[aligner],
    message:
        f"Remapping {{wildcards.lib}}_{{wildcards.rep}} polyU-trimmed unmapped reads on {genome}."
    log:
        log = OPJ(log_dir, "remap_on_genome_{lib}_{rep}.log"),
        err = OPJ(log_dir, "remap_on_genome_{lib}_{rep}.err")
    resources:
        io=5
    threads:
        4
    #shell:
    #    mapping_command(aligner)
    wrapper:
        "file:///pasteur/homes/bli/src/bioinfo_utils/snakemake_wrappers/map_on_genome"


def source_sam(wildcards):
    mapped_type = wildcards.mapped_type
    if mapped_type == f"on_{genome}":
        return rules.map_on_genome.output.sam
    elif mapped_type == f"polyU_on_{genome}":
        return rules.remap_on_genome.output.sam
    else:
        raise NotImplementedError(f"{mapped_type} not handled (yet?)")


rule sam2indexedbam:
    input:
        # sam = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}.sam"),
        sam = source_sam,
    output:
        # sorted_bam = protected(OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam")),
        # index = protected(OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam.bai")),
        sorted_bam = protected(OPJ(output_dir, aligner, f"mapped_{genome}", "{lib}_{rep}_{mapped_type}_sorted.bam")),
        index = protected(OPJ(output_dir, aligner, f"mapped_{genome}", "{lib}_{rep}_{mapped_type}_sorted.bam.bai")),
    message:
        "Sorting and indexing sam file for {wildcards.lib}_{wildcards.rep}."
    log:
        log = OPJ(log_dir, "sam2indexedbam_{lib}_{rep}.log"),
        err = OPJ(log_dir, "sam2indexedbam_{lib}_{rep}.err"),
    resources:
        io=45
    threads:
        8
    wrapper:
        "file:///pasteur/homes/bli/src/bioinfo_utils/snakemake_wrappers/sam2indexedbam"


rule compute_mapping_stats:
    input:
        sorted_bam = rules.sam2indexedbam.output.sorted_bam,
        #sorted_bam = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam"),
    output:
        stats = OPJ(output_dir, aligner, f"mapped_{genome}", "{lib}_{rep}_{mapped_type}_samtools_stats.txt"),
    shell:
        """samtools stats {input.sorted_bam} > {output.stats}"""


rule get_nb_mapped:
    """Extracts the number of mapped reads from samtools stats results."""
    input:
        stats = rules.compute_mapping_stats.output.stats,
        #stats = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_samtools_stats.txt"),
    output:
        nb_mapped = OPJ(output_dir, aligner, f"mapped_{genome}", "{lib}_{rep}_{mapped_type}_nb_mapped.txt"),
    shell:
        """samstats2mapped {input.stats} {output.nb_mapped}"""


rule compute_coverage:
    input:
        sorted_bam = rules.sam2indexedbam.output.sorted_bam,
        #sorted_bam = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam"),
    output:
        coverage = OPJ(output_dir, aligner, f"mapped_{genome}", "{lib}_{rep}_{mapped_type}_coverage.txt"),
    params:
        genomelen = genomelen,
    shell:
        """
        bases=$(samtools depth {input.sorted_bam} | awk '{{sum += $3}} END {{print sum}}') || error_exit "samtools depth failed"
        python3 -c "print(${{bases}} / {params.genomelen})" > {output.coverage}
        """


rule assemble_transcripts:
    input:
        sorted_bam = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam"),
    output:
        gtf = OPJ(output_dir, aligner, f"mapped_{genome}", "stringtie", f"{{lib}}_{{rep}}_on_{genome}.gtf"),
    params:
        annot = OPJ(annot_dir, "genes.gtf"),
    message:
        "Assembling transcripts for {wildcards.lib}_{wildcards.rep} with stringtie."
    log:
        log = OPJ(log_dir, "assemble_transcripts_{lib}_{rep}.log"),
        err = OPJ(log_dir, "assemble_transcripts_{lib}_{rep}.err")
    resources:
        io=5
    threads:
        4
    shell:
        """
        cmd="stringtie -p {threads} -G {params.annot} -o {output} -l {wildcards.lib}_{wildcards.rep} {input}"
        echo ${{cmd}} 1> {log.log}
        eval ${{cmd}} 1>> {log.log} 2> {log.err}
        """


rule merge_transcripts:
    input:
        expand(OPJ(output_dir, aligner, f"mapped_{genome}", "stringtie", f"{{lib}}_{{rep}}_on_{genome}.gtf"), lib=LIBS, rep=REPS),
    output:
        merged = OPJ(transcriptomes_dir, aligner, f"mapped_{genome}", "stringtie", "merged.gtf"),
    params:
        annot = OPJ(annot_dir, "genes.gtf"),
    message:
        "Merging transcripts with stringtie."
    log:
        log = OPJ(log_dir, "merge_transcripts.log"),
        err = OPJ(log_dir, "merge_transcripts.err")
    resources:
        io=5
    shell:
        """
        cmd="stringtie --merge -G {params.annot} -o {output} {input}"
        echo ${{cmd}} 1> {log.log}
        eval ${{cmd}} 1>> {log.log} 2> {log.err}
        """


rule gffcompare:
    input:
        merged = OPJ(transcriptomes_dir, aligner, f"mapped_{genome}", "stringtie", "merged.gtf"),
    output:
        stats = OPJ(transcriptomes_dir, aligner, f"mapped_{genome}", "stringtie", "compare.stats"),
    params:
        compare = OPJ(transcriptomes_dir, aligner, f"mapped_{genome}", "stringtie", "compare"),
        annot = OPJ(annot_dir, "genes.gtf"),
    message:
        "Merging transcripts with stringtie."
    log:
        log = OPJ(log_dir, "gffcompare.log"),
        err = OPJ(log_dir, "gffcompare.err")
    resources:
        io=5
    shell:
        """
        cmd="gffcompare -r {params.annot} -o {params.compare} {input.merged}"
        echo ${{cmd}} 1> {log.log}
        eval ${{cmd}} 1>> {log.log} 2> {log.err}
        """


rule quantify_transcripts:
    input:
        merged = OPJ(transcriptomes_dir, aligner, f"mapped_{genome}", "stringtie", "merged.gtf"),
        sorted_bam = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam"),
    output:
        gtf = OPJ(output_dir, aligner, f"mapped_{genome}", "ballgown", f"{{lib}}_{{rep}}_on_{genome}.gtf"),
    message:
        "Quantifying transcripts for {wildcards.lib}_{wildcards.rep} with stringtie."
    log:
        OPJ(log_dir, "quantify_transcripts_{lib}_{rep}.log")
    threads:
        4
    resources:
        io=5
    shell:
        """
        cmd="stringtie -e -B -p {threads} -G {input.merged} -o {output} {input.sorted_bam}"
        echo ${{cmd}} 1> {log}
        eval ${{cmd}} 1>> {log}
        """


# def htseq_orientation2stranded(wildcards):
#     orientation = wildcards.orientation
#     if orientation == "fwd":
#         if LIB_TYPE[-2:] == "SF":
#             return "yes"
#         elif LIB_TYPE[-2:] == "SR":
#             return "reverse"
#         else:
#             raise ValueError(f"{LIB_TYPE} library type not compatible with strand-aware read counting.")
#     elif orientation == "rev":
#         if LIB_TYPE[-2:] == "SF":
#             return "reverse"
#         elif LIB_TYPE[-2:] == "SR":
#             return "yes"
#         else:
#             raise ValueError(f"{LIB_TYPE} library type not compatible with strand-aware read counting.")
#     elif orientation == "all":
#         return "no"
#     else:
#         exit("Orientation is to be among \"fwd\", \"rev\" and \"all\".")


def biotype2annot(wildcards):
    if wildcards.biotype.endswith("_rmsk_families"):
        biotype = wildcards.biotype[:-9]
    else:
        biotype = wildcards.biotype
    return OPJ(annot_dir, f"{biotype}.gtf")


# rule htseq_count_reads:
#     input:
#         sorted_bam = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam"),
#         bai = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam.bai"),
#     output:
#         counts = OPJ(output_dir, aligner, f"mapped_{genome}", "htseq_count", f"{{lib}}_{{rep}}_on_{genome}", "{biotype}_{orientation}_counts.txt"),
#         counts_converted = OPJ(output_dir, aligner, f"mapped_{genome}", "htseq_count", f"{{lib}}_{{rep}}_on_{genome}", "{biotype}_{orientation}_counts_gene_names.txt"),
#     params:
#         stranded = htseq_orientation2stranded,
#         mode = "union",
#         annot = biotype2annot,
#     message:
#         "Counting {wildcards.orientation} {wildcards.biotype} reads for {wildcards.lib}_{wildcards.rep} with htseq-count."
#     benchmark:
#         OPJ(log_dir, "htseq_count_reads", "{lib}_{rep}_{biotype}_{orientation}_benchmark.txt")
#     log:
#         log = OPJ(log_dir, "htseq_count_reads", "{lib}_{rep}_{biotype}_{orientation}.log"),
#         err = OPJ(log_dir, "htseq_count_reads", "{lib}_{rep}_{biotype}_{orientation}.err"),
#     wrapper:
#         "file:///pasteur/homes/bli/src/bioinfo_utils/snakemake_wrappers/htseq_count_reads"


rule feature_count_reads:
    input:
        sorted_bam = rules.sam2indexedbam.output.sorted_bam,
        index = rules.sam2indexedbam.output.index,
    output:
        counts = OPJ(output_dir, aligner, f"mapped_{genome}", "feature_count", "{lib}_{rep}_{mapped_type}", "{biotype}_{orientation}_counts.txt"),
        counts_converted = OPJ(output_dir, aligner, f"mapped_{genome}", "feature_count", "{lib}_{rep}_{mapped_type}", "{biotype}_{orientation}_counts_gene_names.txt"),
    params:
        stranded = feature_orientation2stranded(LIB_TYPE),
        annot = biotype2annot,
        # pickled dictionary that associates gene ids to gene names
        converter = genome_dict["converter"]
    message:
        "Counting {wildcards.orientation} {wildcards.biotype} reads for {wildcards.lib}_{wildcards.rep} {wildcards.mapped_type} with featureCounts."
    benchmark:
        OPJ(log_dir, "feature_count_reads", "{lib}_{rep}_{mapped_type}_{biotype}_{orientation}_benchmark.txt")
    log:
        log = OPJ(log_dir, "feature_count_reads", "{lib}_{rep}_{mapped_type}_{biotype}_{orientation}.log"),
        err = OPJ(log_dir, "feature_count_reads", "{lib}_{rep}_{mapped_type}_{biotype}_{orientation}.err"),
    shell:
        """
        tmpdir=$(mktemp -dt "feature_{wildcards.lib}_{wildcards.rep}_{wildcards.mapped_type}_{wildcards.biotype}_{wildcards.orientation}.XXXXXXXXXX")
        cmd="featureCounts -a {params.annot} -o {output.counts} -t transcript -g "gene_id" -O -M --primary -s {params.stranded} --fracOverlap 0 --tmpDir ${{tmpdir}} {input.sorted_bam}"
        featureCounts -v 2> {log.log}
        echo ${{cmd}} 1>> {log.log}
        eval ${{cmd}} 1>> {log.log} 2> {log.err} || error_exit "featureCounts failed"
        rm -rf ${{tmpdir}}
        cat {output.counts} | id2name.py {params.converter} > {output.counts_converted}
        """


def biotype2merged(wildcards):
    if wildcards.biotype.endswith("_rmsk_families"):
        biotype = wildcards.biotype[:-9]
    else:
        biotype = wildcards.biotype
    return OPJ(annot_dir, f"{biotype}_merged.bed")


def intersect_orientation2stranded(wildcards):
    orientation = wildcards.orientation
    if orientation == "fwd":
        if LIB_TYPE[-2:] == "SF":
            return "-s"
        elif LIB_TYPE[-2:] == "SR":
            return "-S"
        else:
            raise ValueError(f"{LIB_TYPE} library type not compatible with strand-aware read counting.")
    elif orientation == "rev":
        if LIB_TYPE[-2:] == "SF":
            return "-S"
        elif LIB_TYPE[-2:] == "SR":
            return "-s"
        else:
            raise ValueError(f"{LIB_TYPE} library type not compatible with strand-aware read counting.")
    elif orientation == "all":
        return ""
    else:
        exit("Orientation is to be among \"fwd\", \"rev\" and \"all\".")


rule intersect_count_reads:
    input:
        sorted_bam = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam"),
        bai = OPJ(output_dir, aligner, f"mapped_{genome}", f"{{lib}}_{{rep}}_on_{genome}_sorted.bam.bai"),
    output:
        counts = OPJ(output_dir, aligner, f"mapped_{genome}", "intersect_count", f"{{lib}}_{{rep}}_on_{genome}", "{biotype}_{orientation}_counts.txt"),
        counts_converted = OPJ(output_dir, aligner, f"mapped_{genome}", "intersect_count", f"{{lib}}_{{rep}}_on_{genome}", "{biotype}_{orientation}_counts_gene_names.txt"),
    params:
        stranded = intersect_orientation2stranded,
        annot = biotype2merged,
    message:
        "Counting {wildcards.orientation} {wildcards.biotype} reads for {wildcards.lib}_{wildcards.rep} with bedtools intersect."
    benchmark:
        OPJ(log_dir, "intersect_count_reads", "{lib}_{rep}_{biotype}_{orientation}_benchmark.txt")
    log:
        log = OPJ(log_dir, "intersect_count_reads", "{lib}_{rep}_{biotype}_{orientation}.log"),
        err = OPJ(log_dir, "intersect_count_reads", "{lib}_{rep}_{biotype}_{orientation}.err"),
    threads: 4  # to limit memory usage, actually
    wrapper:
        "file:///pasteur/homes/bli/src/bioinfo_utils/snakemake_wrappers/intersect_count_reads"


rule summarize_counts:
    """For a given library, write a summary of the read counts for the various biotypes."""
    input:
        biotype_counts_files = expand(OPJ(output_dir, aligner, f"mapped_{genome}", "{{counter}}", "{{lib}}_{{rep}}_{{mapped_type}}", "{biotype}_{{orientation}}_counts.txt"), biotype=BIOTYPES),
    output:
        summary = OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}", "summaries", "{lib}_{rep}_{mapped_type}_{orientation}_counts.txt")
    run:
        if wildcards.counter == "htseq_count":
            sum_counter = sum_htseq_counts
        elif wildcards.counter == "intersect_count":
            sum_counter = sum_intersect_counts
        elif wildcards.counter == "feature_count":
            sum_counter = sum_feature_counts
        else:
            raise NotImplementedError(f"{wildcards.counter} not handled (yet?)")
        with open(output.summary, "w") as summary_file:
            header = "\t".join(BIOTYPES)
            summary_file.write("%s\n" % header)
            sums = "\t".join((str(sum_counter(counts_file)) for counts_file in input.biotype_counts_files))
            summary_file.write("%s\n" % sums)


rule gather_read_counts_summaries:
    """Gather read count summaries across libraries: lib_rep -> all."""
    input:
        #summary_tables = expand(OPJ(output_dir, aligner, f"mapped_{genome}", "{{counter}}", "summaries", f"{{lib}}_{{rep}}_on_{genome}_{{{orientation}}}_counts.txt"), lib=LIBS, rep=REPS),
        summary_tables = expand(OPJ(
            output_dir, aligner, f"mapped_{genome}", "{{counter}}", "summaries",
            "{name}_on_%s_{{orientation}}_counts.txt" % genome), name=COND_NAMES),
    output:
        summary_table = OPJ(
            output_dir, aligner, f"mapped_{genome}", "{counter}", "summaries",
            f"all_on_{genome}_{{orientation}}_counts.txt"),
    run:
        #summary_files = (OPJ(
        summary_files = [OPJ(
            output_dir, aligner, f"mapped_{genome}", wildcards.counter, "summaries",
            f"{{name}}_on_{genome}_{{orientation}}_counts.txt".format(
                name=cond_name,
                orientation=wildcards.orientation)) for cond_name in COND_NAMES]
                #orientation=wildcards.orientation)) for cond_name in COND_NAMES)
        try:
            summaries = pd.concat((pd.read_table(summary_file).T.astype(int) for summary_file in summary_files), axis=1)
        except ValueError:
            warnings.warn("Some data may be missing. Using floats instead of ints.\n")
            summaries = pd.concat((pd.read_table(summary_file).T.astype(float) for summary_file in summary_files), axis=1)
        summaries.columns = COND_NAMES
        summaries.to_csv(output.summary_table, na_rep="NA", sep="\t")


rule gather_counts:
    """For a given biotype, gather counts from all libraries in one table."""
    input:
        counts_tables = expand(OPJ(output_dir, aligner, f"mapped_{genome}", "{{counter}}", f"{{lib}}_{{rep}}_on_{genome}", "{{biotype}}_{{orientation}}_counts.txt"), lib=LIBS, rep=REPS),
    output:
        counts_table = OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}", f"all_on_{genome}", "{biotype}_{orientation}_counts.txt"),
    wildcard_constraints:
        # Avoid ambiguity with join_all_counts
        biotype = "|".join(BIOTYPES)
    run:
        # Gathering the counts data
        ############################
        #counts_files = (OPJ(
        counts_files = [OPJ(
            output_dir,
            aligner,
            f"mapped_{genome}",
            wildcards.counter,
            f"{cond_name}_on_{genome}",
            "{biotype}_{orientation}_counts.txt".format(biotype=wildcards.biotype, orientation=wildcards.orientation)) for cond_name in COND_NAMES]
            #"{biotype}_{orientation}_counts.txt".format(biotype=wildcards.biotype, orientation=wildcards.orientation)) for cond_name in COND_NAMES)
        if wildcards.counter == "htseq_count":
            counts_data = pd.concat(
                map(read_htseq_counts, counts_files),
                axis=1).fillna(0).astype(int)
        elif wildcards.counter == "intersect_count":
            counts_data = pd.concat(
                map(read_intersect_counts, counts_files),
                axis=1).fillna(0).astype(int)
        elif wildcards.counter == "feature_count":
            counts_data = pd.concat(
                map(read_feature_counts, counts_files),
                axis=1).fillna(0).astype(int)
        else:
            raise NotImplementedError(f"{wildcards.counter} not handled (yet?)")
        counts_data.columns = COND_NAMES
        # Simple_repeat|Simple_repeat|(TTTTTTG)n:1
        # Simple_repeat|Simple_repeat|(TTTTTTG)n:2
        # Simple_repeat|Simple_repeat|(TTTTTTG)n:3
        # Simple_repeat|Simple_repeat|(TTTTTTG)n:4
        # -> Simple_repeat|Simple_repeat|(TTTTTTG)n
        if wildcards.biotype.endswith("_rmsk_families"):
            repeat_families = [":".join(name.split(":")[:-1]) for name in counts_data.index]
            # Sum the counts for a given repeat family
            counts_data = counts_data.assign(family=repeat_families).groupby("family").sum()
        counts_data.index.names = ["gene"]
        counts_data.to_csv(output.counts_table, na_rep="NA", sep="\t")


# TODO: plot log(pseudocounts) vs log(concentration) as in figs sup8 and sup9 of Risso et al. 2014
# Colour-code the spike-in length
# Also do the plot with RPK instead of pseudocounts
rule compute_RPK:
    """For a given biotype, compute the corresponding RPK value (reads per kilobase)."""
    input:
        # TODO: Why wildcards seems to be None?
        #counts_data = rules.gather_counts.output.counts_table,
        counts_data = OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{biotype}_{orientation}_counts.txt"),
    output:
        rpk_file = OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{biotype}_{orientation}_RPK.txt"),
    params:
        feature_lengths_file = OPJ(annot_dir, "union_exon_lengths.txt"),
    run:
        counts_data = pd.read_table(input.counts_data, index_col="gene")
        feature_lengths = pd.read_table(params.feature_lengths_file, index_col="gene")
        common = counts_data.index.intersection(feature_lengths.index)
        rpk = 1000 * counts_data.loc[common].div(feature_lengths.loc[common]["union_exon_len"], axis="index")
        rpk.to_csv(output.rpk_file, sep="\t")


rule compute_sum_million_RPK:
    input:
        rpk_file = rules.compute_RPK.output.rpk_file,
    output:
        sum_rpk_file = OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{biotype}_{orientation}_sum_million_RPK.txt"),
    run:
        sum_rpk = pd.read_table(
            input.rpk_file,
            index_col=0).sum()
        (sum_rpk / 1000000).to_csv(output.sum_rpk_file, sep="\t")


# Compute TPM using total number of mappers divided by genome length
# (not sum of RPK across biotypes: some mappers may not be counted)
# No, doesn't work: mappers / genome length not really comparable
# Needs to be done on all_types
rule compute_TPM:
    """For a given biotype, compute the corresponding TPM value (reads per kilobase per million mappers)."""
    input:
        rpk_file = rules.compute_RPK.output.rpk_file
    output:
        tpm_file = OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{biotype}_{orientation}_TPM.txt"),
    # The sum must be done over all counted features
    wildcard_constraints:
        biotype = "alltypes"
    run:
        rpk = pd.read_table(input.rpk_file, index_col="gene")
        tpm = 1000000 * rpk / rpk.sum()
        tpm.to_csv(output.tpm_file, sep="\t")


# Useful to compute translation efficiency in the Ribo-seq pipeline
rule compute_mean_TPM:
    input:
        all_tmp_file = rules.compute_TPM.output.tpm_file
    output:
        tpm_file = OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            "{lib}_mean_on_%s" % genome, "{biotype}_{orientation}_TPM.txt"),
    run:
        # TODO
        tpm = pd.read_table(
            input.all_tmp_file, index_col="gene",
            usecols=["gene", *[f"{wildcards.lib}_{rep}" for rep in REPS]])
        tpm_mean = tpm.mean(axis=1)
        tpm_mean.columns = [wildcards.lib]
        tpm_mean.to_csv(output.tpm_file, sep="\t")


rule compute_RPM:
    input:
        counts_data = OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{biotype}_{orientation}_counts.txt"),
        #summary_table = rules.gather_read_counts_summaries.output.summary_table,
        summary_table = OPJ(
            output_dir, aligner, f"mapped_{genome}", "{counter}", "summaries",
            f"all_on_{genome}_fwd_counts.txt"),
    output:
        rpm_file = OPJ(output_dir, aligner, f"mapped_{genome}", "{counter}",
            f"all_on_{genome}", "{biotype}_{orientation}_RPM.txt"),
    run:
        # Reading column counts from {input.counts_table}