1__author__ = "Christopher Schröder, Johannes Köster, Julian de Ruiter"
2__copyright__ = (
3    "Copyright 2020, Christopher Schröder, Johannes Köster and Julian de Ruiter"
4)
5__email__ = "christopher.schroeder@tu-dortmund.de koester@jimmy.harvard.edu, julianderuiter@gmail.com"
6__license__ = "MIT"
7
8
9from os import path
10
11from snakemake.shell import shell
12
13
14# Extract arguments.
15extra = snakemake.params.get("extra", "")
16
17sort = snakemake.params.get("sort", "none")
18sort_order = snakemake.params.get("sort_order", "coordinate")
19sort_extra = snakemake.params.get("sort_extra", "")
20
21log = snakemake.log_fmt_shell(stdout=False, stderr=True)
22
23# Check inputs/arguments.
24if not isinstance(snakemake.input.reads, str) and len(snakemake.input.reads) not in {
25    1,
26    2,
27}:
28    raise ValueError("input must have 1 (single-end) or 2 (paired-end) elements")
29
30if sort_order not in {"coordinate", "queryname"}:
31    raise ValueError("Unexpected value for sort_order ({})".format(sort_order))
32
33# Determine which pipe command to use for converting to bam or sorting.
34if sort == "none":
35
36    # Simply convert to bam using samtools view.
37    pipe_cmd = "samtools view -Sbh -o {snakemake.output[0]} -"
38
39elif sort == "samtools":
40
41    # Sort alignments using samtools sort.
42    pipe_cmd = "samtools sort {sort_extra} -o {snakemake.output[0]} -"
43
44    # Add name flag if needed.
45    if sort_order == "queryname":
46        sort_extra += " -n"
47
48    prefix = path.splitext(snakemake.output[0])[0]
49    sort_extra += " -T " + prefix + ".tmp"
50
51elif sort == "picard":
52
53    # Sort alignments using picard SortSam.
54    pipe_cmd = (
55        "picard SortSam {sort_extra} INPUT=/dev/stdin"
56        " OUTPUT={snakemake.output[0]} SORT_ORDER={sort_order}"
57    )
58
59else:
60    raise ValueError("Unexpected value for params.sort ({})".format(sort))
61
62shell(
63    "(bwa-mem2 mem"
64    " -t {snakemake.threads}"
65    " {extra}"
66    " {snakemake.params.index}"
67    " {snakemake.input.reads}"
68    " | " + pipe_cmd + ") {log}"
69)
70