shell:
    """
    LOGPATH="$PWD/{log}"
    date > "$LOGPATH"
    env >> "$LOGPATH"
    # Check if a proper conda environment is available
    if [ -z "${{CONDA_PREFIX+_}}" ] || [ ! -d "$CONDA_PREFIX"/opt/hla-la ]; then
        echo "--use-conda is required to run the HLA-LA rule." | tee -a "$LOGPATH" >&2
        exit 1
    fi

    # Create and prepare the temp directory
    rm -rf {params.tempdir:q}
    mkdir -p {params.tempdir:q}/src {params.tempdir:q}/bin >>"$LOGPATH"

    # PREPARE DIRECTORY SRC
    cd {params.tempdir:q}/src
    ln -s "$CONDA_PREFIX"/opt/hla-la/src/* . &>>"$LOGPATH"
    mv HLA-LA.pl HLA-LA.pl~
    cp "$CONDA_PREFIX"/opt/hla-la/src/HLA-LA.pl . &>>"$LOGPATH"
    cd ..

    # PREPARE DIRECTORY BIN
    ln -s "$CONDA_PREFIX"/opt/hla-la/bin/HLA-LA bin/ &>>"$LOGPATH"

    # PREPARE DIRECTORY GRAPHS
    ln -s ../../../{input.graphs:q} graphs &>>"$LOGPATH"

    # RUN MAIN PROGRAM
    ./src/HLA-LA.pl --sampleID sample --BAM ../../../{input.bam:q} --workingDir "$PWD" --graph PRG_MHC_GRCh38_withIMGT &>>"$LOGPATH"

    # MOVE RESULTS INTO PLACE
    cd ../../../
    rm -f {output.outdir:q} &>>"$LOGPATH"
    mv {params.tempdir:q}/sample {output.outdir:q} &>>"$LOGPATH"
    ln -s ../../{output.outdir:q}/hla/R1_bestguess_G.txt {output.main:q} &>>"$LOGPATH"

    # REMOVE TEMPORARY DIRECTORY
    rm -rf {params.tempdir:q} &>>"$LOGPATH"
    """

shell:
    """
    # Extract a list of read name that were aligned to HLA loci (HLA-A, B, C, DM, DO, DP, DQ, DR, E, F, G, H, J, K, L, P, V, MIC, and TAP)
    samtools view -@ {threads} {input.bam:q} chr6:29907037-29915661 chr6:31319649-31326989 chr6:31234526-31241863 \
        chr6:32914391-32922899 chr6:32900406-32910847 chr6:32969960-32979389 chr6:32778540-32786825 \
        chr6:33030346-33050555 chr6:33041703-33059473 chr6:32603183-32613429 chr6:32707163-32716664 \
        chr6:32625241-32636466 chr6:32721875-32733330 chr6:32405619-32414826 chr6:32544547-32559613 \
        chr6:32518778-32554154 chr6:32483154-32559613 chr6:30455183-30463982 chr6:29689117-29699106 \
        chr6:29792756-29800899 chr6:29793613-29978954 chr6:29855105-29979733 chr6:29892236-29899009 \
        chr6:30225339-30236728 chr6:31369356-31385092 chr6:31460658-31480901 chr6:29766192-29772202 \
        chr6:32810986-32823755 chr6:32779544-32808599 chr6:29756731-29767588 \
        2> {log:q} | awk '{{print $1}}' | LC_ALL=C sort | uniq > {output.locus_ids:q}

    echo "'samtools view' finished extracting chromosome 6 reads." >>{log:q}

    mkdir -p {output.idx_dir:q}

    java -jar -Xmx32g -Xms32g '{params[hla_vbseq_bam_name_index_jar]}' index {input.bam:q} --indexFile {output.idx_dir:q}/index &>>{log:q}
    echo "bamNameIndex.jar [1/2] finished." >>{log:q}

    java -jar '{params[hla_vbseq_bam_name_index_jar]}' search {input.bam:q} --indexFile {output.idx_dir:q}/index --name {output.locus_ids:q} --output {output.partial_sam:q} &>>{log:q}
    echo "bamNameIndex.jar [2/2] finished." >>{log:q}

    picard SamToFastq I={output.partial_sam:q} F={output.partial_fq_1:q} F2={output.partial_fq_2:q} &>>{log:q}
    echo "picard SamToFastq [1/2] finished." >>{log:q}

    samtools view -@ {threads} -bh -f 12 {input.bam:q} >{output.bam_unmapped:q} 2>>{log:q}
    echo "'samtools view' finished extracting unmapped reads." >>{log:q}

    picard SamToFastq I={output.bam_unmapped:q} F={output.unmapped_fq_1:q} F2={output.unmapped_fq_2:q} &>>{log:q}
    echo "picard SamToFastq [2/2] finished." >>{log:q}

    cat {output.partial_fq_1:q} {output.unmapped_fq_1:q} >{output.cat_fq_1} 2>>{log:q}
    cat {output.partial_fq_2:q} {output.unmapped_fq_2:q} >{output.cat_fq_2} 2>>{log:q}
    echo "Generating input FASTQ files finished." >>{log:q}

    bwa index '{params[hla_vbseq_ref_fa]}' &>>{log:q}
    bwa mem -t {threads} -P -L 10000 -a '{params[hla_vbseq_ref_fa]}' {output.cat_fq_1} {output.cat_fq_2} >{output.sam_final:q} 2>>{log:q}
    echo "Mapping against HLA reference finished." >>{log:q}

    # Run VBSeq
    java -jar '{params[hla_vbseq_main_jar]}'  '{params[hla_vbseq_ref_fa]}' {output.sam_final:q} {output.result_txt:q} --alpha_zero 0.01 --is_paired &>>{log:q}
    echo "VBSeq finished." >>{log:q}

    # Call alleles
    python '{params[hla_vbseq_call_hla_digits]}' -v {output.result_txt:q} -a '{params[hla_vbseq_ref_txt]}'  -r 140 -d 4 --ispaired >{output.result_tsv:q} 2>>{log:q}
    echo "Calling HLA alleles finished." >>{log:q}

    echo "" >>{log:q}
    echo "Rule 'hla_vbseq_typing' finished." >>{log:q}
    """

