# Adapted from Sebastian Mueller's Snakemake workflow for mapping small RNA-seq 
# libraries to a reference genome
# https://github.com/seb-mueller/snakemake_sRNAseq

# Chromosome sizes file below ("data/index/genome.fa.sizes") must exist
# before running Snakefile
# e.g., in "data/index/" run:
# samtools faidx genome.fa; cut -f1,2 genome.fa.fai > genome.fa.sizes

# Usage (snakemake --cores should reflect available cores):
# conda env create --file environment.yaml --name srna_mapping
# source activate srna_mapping
# snakemake -p --cores 48
# source deactivate

import pandas as pd
import os

# To make the per_base_coverage rule work with a shell script invoked using the "shell" directive,
# we need to determine the base path of Snakefile since we expect the scripts directory to be there as well
SRCDIR = srcdir("")

# Specify config file parameters
configfile: "config.yaml"
samples    = config["SAMPLES"]
mismatches = config["MAPPING"]["mismatches"]
reference  = config["MAPPING"]["reference"]
refbase    = os.path.basename(reference)
mode       = config["MAPPING"]["mode"]
sRNAsize   = config["sRNA_SIZE"]
genomeBinName = config["COVERAGE"]["genomeBinName"]

# Determine bam index format (bai or csi) based on chromosome sizes
# Genomes with chromosomes longer than ~500 Mb (e.g., in wheat) require a csi index
# E.g., in axolotl: https://sourceforge.net/p/samtools/mailman/message/36249039/
chrSizes = pd.read_table("data/index/" + refbase + ".fa.sizes",
                         header = None)
smallChrs = 0
for x in chrSizes[1]:
    if x < 5e+08:
        smallChrs = smallChrs + 1

if smallChrs < len(chrSizes[1]):
    bamidx = "csi"
else:
    bamidx = "bai"

# Set bowtie mapping parameters based on "mode" specified in config.yaml
# (report uniquely aligned reads or both uniquely and multiply aligned reads)
if mode == "unique":
    bowtie_params = "-m 1"
elif mode == "both":
    bowtie_params = "--best"
else:
    bowtie_params = ""

# Specify the desired end target file(s)
rule all:
    input:
        expand("logs/fastqc/raw/{sample}_R1_fastqc.html",
               sample = samples),
        expand("data/trimmed/{sample}_R1_trimmed.fastq.gz",
               sample = samples),
        expand("logs/fastqc/trimmed/{sample}_R1_trimmed_fastqc.html",
               sample = samples),
        expand("mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam",
               mode = mode,
               sample = samples,
               refbase = refbase),
        expand("mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam",
               mode = mode,
               sample = samples,
               refbase = refbase,
               sRNAsize = sRNAsize),
        expand("mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam.{bamidx}",
               mode = mode,
               sample = samples,
               refbase = refbase,
               bamidx = bamidx),
        expand("mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam.{bamidx}",
               mode = mode,
               sample = samples,
               refbase = refbase,
               sRNAsize = sRNAsize,
               bamidx = bamidx),
        expand("mapped/{mode}/bw/{sample}_MappedOn_{refbase}_{mode}_sort_norm.bw",
               mode = mode,
               sample = samples,
               refbase = refbase),
        expand("mapped/{mode}/bw/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm.bw",
               mode = mode,
               sample = samples,
               refbase = refbase,
               sRNAsize = sRNAsize),
        expand("mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_sort_norm.bedgraph",
               mode = mode,
               sample = samples,
               refbase = refbase),
        expand("mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm.bedgraph",
               mode = mode,
               sample = samples,
               refbase = refbase,
               sRNAsize = sRNAsize),
        expand("mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_sort_norm_binSize{genomeBinName}.bedgraph",
               mode = mode,
               sample = samples,
               refbase = refbase,
               genomeBinName = genomeBinName),
        expand("mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm_binSize{genomeBinName}.bedgraph",
               mode = mode,
               sample = samples,
               refbase = refbase,
               sRNAsize = sRNAsize,
               genomeBinName = genomeBinName)

# Run fastqc on R1 raw data
rule fastqc_R1_raw:
    """Create fastqc report"""
    input:
        "data/{sample}_R1.fastq.gz"
    output:
        html = "logs/fastqc/raw/{sample}_R1_fastqc.html",
        zip  = "logs/fastqc/raw/{sample}_R1_fastqc.zip"
    params: "--extract"
    log:
        "logs/fastqc/raw/{sample}_R1.log"
    wrapper:
        "0.27.1/bio/fastqc"

