# vim: syntax=python tabstop=4 expandtab
# coding: utf-8


"""
ALPACA SNV calling.

For usage, include this in your workflow.

Expects a global variable config (e.g. parsed from json with configfile keyword)
of at least the following structure, assuming that the desired reference sequence
is some genome to be found under the given path, and two units A and B have been sequenced.
The heterozygosity parameter shows the default values for homo sapiens.

{
    "references": {
        "genome": "path/to/genome.fasta"
    },
    "samples": {
        "A": ["A"],
        "B": ["B"]
    },
    "units": {
        "A":
            ["path/to/A_R1.fastq.gz", "path/to/A_R2.fastq.gz"],
        "B":
            ["path/to/B.fastq.gz"]
    },
    "known_variants": {
        "dbsnp": "path/to/dbsnp.vcf"
    },
    "ploidy": 2,
    "heterozygosity": 0.001,
    "fdr": 0.05,
    "max_strandbias": 50,
    "alpaca_queries": {
        "somatic": "A - B"
    }
}

Note the separation between samples and units that allows to have more than
one sequencing run for each sample, or multiple lanes per sample.
"""


__author__ = "Johannes Köster (http://johanneskoester.bitbucket.org)"
__license__ = "MIT"


def _sample_units(pattern):
    def apply(wildcards):
        return expand(
            pattern, reference=wildcards.reference,
            unit=config["samples"][wildcards.sample])
    return apply


def _get_ref(wildcards):
    return config["references"][wildcards.reference]


rule alpaca_index:
    input:
        _sample_units("mapping/{reference}/units/{unit}.sorted.bam.bai"),
        bams=_sample_units("mapping/{reference}/units/{unit}.sorted.bam"),
        ref=_get_ref
    output:
        "snv_calling/{reference}/{sample}.index.hdf5"
    log:
        "snv_calling/log/{reference}/{sample}.index.log"
    threads: 8
    resources: gpu=1
    shell:
        "alpaca --dev gpu --threads {threads} index --ploidy {config[ploidy]} "
        "--sample-name {wildcards.sample} {input.ref} "
        "{input.bams} {output} 2> {log}"


rule alpaca_merge:
    input:
        _get_ref,
        expand(
            "snv_calling/{{reference}}/{sample}.index.hdf5",
            sample=config["samples"])
    output:
        "snv_calling/{reference}/index.hdf5"
    log:
        "snv_calling/log/{reference}/merge.log"
    resources: gpu=1
    shell:
        "alpaca --dev gpu merge {input} {output} 2> {log}"


rule alpaca_call:
    input:
        "snv_calling/{reference}/index.hdf5"
    output:
        "snv_calling/{reference}/{query}.vcf"
    log:
        "snv_calling/log/{reference}/{query}.call.log"
    params:
        query=lambda wildcards: config["alpaca_queries"][wildcards.query]
    resources: gpu=1
    shell:
        "alpaca --dev gpu call --fdr {config[fdr]} "
        "--max-strand-bias {config[max_strandbias]} "
        "--heterozygosity {config[heterozygosity]} "
        "{input} '{params.query}' > {output} 2> {log}"