run:
    # Parse JSON output of xHLA
    import json
    with open(input[0], 'r') as infile:
        alleles = json.load(infile)['hla']['alleles']

    # Write multihla output
    import collections
    with open(output[0], 'w') as outfile:
        print('Dataset\tSample\tMethod\tOptions\tGene\tAllele1\tAllele2', file=outfile)
        data = collections.defaultdict(lambda : [])
        for allele in alleles:
            gene, allele = allele.split('*')
            data[gene].append(allele)
        for gene, alleles in data.items():
            if len(alleles) == 2:
                print(f'{wildcards.dataset}\t{wildcards.sample}\txHLA\t{params.opts}\t{gene}\t{alleles[0]}\t{alleles[1]}', file = outfile)
            elif len(alleles) == 1:
                print(f'{wildcards.dataset}\t{wildcards.sample}\txHLA\t{params.opts}\t{gene}\t{alleles[0]}\t{alleles[0]}', file = outfile)
            else:
                raise RuntimeError(f'xHLA: More than two alleles reported for gene {gene}.')

shell:
    """
    # Run trimgalore
    trim_galore --output_dir trim --fastqc -j {threads} --paired {input[0]:q} {input[1]:q} &>{log:q}

    # trimgalore compresses output only if input was compressed
    [ "{params.input_ext}" = ".gz" ] || gzip trim/*.fq &>>{log:q}

    # Move output into place
    mv -f {params.r1_tmp:q} {output.r1:q}
    mv -f {params.r2_tmp:q} {output.r2:q}
    mv -f {params.r1fastqc_tmp:q} {output.r1fastqc:q}
    mv -f {params.r2fastqc_tmp:q} {output.r2fastqc:q}
    """

shell:
    """
    mkdir -p {output.bindir:q}
    mkdir -p {output.tdir:q}

    # This is a hack to circumvent that xHLA does not allow for any control
    # over the number of threads it consumes - it will always consume
    # everything the machine reports. Two components seem to use
    # multithreading, the invocation of diamond and the lpsolver in R. Here
    # we fix the diamond issue by integrating a wrapper script called
    # 'diamond' in the PATH that sets the max threads paramter of diamond
    # accordingly. Currently there is no fix for the R code, which makes
    # this rule vulnerable to cluster engines that enforce the SMP limit.

    echo '#!/bin/bash' > {output.bindir:q}/diamond
    echo 'shift' >> {output.bindir:q}/diamond
    echo '/usr/bin/diamond blastx -p {threads} "$@"' >> {output.bindir:q}/diamond
    chmod +x {output.bindir:q}/diamond
    export PATH="{output.bindir}:$PATH"

    export LC_ALL=C
    ecode=0
    if python /opt/bin/run.py --sample_id sample.$HOSTNAME.$$ --input_bam_path {input.bam:q} --output_path {output.tdir:q} &>{log:q}
    then
        mv -v {output.tdir:q}/report-sample.$HOSTNAME.$$-hla.json {output.json:q} &>>{log:q}
    else
        ecode=1
    fi
    rm -rfv hla-sample.$HOSTNAME.$$ &>>{log:q}
    exit $ecode
    """

import pandas as pd