# Trim off adapters
rule cutadapt:
    """Remove adapters"""
    input:
        "data/{sample}_R1.fastq.gz"
    output:
        fastq = "data/trimmed/{sample}_R1_trimmed.fastq.gz",
        qc    = "data/trimmed/{sample}_trimmed.qc.txt"
    params:
        " -a " +     config["FILTER"]["cutadapt"]["adapter"] +
        " -q " + str(config["FILTER"]["cutadapt"]["quality-filter"]) +
        " -m " + str(config["FILTER"]["cutadapt"]["minimum-length"]) +
        " -M " + str(config["FILTER"]["cutadapt"]["maximum-length"]) +
        " -O " + str(config["FILTER"]["cutadapt"]["minimum-overlap"]) +
        " --cores=0"
    log:
        "logs/cutadapt/{sample}_trimmed.log"
    wrapper:
        "0.27.1/bio/cutadapt/se"

# Run fastqc on R1 trimmed data
rule fastqc_R1_trimmed:
    """Create fastqc report"""
    input:
        "data/trimmed/{sample}_R1_trimmed.fastq.gz"
    output:
        html = "logs/fastqc/trimmed/{sample}_R1_trimmed_fastqc.html",
        zip  = "logs/fastqc/trimmed/{sample}_R1_trimmed_fastqc.zip"
    params: "--extract"
    log:
        "logs/fastqc/trimmed/{sample}_R1_trimmed.log"
    wrapper:
        "0.27.1/bio/fastqc"

# Align to reference genome
rule bowtie:
    """Map small RNAs using bowtie and sort them using samtools"""
    input:
        "data/trimmed/{sample}_R1_trimmed.fastq.gz"
    output:
        protected("mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam")
    log:
        "logs/bowtie/{sample}_MappedOn_{refbase}_{mode}_sort.log"
    params:
        sortMemory = config["MAPPING"]["sortMemory"]
    threads: config["THREADS"]  
    shell:
        "(bowtie {reference} --threads {threads} -v {mismatches}"
        " {bowtie_params} -q {input} -S "
        "| samtools view -u -F 4 - "
        "| samtools sort -@ {threads} -m {params.sortMemory} -o {output}) 2> {log}"

# Extract differently sized sRNA populations as BAM files
rule sRNAsize_extract:
    """Extract differently sized sRNA populations as BAM files"""
    input:
        "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam"
    output:
        protected("mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam")
    params:
        sortMemory = config["MAPPING"]["sortMemory"]
    threads: config["THREADS"]
    log:
        "logs/samtools/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.log"
    shell:
        "(samtools view -h {input} "
        "| awk 'BEGIN{{OFS=FS=\"\\t\"}}/^@/{{print}}!/^@/{{if(length($10)=={sRNAsize}) print}}' "
        "| samtools view -u - "
        "| samtools sort -@ {threads} -m {params.sortMemory} -o {output}) 2> {log}"

