__author__ = "Julian de Ruiter"
__copyright__ = "Copyright 2017, Julian de Ruiter"
__email__ = "[email protected]"
__license__ = "MIT"


from snakemake.shell import shell


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

shell(
    "cutadapt"
    " {snakemake.params}"
    " -j {snakemake.threads}"
    " -o {snakemake.output.fastq}"
    " {snakemake.input[0]}"
    " > {snakemake.output.qc} {log}"
)

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

import subprocess as sp
import sys
from itertools import product
from snakemake.shell import shell

species = snakemake.params.species.lower()
release = int(snakemake.params.release)
build = snakemake.params.build

branch = ""
if release >= 81 and build == "GRCh37":
    # use the special grch37 branch for new releases
    branch = "grch37/"

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

spec = ("{build}" if int(release) > 75 else "{build}.{release}").format(
    build=build, release=release
)

suffixes = ""
datatype = snakemake.params.get("datatype", "")
if datatype == "dna":
    suffixes = ["dna.primary_assembly.fa.gz", "dna.toplevel.fa.gz"]
elif datatype == "cdna":
    suffixes = ["cdna.all.fa.gz"]
elif datatype == "cds":
    suffixes = ["cds.all.fa.gz"]
elif datatype == "ncrna":
    suffixes = ["ncrna.fa.gz"]
elif datatype == "pep":
    suffixes = ["pep.all.fa.gz"]
else:
    raise ValueError("invalid datatype, must be one of dna, cdna, cds, ncrna, pep")

success = False
for suffix in suffixes:
    url = "ftp://ftp.ensembl.org/pub/{branch}release-{release}/fasta/{species}/{datatype}/{species_cap}.{spec}.{suffix}".format(
        release=release,
        species=species,
        datatype=datatype,
        spec=spec.format(build=build, release=release),
        suffix=suffix,
        species_cap=species.capitalize(),
        branch=branch,
    )

    try:
        shell("curl -sSf {url} > /dev/null 2> /dev/null")
    except sp.CalledProcessError:
        continue

    shell("(curl -L {url} | gzip -d > {snakemake.output[0]}) {log}")
    success = True
    break

if not success:
    print(
        "Unable to download requested sequence data from Ensembl. "
        "Did you check that this combination of species, build, and release is actually provided?",
        file=sys.stderr,
    )
    exit(1)

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

import tempfile
import subprocess
import sys
import os
from snakemake.shell import shell
from snakemake.exceptions import WorkflowError

species = snakemake.params.species.lower()
release = int(snakemake.params.release)
build = snakemake.params.build
type = snakemake.params.type

if release < 98:
    print("Ensembl releases <98 are unsupported.", file=open(snakemake.log[0], "w"))
    exit(1)

branch = ""
if release >= 81 and build == "GRCh37":
    # use the special grch37 branch for new releases
    branch = "grch37/"

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

if type == "all":
    if species == "homo_sapiens" and release >= 93:
        suffixes = [
            "-chr{}".format(chrom) for chrom in list(range(1, 23)) + ["X", "Y", "MT"]
        ]
    else:
        suffixes = [""]
elif type == "somatic":
    suffixes = ["_somatic"]
elif type == "structural_variations":
    suffixes = ["_structural_variations"]
else:
    raise ValueError(
        "Unsupported type {} (only all, somatic, structural_variations are allowed)".format(
            type
        )
    )

species_filename = species if release >= 91 else species.capitalize()

urls = [
    "ftp://ftp.ensembl.org/pub/{branch}release-{release}/variation/vcf/{species}/{species_filename}{suffix}.{ext}".format(
        release=release,
        species=species,
        suffix=suffix,
        species_filename=species_filename,
        branch=branch,
        ext=ext,
    )
    for suffix in suffixes
    for ext in ["vcf.gz", "vcf.gz.csi"]
]
names = [os.path.basename(url) for url in urls if url.endswith(".gz")]

try:
    gather = "curl {urls}".format(urls=" ".join(map("-O {}".format, urls)))
    workdir = os.getcwd()
    with tempfile.TemporaryDirectory() as tmpdir:
        if snakemake.input.get("fai"):
            shell(
                "(cd {tmpdir}; {gather} && "
                "bcftools concat -Oz --naive {names} > concat.vcf.gz && "
                "bcftools reheader --fai {workdir}/{snakemake.input.fai} concat.vcf.gz "
                "> {workdir}/{snakemake.output}) {log}"
            )
        else:
            shell(
                "(cd {tmpdir}; {gather} && "
                "bcftools concat -Oz --naive {names} "
                "> {workdir}/{snakemake.output}) {log}"
            )
