# The Snakemake tutorial's workflow as its "Basics" end it (docs: tutorial/basics, advanced),
# with a conda environment per step so it runs in a stock image. Its data comes from
# snakemake/snakemake-tutorial-data. One change: bwa_map lists the genome's bwa index, which bwa
# reads from next to genome.fa, as inputs (named, so the command passes only the genome and the
# reads). Locally the index is there anyway; a job on Batch only gets the inputs its rule names.
configfile: "config.yaml"


rule all:
    input:
        "plots/quals.svg",


def get_bwa_map_input_fastqs(wildcards):
    return config["samples"][wildcards.sample]


rule bwa_map:
    input:
        fa="data/genome.fa",
        reads=get_bwa_map_input_fastqs,
        index=multiext("data/genome.fa", ".amb", ".ann", ".bwt", ".pac", ".sa"),
    output:
        temp("mapped_reads/{sample}.bam"),
    params:
        rg=r"@RG\tID:{sample}\tSM:{sample}",
    log:
        "logs/bwa_mem/{sample}.log",
    threads: 8
    conda:
        "envs/mapping.yaml"
    shell:
        "(bwa mem -R '{params.rg}' -t {threads} {input.fa} {input.reads} | "
        "samtools view -Sb - > {output}) 2> {log}"


rule samtools_sort:
    input:
        "mapped_reads/{sample}.bam",
    output:
        protected("sorted_reads/{sample}.bam"),
    conda:
        "envs/mapping.yaml"
    shell:
        "samtools sort -T sorted_reads/{wildcards.sample} -O bam {input} > {output}"


rule samtools_index:
    input:
        "sorted_reads/{sample}.bam",
    output:
        "sorted_reads/{sample}.bam.bai",
    conda:
        "envs/mapping.yaml"
    shell:
        "samtools index {input}"


rule bcftools_call:
    input:
        fa="data/genome.fa",
        bam=expand("sorted_reads/{sample}.bam", sample=config["samples"]),
        bai=expand("sorted_reads/{sample}.bam.bai", sample=config["samples"]),
    output:
        "calls/all.vcf",
    params:
        rate=config["prior_mutation_rate"],
    log:
        "logs/bcftools_call/all.log",
    conda:
        "envs/calling.yaml"
    shell:
        "(bcftools mpileup -f {input.fa} {input.bam} | "
        "bcftools call -mv -P {params.rate} - > {output}) 2> {log}"


rule plot_quals:
    input:
        "calls/all.vcf",
    output:
        "plots/quals.svg",
    conda:
        "envs/plot.yaml"
    script:
        "scripts/plot-quals.py"