# Postmapping steps:
# Index BAM files (index format [bai or csi] depends on chromosome sizes)
# Generate samtools flagstat and idxstats
# Calculate library-size-normalized coverage
if bamidx == "bai":
    rule postmapping:
        """bam.bai samtools flagstat idxstats"""
        input:
            mapped = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam"
        output:
            mapped = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam.{bamidx}"
        log:
            mappedflagstat = "logs/samtools/stats/{sample}_MappedOn_{refbase}_{mode}_sort_flagstat.log",
            mappedidxstats = "logs/samtools/stats/{sample}_MappedOn_{refbase}_{mode}_sort_idxstats.log"
        shell:
            """
            samtools index    {input.mapped}
            samtools flagstat {input.mapped} > {log.mappedflagstat}
            samtools idxstats {input.mapped} > {log.mappedidxstats}
            """
    rule calc_coverage:
        """Calculate library-size-normalized coverage"""
        input:
            mappedBAM    = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam",
            mappedBAMidx = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam.bai"
        output:
            mappedBW = "mapped/{mode}/bw/{sample}_MappedOn_{refbase}_{mode}_sort_norm.bw",
            mappedBG = "mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_sort_norm.bedgraph"
        params:
            normalizeUsing         = config["COVERAGE"]["normalizeUsing"],
            ignoreForNormalization = config["COVERAGE"]["ignoreForNormalization"],
            binSize                = config["COVERAGE"]["binSize"]
        threads: config["THREADS"]  
        log:
            mapped = "logs/bamCoverage/{sample}_MappedOn_{refbase}_{mode}_sort_norm.log"
        shell:
            "(bamCoverage -b {input.mappedBAM} -o {output.mappedBW}"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.binSize} -p {threads}; "
            "bamCoverage -b {input.mappedBAM} -o {output.mappedBG} -of bedgraph"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.binSize} -p {threads}) 2> {log.mapped}"
    rule calc_coverage_genome:
        """Calculate library-size-normalized coverage in adjacent windows"""
        input:
            mappedBAM    = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam",
            mappedBAMidx = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam.bai"
        output:
            mappedBGgenome = "mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_sort_norm_binSize{genomeBinName}.bedgraph"
        params:
            normalizeUsing         = config["COVERAGE"]["normalizeUsing"],
            ignoreForNormalization = config["COVERAGE"]["ignoreForNormalization"],
            genomeBinSize          = config["COVERAGE"]["genomeBinSize"]
        threads: config["THREADS"]  
        log:
            mapped = "logs/bamCoverage/{sample}_MappedOn_{refbase}_{mode}_sort_norm_binSize{genomeBinName}.log"
        shell:
            "(bamCoverage -b {input.mappedBAM} -o {output.mappedBGgenome} -of bedgraph"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.genomeBinSize} -p {threads}) 2> {log.mapped}"
    rule postmapping_sRNAsize:
        """bam.bai samtools flagstat idxstats"""
        input:
            subset = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam"
        output:
            subset = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam.{bamidx}"
        log:
            subsetflagstat = "logs/samtools/stats/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_flagstat.log",
            subsetidxstats = "logs/samtools/stats/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_idxstats.log"
        shell:
            """
            samtools index    {input.subset}
            samtools flagstat {input.subset} > {log.subsetflagstat}
            samtools idxstats {input.subset} > {log.subsetidxstats}
            """
    rule calc_coverage_sRNAsize:
        """Calculate library-size-normalized coverage"""
        input:
            subsetBAM    = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam",
            subsetBAMidx = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam.bai"
        output:
            subsetBW = "mapped/{mode}/bw/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm.bw",
            subsetBG = "mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm.bedgraph"
        params:
            normalizeUsing         = config["COVERAGE"]["normalizeUsing"],
            ignoreForNormalization = config["COVERAGE"]["ignoreForNormalization"],
            binSize                = config["COVERAGE"]["binSize"]
        threads: config["THREADS"]  
        log:
            subset = "logs/bamCoverage/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm.log"
        shell:
            "(bamCoverage -b {input.subsetBAM} -o {output.subsetBW}"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.binSize} -p {threads}; "
            "bamCoverage -b {input.subsetBAM} -o {output.subsetBG} -of bedgraph"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.binSize} -p {threads}) 2> {log.subset}"
    rule calc_coverage_genome_sRNAsize:
        """Calculate library-size-normalized coverage in adjacent windows"""
        input:
            subsetBAM    = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam",
            subsetBAMidx = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam.bai"
        output:
            subsetBGgenome = "mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm_binSize{genomeBinName}.bedgraph"
        params:
            normalizeUsing         = config["COVERAGE"]["normalizeUsing"],
            ignoreForNormalization = config["COVERAGE"]["ignoreForNormalization"],
            genomeBinSize          = config["COVERAGE"]["genomeBinSize"]
        threads: config["THREADS"]  
        log:
            subset = "logs/bamCoverage/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm_binSize{genomeBinName}.log"
        shell:
            "(bamCoverage -b {input.subsetBAM} -o {output.subsetBGgenome} -of bedgraph"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.genomeBinSize} -p {threads}) 2> {log.subset}"