except subprocess.CalledProcessError as e:
    if snakemake.log:
        sys.stderr = open(snakemake.log[0], "a")
    print(
        "Unable to download variation data from Ensembl. "
        "Did you check that this combination of species, build, and release is actually provided? ",
        file=sys.stderr,
    )
    exit(1)

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


from snakemake.shell import shell


shell("samtools index {snakemake.params} {snakemake.input[0]} {snakemake.output[0]}")

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

from pathlib import Path
from snakemake.shell import shell

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

shell(
    "vep_install --AUTO cf "
    "--SPECIES {snakemake.params.species} "
    "--ASSEMBLY {snakemake.params.build} "
    "--CACHE_VERSION {snakemake.params.release} "
    "--CACHEDIR {snakemake.output} "
    "--CONVERT "
    "--NO_UPDATE {log}"
)

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

import sys
from pathlib import Path
from urllib.request import urlretrieve
from zipfile import ZipFile
from tempfile import NamedTemporaryFile

if snakemake.log:
    sys.stderr = open(snakemake.log[0], "w")

outdir = Path(snakemake.output[0])
outdir.mkdir()

with NamedTemporaryFile() as tmp:
    urlretrieve(
        "https://github.com/Ensembl/VEP_plugins/archive/release/{release}.zip".format(
            release=snakemake.params.release
        ),
        tmp.name,
    )

    with ZipFile(tmp.name) as f:
        for member in f.infolist():
            memberpath = Path(member.filename)
            if len(memberpath.parts) == 1:
                # skip root dir
                continue
            targetpath = outdir / memberpath.relative_to(memberpath.parts[0])
            if member.is_dir():
                targetpath.mkdir()
            else:
                with open(targetpath, "wb") as out:
                    out.write(f.read(member.filename))

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


from snakemake.shell import shell


shell("samtools view {snakemake.params} {snakemake.input[0]} > {snakemake.output[0]}")

__author__ = "Christopher Schröder, Johannes Köster, Julian de Ruiter"
__copyright__ = (
    "Copyright 2020, Christopher Schröder, Johannes Köster and Julian de Ruiter"
)
__email__ = "[email protected] [email protected], [email protected]"
__license__ = "MIT"


from os import path

from snakemake.shell import shell


# Extract arguments.
extra = snakemake.params.get("extra", "")

sort = snakemake.params.get("sort", "none")
sort_order = snakemake.params.get("sort_order", "coordinate")
sort_extra = snakemake.params.get("sort_extra", "")

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

# Check inputs/arguments.
if not isinstance(snakemake.input.reads, str) and len(snakemake.input.reads) not in {
    1,
    2,
}:
    raise ValueError("input must have 1 (single-end) or 2 (paired-end) elements")

if sort_order not in {"coordinate", "queryname"}:
    raise ValueError("Unexpected value for sort_order ({})".format(sort_order))

# Determine which pipe command to use for converting to bam or sorting.
if sort == "none":

    # Simply convert to bam using samtools view.
    pipe_cmd = "samtools view -Sbh -o {snakemake.output[0]} -"

elif sort == "samtools":

    # Sort alignments using samtools sort.
    pipe_cmd = "samtools sort {sort_extra} -o {snakemake.output[0]} -"

    # Add name flag if needed.
    if sort_order == "queryname":
        sort_extra += " -n"

    prefix = path.splitext(snakemake.output[0])[0]
    sort_extra += " -T " + prefix + ".tmp"

elif sort == "picard":

    # Sort alignments using picard SortSam.
    pipe_cmd = (
        "picard SortSam {sort_extra} INPUT=/dev/stdin"
        " OUTPUT={snakemake.output[0]} SORT_ORDER={sort_order}"
    )

else:
    raise ValueError("Unexpected value for params.sort ({})".format(sort))

shell(
    "(bwa-mem2 mem"
    " -t {snakemake.threads}"
    " {extra}"
    " {snakemake.params.index}"
    " {snakemake.input.reads}"
    " | " + pipe_cmd + ") {log}"
)

wrapper:
    "0.59.2/bio/vep/annotate"

shell:
    "(bcftools view --threads {threads} {input.bcf} {params.pipes} | bcftools view --threads {threads} -Ob > {output}) 2> {log}"

script:
    "../scripts/render-scenario.py"

shell:
    "varlociraptor "
    "call variants generic --obs {params.obs} "
    "--scenario {input.scenario} > {output} 2> {log}"

wrapper:
    "0.59.2/bio/bcftools/concat"

