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


"""
GATK haplotype calling using the workflow suggested in 
https://www.broadinstitute.org/gatk/guide/topic?name=methods#methods3893
Needs at least GATK 3.0, available in your PATH as "gatk".

For usage, include this in your workflow.

Expects a global variable config
(see https://bitbucket.org/johanneskoester/snakemake/wiki/Documentation#markdown-header-configuration)
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 parameters show 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"
    },
    "heterozygosity": 0.001,
    "indel_heterozygosity": 1.25E-4
}

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 _get_ref(wildcards):
    return config["references"][wildcards.reference]


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


# def _gatk_multi_arg(flag, files):
#     flag += " "
#     return " ".join(flag + f for f in files)



"""
gatk_haplotype caller changes
--variant_index_parameter 128000 removed
--variant_index_type LINEAR --variant_index_parameter 128000 removed
--both not neccisary with current verison of GATK

removed from haplotype caller
--heterozygosity {config[heterozygosity]} \
--indel_heterozygosity {config[indel_heterozygosity]} \

"""

rule gatk_haplotype_caller:
    input:
        ref=config["references"],
        bams="gatk_recalibrated/{prefix}.recalibrated.bam"
    output:
        gvcf="gatk/variant_calling/{prefix}.gvcf",
        idx="gatk/variant_calling/{prefix}.gvcf.idx"
    params:
        interval=config.get("interval", "")
    log:
        "gatk/variant_calling/log/{prefix}.genotype_info.log"
    threads: 2  # haplotype caller currently does not reliably use more than 2 threads
    run:
        bams = _gatk_multi_arg("-I", input.bams)
        shell(
            """
            module load GATK
            java -Xmx8g -Djava.io.tmpdir=/lscratch/${{prefix}} -jar $GATK_HOME/GenomeAnalysisTK.jar
            -T HaplotypeCaller \
            -R {input.ref} \
            -I {input.bams} \
            -L {params.interval} \
            -emitRefConfidence GVCF \
            -dbsnp {config[known_variants][dbsnp]} \
            -nct {threads} \
            -o {output.gvcf} >& {log}
            """





"""
 collects all gvcf files (which will be in gatk/variant_calling/{unit}/{prefix}.recalibrated bams)
ouput is combined vcf in the root folder of gatk/variant_calling/
"""
rule gatk_genotyping:
    input:
        ref=config["references"],
        gvcfs=expand(
            "gatk/variant_calling/{samples}.gvcf",
            sample=config["samples"])
    output:
        "gatk/variant_calling/all.vcf"
    log:
        "gatk/variant_calling/log/all.genotype.log"
    threads: 16
    shell:
        """
        module load GATK
        java -Xmx8g -Djava.io.tmpdir=/lscratch/${{prefix}} -jar $GATK_HOME/GenomeAnalysisTK.jar \
        -T GenotypeGVCFs {gvcfs} \
        -V {input.gvcfs} \
        -nt {threads} \
        -R {input.ref} \
        --dbsnp {config[known_variants][dbsnp]} \
        --heterozygosity {config[heterozygosity]} \
        --indel_heterozygosity {config[indel_heterozygosity]} \
        -o {output} >& {log}
            
        """