# Read the ground truth data
gtruth = pd.read_csv(snakemake.input.gtruth, '\t')
gtruth['Dataset'] = snakemake.wildcards.dataset
gtruth['Method'] = 'ground_truth'
gtruth['Options'] = ''

# Read the typing results
typing = pd.read_csv(snakemake.input.typing, '\t')

# Concatenate and pivot the data
compare = pd.concat([gtruth, typing])

# Create result column with set of called alleles per gene and drop original, individual columns
compare['Result'] = compare.apply(axis='columns', func=lambda x : {x.Allele1, x.Allele2})
compare.drop(columns=['Allele1','Allele2'], inplace = True)

# Pivot the table to have one column per method, drop constant column index level
compare = compare.pivot_table(index=['Dataset','Sample','Gene'], columns=['Method','Options'], aggfunc=lambda x : set.union(*x))
compare = compare.droplevel(0, axis='columns')

# Make the ground truth the first column
compare = compare[ [c for c in compare.columns if c[0]=='ground_truth'] + [c for c in compare.columns if c[0]!='ground_truth'] ]

def annotate(row, method):
    if pd.isna(row[method]) or len(row[method])==0:
        return 'no result'
    gt = row[('ground_truth', '')]
    if row[method] == gt:
        return 'correct'
    if pd.isna(gt) or len(gt)==0:
        return f'no ground truth: {row[method]}'
    if row[method] & gt: # Intersection
        return f'partial match: {row[method]}'
    else:
        return f'mismatch: {row[method]}'

# Create copy with annotated version of results
annotated = compare.copy()
for col in [ c for c in annotated.columns if c[0] != 'ground_truth' ]:
    annotated[col] = annotated.apply(lambda x : annotate(x, col), axis = 'columns')

# Write results
compare.to_csv(snakemake.output.raw, '\t', index = True)
annotated.to_csv(snakemake.output.annotated, '\t', index = True)

from Bio import SeqIO
import gzip

whitelist = [ 'chr' + e for e in [ str(num) for num in range(1,23) ] + ['X','Y'] ]

with open(snakemake.output[0], 'w') as outfile:
    with gzip.open(snakemake.input[0], 'rt') as infile:
        generator = ( record for record in SeqIO.parse(infile, "fasta") if record.id in whitelist )
        SeqIO.write(generator, outfile, 'fasta')

__author__ = "Christopher Schröder, Patrik Smeds"
__copyright__ = "Copyright 2020, Christopher Schröder, Patrik Smeds"
__email__ = "[email protected], [email protected]"
__license__ = "MIT"

from os import path

from snakemake.shell import shell

log = snakemake.log_fmt_shell(stdout=False, stderr=True)

# Check inputs/arguments.
if len(snakemake.input) == 0:
    raise ValueError("A reference genome has to be provided.")
elif len(snakemake.input) > 1:
    raise ValueError("Please provide exactly one reference genome as input.")

# Prefix that should be used for the database
prefix = snakemake.params.get("prefix", "")

if len(prefix) > 0:
    prefix = "-p " + prefix

shell("bwa-mem2 index" " {prefix}" " {snakemake.input[0]}" " {log}")

__author__ = "Johannes Köster"
__copyright__ = "Copyright 2016, Johannes Köster"
__email__ = "[email protected]"
__license__ = "MIT"


import os
from snakemake.shell import shell

extra = snakemake.params.get("extra", "")
log = snakemake.log_fmt_shell(stdout=True, stderr=True)

fq1 = snakemake.input.get("fq1")
assert fq1 is not None, "input-> fq1 is a required input parameter"
fq1 = (
    [snakemake.input.fq1]
    if isinstance(snakemake.input.fq1, str)
    else snakemake.input.fq1
)
fq2 = snakemake.input.get("fq2")
if fq2:
    fq2 = (
        [snakemake.input.fq2]
        if isinstance(snakemake.input.fq2, str)
        else snakemake.input.fq2
    )
    assert len(fq1) == len(
        fq2
    ), "input-> equal number of files required for fq1 and fq2"
input_str_fq1 = ",".join(fq1)
input_str_fq2 = ",".join(fq2) if fq2 is not None else ""
input_str = " ".join([input_str_fq1, input_str_fq2])

if fq1[0].endswith(".gz"):
    readcmd = "--readFilesCommand zcat"
else:
    readcmd = ""

outprefix = os.path.dirname(snakemake.output[0]) + "/"

shell(
    "STAR "
    "{extra} "
    "--runThreadN {snakemake.threads} "
    "--genomeDir {snakemake.params.index} "
    "--readFilesIn {input_str} "
    "{readcmd} "
    "--outFileNamePrefix {outprefix} "
    "--outStd Log "
    "{log}"
)