wrapper:
    "0.60.0/bio/freebayes"

wrapper:
    "0.60.0/bio/delly"

shell:
    "(bcftools view {input} | filter_vep --filter \"{params.filter}\" --vcf_info_field ANN --only_matched | bcftools view -Ob > {output}) 2> {log}"

shell:
    "varlociraptor filter-calls posterior-odds --events {params.events} --odds barely < {input} > {output} 2> {log}"

shell:
    "varlociraptor filter-calls control-fdr {input} --var {wildcards.vartype} "
    "--events {params.events} --fdr {params.threshold} > {output} 2> {log}"

wrapper:
    "0.59.2/bio/bcftools/concat"

wrapper:
    "0.64.0/bio/bwa-mem2/mem"

wrapper:
    "0.59.2/bio/picard/markduplicates"

wrapper:
    "0.59.2/bio/gatk/baserecalibrator"

shell:
    "rbt oncoprint --vep-annotation {params.groups} > {output} 2> {log}"

shell:
    "yara_mapper -t {threads} -ll {params.library_len} -ld {params.library_error} -o {output} {params.ref_prefix} {input.reads} > {log}"

wrapper:
    "0.61.0/bio/samtools/view"

shell:
    "samtools sort -n {input} | bamToBed -i - -bedpe > {output} 2> {log}"

shell:
    "samtools sort -n {input} | bamToBed -i - > {output} 2> {log}"

script:
    "../scripts/build_primer_regions.py"

script:
    "../scripts/build_target_regions.py"

shell:
    "sort -k1,1 -k2,2n {input} | mergeBed -i - > {output} 2> {log}"

shell:
    "(complementBed -i {input.target_regions} -g <(head "
    "-n {params.chroms} {input.genome_index} | cut "
    "-f 1,2 | sort -k1,1 -k 2,2n) > {output}) 2> {log}"

shell:
    "samtools view -h -b -L {input.regions} {input.bam} > {output} 2> {log}"

wrapper:
    "0.59.2/bio/reference/ensembl-sequence"

wrapper:
    "0.59.2/bio/samtools/faidx"

shell:
    "samtools dict {input} > {output} 2> {log} "

wrapper:
    "0.59.2/bio/reference/ensembl-variation"

shell:
    "rbt vcf-fix-iupac-alleles < {input} | bcftools view -Oz > {output}"

wrapper:
    "0.64.0/bio/bwa-mem2/index"

wrapper:
    "0.59.2/bio/vep/cache"

wrapper:
    "0.59.2/bio/vep/plugins"

shell:
    "bcftools view {input.bcf} > {output}.vcf;"
    "create_report {output}.vcf {input.ref} --flanking 100 "
    "--info-columns ANN dgiDB_drugs cosmic_LEGACY_ID --info-columns-prefixes PROB_ --sample-columns DP AF OBS"
    " --template {params} --tracks {input.bams} --output {output} --standalone 2>&1 > {log}; "
    "rm {output}.vcf"

shell:
    "varlociraptor estimate tmb "
    "--plot-mode {wildcards.mode} "
    "--coding-genome-size {params.coding_genome_size} "
    "--somatic-tumor-events {params.somatic_events} "
    "--tumor-sample {params.tumor_sample} "
    "< {input} > {output} 2> {log}"

wrapper:
    "0.56.0/bio/sra-tools/fasterq-dump"

shell:
    "cat {input} > {output} 2> {log}"

wrapper:
    "0.59.2/bio/cutadapt/pe"

wrapper:
    "0.59.2/bio/cutadapt/se"

shell:
    "cat {input} > {output} 2> {log}"

shell:
    "bcftools index {input} 2> {log}"

wrapper:
    "0.59.2/bio/samtools/index"

wrapper:
    "0.59.2/bio/tabix"

shell:
    "vl2svg {input} {output} 2> {log}"

import pandas as pd

chunksize = 10 ** 6
with open(snakemake.output[0], "w") as out:
    with open(snakemake.log[0], "w") as log_file:
        for data_primers in pd.read_csv(
            snakemake.input[0],
            sep="\t",
            header=None,
            chunksize=chunksize,
            usecols=[0, 1, 3, 5, 6],
        ):
            valid_primers = data_primers[0] == data_primers[3]
            print(
                data_primers[valid_primers]
                .drop(columns=[3, 6])
                .to_csv(sep="\t", index=False, header=False),
                file=out,
            )
            print(
                data_primers[~valid_primers].to_csv(
                    sep="\t", index=False, header=False
                ),
                file=log_file,
            )