else:
    rule postmapping:
        """bam.csi samtools flagstat idxstats"""
        input:
            mapped = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam"
        output:
            mapped = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam.{bamidx}"
        log:
            mappedflagstat = "logs/samtools/stats/{sample}_MappedOn_{refbase}_{mode}_sort_flagstat.log",
            mappedidxstats = "logs/samtools/stats/{sample}_MappedOn_{refbase}_{mode}_sort_idxstats.log"
        shell:
            """
            samtools index -c -m 14 {input.mapped}
            samtools flagstat       {input.mapped} > {log.mappedflagstat}
            samtools idxstats       {input.mapped} > {log.mappedidxstats}
            """
    rule calc_coverage:
        """Calculate library-size-normalized coverage"""
        input:
            mappedBAM    = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam",
            mappedBAMidx = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam.csi"
        output:
            mappedBW = "mapped/{mode}/bw/{sample}_MappedOn_{refbase}_{mode}_sort_norm.bw",
            mappedBG = "mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_sort_norm.bedgraph"
        params:
            normalizeUsing         = config["COVERAGE"]["normalizeUsing"],
            ignoreForNormalization = config["COVERAGE"]["ignoreForNormalization"],
            binSize                = config["COVERAGE"]["binSize"]
        threads: config["THREADS"]  
        log:
            mapped = "logs/bamCoverage/{sample}_MappedOn_{refbase}_{mode}_sort_norm.log"
        shell:
            "(bamCoverage -b {input.mappedBAM} -o {output.mappedBW}"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.binSize} -p {threads}; "
            "bamCoverage -b {input.mappedBAM} -o {output.mappedBG} -of bedgraph"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.binSize} -p {threads}) 2> {log.mapped}"
    rule calc_coverage_genome:
        """Calculate library-size-normalized coverage in adjacent windows"""
        input:
            mappedBAM    = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam",
            mappedBAMidx = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_sort.bam.csi"
        output:
            mappedBGgenome = "mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_sort_norm_binSize{genomeBinName}.bedgraph"
        params:
            normalizeUsing         = config["COVERAGE"]["normalizeUsing"],
            ignoreForNormalization = config["COVERAGE"]["ignoreForNormalization"],
            genomeBinSize          = config["COVERAGE"]["genomeBinSize"]
        threads: config["THREADS"]  
        log:
            mapped = "logs/bamCoverage/{sample}_MappedOn_{refbase}_{mode}_sort_norm_binSize{genomeBinName}.log"
        shell:
            "(bamCoverage -b {input.mappedBAM} -o {output.mappedBGgenome} -of bedgraph"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.genomeBinSize} -p {threads}) 2> {log.mapped}"
    rule postmapping_sRNAsize:
        """bam.csi samtools flagstat idxstats"""
        input:
            subset = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam"
        output:
            subset = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam.{bamidx}"
        log:
            subsetflagstat = "logs/samtools/stats/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_flagstat.log",
            subsetidxstats = "logs/samtools/stats/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_idxstats.log"
        shell:
            """
            samtools index -c -m 14 {input.subset}
            samtools flagstat       {input.subset} > {log.subsetflagstat}
            samtools idxstats       {input.subset} > {log.subsetidxstats}
            """
    rule calc_coverage_sRNAsize:
        """Calculate library-size-normalized coverage"""
        input:
            subsetBAM    = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam",
            subsetBAMidx = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam.csi"
        output:
            subsetBW = "mapped/{mode}/bw/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm.bw",
            subsetBG = "mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm.bedgraph"
        params:
            normalizeUsing         = config["COVERAGE"]["normalizeUsing"],
            ignoreForNormalization = config["COVERAGE"]["ignoreForNormalization"],
            binSize                = config["COVERAGE"]["binSize"]
        threads: config["THREADS"]  
        log:
            subset = "logs/bamCoverage/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm.log"
        shell:
            "(bamCoverage -b {input.subsetBAM} -o {output.subsetBW}"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.binSize} -p {threads}; "
            "bamCoverage -b {input.subsetBAM} -o {output.subsetBG} -of bedgraph"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.binSize} -p {threads}) 2> {log.subset}"
    rule calc_coverage_genome_sRNAsize:
        """Calculate library-size-normalized coverage in adjacent windows"""
        input:
            subsetBAM    = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam",
            subsetBAMidx = "mapped/{mode}/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort.bam.csi"
        output:
            subsetBGgenome = "mapped/{mode}/bg/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm_binSize{genomeBinName}.bedgraph"
        params:
            normalizeUsing         = config["COVERAGE"]["normalizeUsing"],
            ignoreForNormalization = config["COVERAGE"]["ignoreForNormalization"],
            genomeBinSize          = config["COVERAGE"]["genomeBinSize"]
        threads: config["THREADS"]  
        log:
            subset = "logs/bamCoverage/{sample}_MappedOn_{refbase}_{mode}_{sRNAsize}nt_sort_norm_binSize{genomeBinName}.log"
        shell:
            "(bamCoverage -b {input.subsetBAM} -o {output.subsetBGgenome} -of bedgraph"
            " --normalizeUsing {params.normalizeUsing}"
            " --ignoreForNormalization {params.ignoreForNormalization}"
            " --binSize {params.genomeBinSize} -p {threads}) 2> {log.subset}"
