"""
Author: C. Laczny
Affiliation: ESB group LCSB UniLU
Aim:
Date: [25 June 2019]
Run: snakemake -s Snakefile
Modified on: 2020-04-23
Modified for: New_Basecalling, Mapping to Hybrid assembly, metaT mapping, Binning, CheckM and GTDBTK
TODO:
s. TODO.md
"""

#shell.executable("/bin/bash")
#shell.prefix("source ~/.bashrc; ")
import os
from tempfile import TemporaryDirectory

configfile: "CONFIG.yaml"
DATA_DIR = config["data_dir"]
RESULTS_DIR = config["results_dir"]
DB_DIR=config["db_dir"]
RUNS=[config["runs"]["first"],
    config["runs"]["second"],
    config["runs"]["third"]]
##RUNS=config["runs"]["third"]
BARCODES=config["barcodes"]
#SAMPLES=["NEB2_MG_S17"]
#SAMPLES_ALL=config["samples"]
REFERENCES=["igc", "hg38"]
IGC_URI=config["igc"]["uri"]
HG38_URI=config["hg38"]["uri"]
ASSEMBLERS=config["assemblers"]
#MAPPERS=["bwa_mem", "minimap2"]


#configfile:"binning_config.yaml"
DATA_DIR=config["data_dir"]
RESULTS_DIR=config["results_dir"]
SAMPLES=config["samples"]
# DATABASE=config["database"]
BINNING_SAMPLES=config["binning_samples"]
HYBRID_ASSEMBLER=config["hybrid_assembler"]

############
rule all:
    #input: expand("{data_dir}/multifast5/{run}", data_dir=DATA_DIR, run=RUNS)
    #input: expand(os.path.join(RESULTS_DIR, "basecalled/guppy/cpu/{run}/basecalling.done"), run=RUNS)
    #input: expand(os.path.join(RESULTS_DIR, "basecalled/{run}/basecalling.done"), run=RUNS)
    #input: expand(os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/NO_MOD_basecalling.done"), run=RUNS)
    #input: expand(os.path.join(RESULTS_DIR, "basecalled/{run}/demux.done"), run=RUNS)
    #input: expand(os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/demux_NO_MOD.done"), run=RUNS)
    #input: expand(os.path.join(RESULTS_DIR, "basecalled/{run}/{barcode}/{barcode}.fastq.gz"), run=RUNS, barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/{barcode}/{barcode}.fastq.gz"), run=RUNS, barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"), barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "basecalled_NO_MOD/merged/{barcode}.fastq.gz"), barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "qc/lr/merged/{barcode}/{barcode}NanoStats.txt"), barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "qc/lr/merged/{barcode}/{barcode}NanoStats_NO_MOD.txt"), barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "qc/lr/{run}/{barcode}/{barcode}NanoStats.txt"), run=RUNS, barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "qc/lr/{run}/{barcode}/{barcode}NanoStats_NO_MOD.txt"), run=RUNS, barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "mapping/guppy/gpu/{barcode}/{barcode}-long_reads_on_short_reads.sam"), barcode=BARCODES)		# checked with CCL. not required (legacy)
    #input: expand(os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}-x-igc.bam"), barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}-x-hg38.bam"), barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}_R1_001.fastp.fq.gz"), sr_sample="NEB2_MG_S17")
    # not required (legacy code) input: expand(os.path.join(RESULTS_DIR, "mapping/sr/{sample}-x-igc.bam"), sample="NEB2_MG_S17")
    # not required (legacy code) input: expand(os.path.join(RESULTS_DIR, "mapping/sr/{sample}-x-hg38.bam"), sample="NEB2_MG_S17")
    #input: expand(os.path.join(RESULTS_DIR, "mapping/sr/{mapper}/{sample}-x-{reference}.bam"), sample="NEB2_MG_S17", reference=REFERENCES, mapper=MAPPERS)
    #input: expand(os.path.join(RESULTS_DIR, "mapping/sr/{mapper}/{sr_sample}_reads-x-{barcode}-{assembler}_contigs.bam"), sr_sample="NEB2_MG_S17", barcode=BARCODES, assembler=ASSEMBLERS, mapper=MAPPERS)
    #input: expand(os.path.join(RESULTS_DIR, "genomecov/lr/merged/{barcode}/{barcode}-x-igc.cov.txt"), barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "genomecov/sr/{mapper}/{sample}-x-igc.avg_cov.txt"), sample="NEB2_MG_S17", mapper=MAPPERS)
    #input: expand(os.path.join(RESULTS_DIR, "genomecov/lr/merged/{barcode}/{barcode}-x-igc.avg_cov.txt"), barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "mapping/lr/merged/{sample}-x-{barcode}.bam"), sample=SAMPLES, barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/{barcode}/assembly.fna"), barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam"), barcode=BARCODES, assembler=ASSEMBLERS)
    #input: expand(os.path.join(RESULTS_DIR, "genomecov/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.avg_cov.txt"), barcode=BARCODES, assembler=ASSEMBLERS)
    #input: expand(os.path.join(RESULTS_DIR, "mapping/sr/{mapper}/{sample}_reads-x-{barcode}-{assembler}_contigs.bam"), sample=SAMPLES, barcode=BARCODES, assembler=ASSEMBLERS, mapper=MAPPERS)
    #input: expand(os.path.join(RESULTS_DIR, "genomecov/sr/{mapper}/{sample}_reads-x-{barcode}-{assembler}_contigs.avg_cov.txt"), sample=SAMPLES, barcode=BARCODES, assembler=ASSEMBLERS, mapper=MAPPERS)
    #input: "assemble_and_coverage.done"
    #input: expand(os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/medaka/round_{round}/lr/merged/{barcode}/assembly_polished.fna"), assembler=ASSEMBLERS, barcode=BARCODES, round="01")
    #input: "basecall_merge_qc.done"
    #input: "basecall_merge_qc_NO_MOD.done"
    #input: "coverage_of_references.done"
    #input: expand(os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index.readdb"), barcode=BARCODES)
    #input: "annotate.done"
    #input: expand(os.path.join(RESULTS_DIR, "qc/lr/{run}/{barcode}/{barcode}NanoStats.txt"), run=RUNS, barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "qc/sr/fastqc/{sr_sample}/{sr_sample}_{suffix}.fastqc.zip"), sr_sample="NEB2_MG_S17",suffix=["R1_001.fastp", "R2_001.fastp"])
    #input: expand(os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/racon/round_{round}/lr/merged/{barcode}/assembly_polished.fna"), barcode=BARCODES, assembler=ASSEMBLERS, round="01")
    #input: expand(os.path.join(RESULTS_DIR, "assembly/metaspades_hybrid/lr_{barcode}-sr_{sr_sample}/contigs.fna"), barcode=BARCODES, sr_sample="NEB2_MG_S17")
    # specifc HYBRID_ASSEBMLERS instead of ASSEMBLERS; not sure if this will solve the issue, though. #input: expand(os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.bam"), barcode=BARCODES, assembler=ASSEMBLERS, sr_sample="NEB2_MG_S17")
    # RUN AFTER line #80 #input: expand(os.path.join(RESULTS_DIR, "mapping/lr/bwa_mem/{barcode}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.bam"), barcode=BARCODES, assembler=ASSEMBLERS, sr_sample="NEB2_MG_S17")
    #input: expand(os.path.join(RESULTS_DIR, "mapping/polished/{polisher}/round_{round}/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_polished_contigs.bam"),  barcode=BARCODES, assembler=ASSEMBLERS, round="01", polisher="medaka")
    #input: expand(os.path.join(RESULTS_DIR, "assembly/megahit/{sr_sample}/final.contigs.fna"), sr_sample="NEB2_MG_S17")
    ### RUN LATER #input: expand(os.path.join(RESULTS_DIR, "mapping/lr/bwa_mem/{barcode}_reads-x-{sr_sample}-{assembler}_contigs.bam"), barcode=BARCODES, assembler=ASSEMBLERS, sr_sample="NEB2_MG_S17")
    #input: expand(os.path.join(RESULTS_DIR, "annotation/nonpareil/{sr_sample}_R{orientation}_{suffix}.fastp.npo"), sr_sample="NEB2_MG_S17", orientation=["1", "2"], suffix="001")
    #input: expand(os.path.join(RESULTS_DIR, "assembly/flye/polished/nanopolish/round_{round}/lr/merged/{barcode}/windows"), round="01", barcode=BARCODES)
    input: expand(os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/nanopolish/round_{round}/lr/merged/{barcode}/assembly_polished.fna"), assembler=ASSEMBLERS, round="01", barcode=BARCODES)
    ### RUN LATER #input: expand(os.path.join(RESULTS_DIR, "annotation/methylation/{assembler}/nanopolish/round_{round}/lr/merged/{barcode}/methylation_frequencies.tsv"), assembler=ASSEMBLERS, round="01", barcode=BARCODES)
    #input: expand(os.path.join(RESULTS_DIR, "annotation/diamond/lr/merged/{barcode}.diamond.daa"), barcode=BARCODES)
    #input: expand(f"{RESULTS_DIR}/annotation/mmseq2/{{assembler1}}__{{assembler2}}_rbh"), zip, assembler1=["flye", "flye", "megahit"], assembler2=["megahit", "megahit_metaspades_hybrid", "megahit_metaspades_hybrid"])


######
# WRAPPER RULES
######
# Added the `shell:` so as to avoid it from being autodetected as `localrule`. This is needed so that an email can be sent upon event changes for this rule.
rule BASECALL_MERGE_QC:
    input:
        expand(os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"), barcode=BARCODES),
        expand(os.path.join(RESULTS_DIR, "qc/lr/merged/{barcode}/{barcode}NanoStats.txt"), barcode=BARCODES),
        expand(os.path.join(RESULTS_DIR, "qc/lr/{run}/{barcode}/{barcode}NanoStats.txt"), run=RUNS, barcode=BARCODES),
        expand(os.path.join(RESULTS_DIR, "preprocessing/sr/{sample}.fastp.{report_type}"), sample=SAMPLES, report_type=["html", "json"]),
        expand(os.path.join(RESULTS_DIR, "qc/sr/fastqc/{sample}/{sample}_R{orientation}_001.fastp.fastqc.html"), sample=SAMPLES, orientation=["1", "2"])
    output: "basecall_merge_qc.done"
    shell:
        """
        touch {output}
        """

rule BASECALL_MERGE_QC_NO_MOD:
    input:
        expand(os.path.join(RESULTS_DIR, "basecalled_NO_MOD/merged/{barcode}.fastq.gz"), barcode=BARCODES),
        expand(os.path.join(RESULTS_DIR, "qc/lr/merged/{barcode}/{barcode}NanoStats_NO_MOD.txt"), barcode=BARCODES),
        expand(os.path.join(RESULTS_DIR, "qc/lr/{run}/{barcode}/{barcode}NanoStats_NO_MOD.txt"), run=RUNS, barcode=BARCODES)
    output: "basecall_merge_qc_NO_MOD.done"
    shell:
        """
        touch {output}
        """

rule COVERAGE_OF_REFERENCES:
    input: 
        expand(os.path.join(RESULTS_DIR, "genomecov/sr/minimap2/{sample}-x-{reference}.avg_cov.txt"), sample=SAMPLES, reference=REFERENCES),
        expand(os.path.join(RESULTS_DIR, "genomecov/sr/bwa_mem/{sample}-x-{reference}.avg_cov.txt"), sample=SAMPLES, reference=REFERENCES),
        expand(os.path.join(RESULTS_DIR, "genomecov/lr/merged/{barcode}/{barcode}-x-{reference}.avg_cov.txt"), barcode=BARCODES, reference=REFERENCES)
    output: "coverage_of_references.done"
    shell:
        """
        touch {output}
        """

rule TEST:
        input:
            # Commented out as it crashes due to a Key error
            #expand(os.path.join(RESULTS_DIR, "assembly/rebaler/lr/merged/{barcode}/lr_{barcode}-sr_contigs_{sample}.fna"), barcode=BARCODES, sample=SAMPLES)
            expand(os.path.join(RESULTS_DIR, "annotation/diamond/lr/merged/{barcode}.diamond.daa"), barcode=BARCODES),
            expand(os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sample}/contigs.fna"), barcode=BARCODES, sample=SAMPLES, assembler="metaspades_hybrid")

rule TEST_PRODIGAL:
        input:
            expand(os.path.join(RESULTS_DIR, "annotation/proteins/{assembler}/lr/merged/{barcode}/assembly.faa"), assembler="flye", barcode=BARCODES),
            expand(os.path.join(RESULTS_DIR, "annotation/proteins/{assembler}/polished/{polisher}/round_{round}/lr/merged/{barcode}/assembly_polished.faa"), assembler="flye", polisher=["medaka", "racon"], round="01", barcode=BARCODES),
            expand(os.path.join(RESULTS_DIR, "annotation/proteins/{assembler}/{sample}/final.contigs.faa"), assembler="megahit", sample=SAMPLES),
            expand(os.path.join(RESULTS_DIR, "annotation/proteins/{assembler}/lr_{barcode}-sr_{sample}/contigs.faa"), assembler="metaspades_hybrid", barcode=BARCODES, sample=SAMPLES)
        output: "prodigal_gene_call.done"
        shell:
            """
            touch {output}
            """

rule TEST_DIAMOND:
        input:
            expand(os.path.join(RESULTS_DIR, "annotation/diamond/{assembler}/lr/merged/{barcode}/assembly.tsv"), assembler="flye", barcode=BARCODES),
            expand(os.path.join(RESULTS_DIR, "annotation/diamond/{assembler}/polished/{polisher}/round_{round}/lr/merged/{barcode}/assembly_polished.tsv"), assembler="flye", polisher=["medaka", "racon"], round="01", barcode=BARCODES),
            expand(os.path.join(RESULTS_DIR, "annotation/diamond/{assembler}/{sample}/final.contigs.tsv"), assembler="megahit", sample=SAMPLES),
            expand(os.path.join(RESULTS_DIR, "annotation/diamond/{assembler}/lr_{barcode}-sr_{sample}/contigs.tsv"), assembler="metaspades_hybrid", barcode=BARCODES, sample=SAMPLES)
        output: "diamond_proteins.done"
        shell:
            """
            touch {output}
            """


rule ASSEMBLE_AND_COVERAGE:
    input:
        # long reads on long read contigs
        expand(os.path.join(RESULTS_DIR, "genomecov/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.avg_cov.txt"), barcode=BARCODES, assembler=ASSEMBLERS),
        # short reads on long read contigs
        expand(os.path.join(RESULTS_DIR, "genomecov/sr/minimap2/{sample}_reads-x-{barcode}-{assembler}_contigs.avg_cov.txt"), sample=SAMPLES, barcode=BARCODES, assembler=ASSEMBLERS),
        # short reads on long read contigs
        expand(os.path.join(RESULTS_DIR, "genomecov/sr/bwa_mem/{sample}_reads-x-{barcode}-{assembler}_contigs.avg_cov.txt"), sample=SAMPLES, barcode=BARCODES, assembler=ASSEMBLERS),
        #expand(os.path.join(RESULTS_DIR, "assembly/operams/lr_{barcode}-sr_{sample}/contigs.fna"), sample=SAMPLES, barcode=BARCODES),
        #expand(os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam"), barcode=BARCODES, assembler="flye"),
        # short reads on hybrid contigs
        #expand(os.path.join(RESULTS_DIR, "genomecov/sr/bwa_mem/{sample}_reads-x-lr_{barcode}_sr_{sample}-{assembler}_contigs.avg_cov.txt"), sample=SAMPLES, barcode=BARCODES, assembler="operams"),
        # long reads on hybrid contigs
        #expand(os.path.join(RESULTS_DIR, "genomecov/lr/bwa_mem/{barcode}_reads-x-lr_{barcode}_sr_{sample}-{assembler}_contigs.avg_cov.txt"), barcode=BARCODES, sample=SAMPLES, assembler="operams"),
        #expand(os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/{barcode}/assembly.fna"), barcode=BARCODES)
        # short reads on short read contigs
        expand(os.path.join(RESULTS_DIR, "genomecov/sr/bwa_mem/{sample}_reads-x-{sample}-{assembler}_contigs.avg_cov.txt"), sample=SAMPLES, assembler="megahit"),
        # long reads on short read contigs
        expand(os.path.join(RESULTS_DIR, "genomecov/lr/bwa_mem/{barcode}_reads-x-{sample}-{assembler}_contigs.avg_cov.txt"), barcode=BARCODES, sample=SAMPLES, assembler="megahit")
    output: "assemble_and_coverage.done"
    shell:
        """
        touch {output}
        """

rule ANNOTATE:
    input:
        # short reads coverage and diversity estimate
        expand(os.path.join(RESULTS_DIR, "annotation/nonpareil/{sample}_R{orientation}_001.fastp.npa"), sample=SAMPLES, orientation=["1", "2"]),
        # compute methylation frequencies
        expand(os.path.join(RESULTS_DIR, "annotation/methylation/{assembler}/nanopolish/round_{round}/lr/merged/{barcode}/methylation_frequencies.tsv"), assembler="flye", barcode=BARCODES, round="01")
    output: "annotate.done"
    shell:
        """
        touch {output}
        """
        
rule POLISH_AND_COVERAGE:
    input:
        # short read-polished contigs from OPERA-MS
        #expand(os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sample}/contigs.fna"), sample=SAMPLES, barcode=BARCODES, assembler="operams"),
        #long read polished contigs
        #expand(os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/{polisher}/round_{round}/lr/merged/{barcode}/assembly_polished.fna"), assembler="flye", polisher="medaka", round="01", barcode=BARCODES),
        # map long reads to long read polished contigs
        expand(os.path.join(RESULTS_DIR, "genomecov/polished/{polisher}/round_{round}/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_polished_contigs.avg_cov.txt"), assembler="flye", polisher="medaka", round="01", barcode=BARCODES),
        # create nanonoplish index
        #expand(os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index.readdb"), barcode=BARCODES)
        # polish flye contigs with nanopolish
        #expand(os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/nanopolish/round_{round}/lr/merged/{barcode}/assembly_polished.fna"), assembler="flye", round="01", barcode=BARCODES),
        # polish using racon
        expand(os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/racon/round_{round}/lr/merged/{barcode}/assembly_polished.fna"), assembler="flye", round="01", barcode=BARCODES)


rule DB_MMSEQ2:
    input:
        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/{assembler}_db"), assembler=["flye", "megahit", "metaspades_hybrid"])

rule COMPARE_MMSEQ2:
    input:
        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_megahit_rbh")),
        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_metaspades_hybrid_rbh")),
        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/megahit_metaspades_hybrid_rbh"))

rule CONVERT_MMSEQ2:
    input:
        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_megahit.m8")),
        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_metaspades_hybrid.m8")),
        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/megahit_metaspades_hybrid.m8"))

rule PREPARE_FILES:
    input:
        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/plot_files_ready.done"))

rule PLOT_MMSEQ2:
    input:
#        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/plot_files_ready.done")),
        expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/upset_plots.done")) 

rule METAT:
    input:
        expand(os.path.join(RESULTS_DIR, "preprocessing/metaT/{metaT_sample}.fastp.{report_type}"), metaT_sample="GDB_2018_metaT", report_type=["html", "json"]),
#	expand(os.path.join(RESULTS_DIR, "mapping/metaT/lr/{metaT_sample}-x-{barcode}.bam"), metaT_sample="GDB_2018_metaT", barcode=BARCODES),
        expand(os.path.join(RESULTS_DIR, "mapping/metaT/sr/{metaT_sample}_reads-x-{sr_sample}-{assembler}_contigs.bam"), metaT_sample="GDB_2018_metaT", sr_sample="NEB2_MG_S17", assembler="megahit"),
        expand(os.path.join(RESULTS_DIR, "mapping/metaT/sr/{metaT_sample}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.bam"), metaT_sample="GDB_2018_metaT", sr_sample="NEB2_MG_S17", barcode=BARCODES, assembler="metaspades_hybrid"),
        expand(os.path.join(RESULTS_DIR, "mapping/metaT/lr/{metaT_sample}_reads-x-{barcode}-{assembler}_contigs.bam"), metaT_sample="GDB_2018_metaT", barcode=BARCODES, assembler=ASSEMBLERS),
        expand(os.path.join(RESULTS_DIR, "genomecov/metaT/sr/{metaT_sample}_reads-x-{sr_sample}-{assembler}_contigs.avg_cov.txt"), metaT_sample="GDB_2018_metaT", sr_sample="NEB2_MG_S17", assembler="megahit"),
        expand(os.path.join(RESULTS_DIR, "genomecov/metaT/sr/{metaT_sample}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.avg_cov.txt"), metaT_sample="GDB_2018_metaT", sr_sample="NEB2_MG_S17", barcode=BARCODES, assembler="metaspades_hybrid"),
        expand(os.path.join(RESULTS_DIR, "genomecov/metaT/lr/{metaT_sample}_reads-x-{barcode}-{assembler}_contigs.avg_cov.txt"), metaT_sample="GDB_2018_metaT", barcode=BARCODES, assembler=ASSEMBLERS)


rule MAPPING:
    input:
        expand(os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.bwt"), hybrid_assembler=HYBRID_ASSEMBLER),
        expand(os.path.join(RESULTS_DIR, "mapping/{mapper}_{reads}_{hybrid_assembler}/{mapper}_{reads}_{hybrid_assembler}.bam"), mapper=["bwa", "mmi"], reads=["sr", "lr"], hybrid_assembler=HYBRID_ASSEMBLER),
        expand(os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{hybrid_assembler}.mmi"), hybrid_assembler=HYBRID_ASSEMBLER),
        expand(os.path.join(RESULTS_DIR, "mapping/{mapper}_merged_{hybrid_assembler}/{mapper}_merged_{hybrid_assembler}.bam"), mapper=["bwa", "mmi"], hybrid_assembler=HYBRID_ASSEMBLER)


rule BINNING:
    input:
#        expand(os.path.join(DATA_DIR, "{sample}/dastool_output/{sample}"), sample=BINNING_SAMPLES),
#        expand(os.path.join(DATA_DIR, "checkm_output/{sample}/lineage.ms"), sample=BINNING_SAMPLES)
#        expand(os.path.join(DATA_DIR, "gtdbtk_output/{sample}/gtdbtk.bac120.summary.tsv"), sample=BINNING_SAMPLES)
        expand(os.path.join(RESULTS_DIR, "assembly/{sample}.fa"), sample=BINNING_SAMPLES),
        expand(os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_proteins.faa"), sample=BINNING_SAMPLES),
        expand(os.path.join(RESULTS_DIR, "Binning/checkm_output/{sample}_output.txt"), sample=BINNING_SAMPLES),
        expand(os.path.join(RESULTS_DIR, "Binning/gtdbtk_output/{sample}/gtdbtk.bac120.summary.tsv"), sample=BINNING_SAMPLES)



######
# RULES
######       
# rule unpack:
#     input: os.path.join(DATA_DIR, "2018-11-08-minion_latestdata.7z")
#     output: os.path.join(DATA_DIR, "raw/raw_unpacked.done")
#     threads: config["p7zip"]["threads"]
#     shell:
#         """
#         date
#         {config[p7zip][bin]} x -o$(dirname {output}) {input} -mmt{threads} &> {output}
#         date
#         """

# Take the per-read FAST5 files and create multi-FAST5 files including a batch of reads (typically not all reads, though).
rule create_multifast5s:
    input: os.path.join(DATA_DIR, "raw/{run}")
    output: protected(directory(os.path.join(DATA_DIR, "multifast5/{run}")))
    threads: config["ont_fast5_api"]["single_to_multi_fast5"]["threads"]
    log: os.path.join(DATA_DIR, "multifast5/{run}/create_multifast5s")
    shell:
        """
        (date &&\
        {config[ont_fast5_api][single_to_multi_fast5][bin]} --input_path {input} --save_path {output} --filename_base multifast5 --batch_size {config[ont_fast5_api][single_to_multi_fast5][batch]} --recursive -t {threads} &&\
        touch {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# FLO-MIN107 SQK-LSK108           dna_r9.5_450bps
#   --verbose_logs                    Enable verbose logs.
#   --qscore_filtering                Enable filtering of reads into PASS/FAIL
#                                     folders based on min qscore.
#   --min_qscore arg                  Minimum acceptable qscore for a read to be
#                                     filtered into the PASS folder
#   --disable_pings                   Disable the transmission of telemetry
#   --num_callers arg                 Number of parallel basecallers to create.
#   --num_barcode_threads arg         Number of worker threads to use for
#                                     barcoding.
#   --barcode_kits arg                Space separated list of barcoding kit(s) or
#                                     expansion kit(s) to detect against. Must be
#                                     in double quotes.
#   --trim_barcodes                   Trim the barcodes from the output sequences
#                                     in the FastQ files.

# Left in for legacy look-up reasons. GPU-based basecalling strongly encouraged.
# rule guppy_cpu_basecall:
#     input: os.path.join(DATA_DIR, "multifast5/{run}")
#     output: os.path.join(RESULTS_DIR, "basecalled/guppy/cpu/{run}/basecalling.done")
#     threads: config["guppy_cpu"]["cpu_threads"]
#     shell:
#         """
#         date
#         {config[guppy_cpu][bin]} --input_path {input} --save_path $(dirname {output}) --config {config[guppy_cpu][config]} --cpu_threads_per_caller {threads} --disable_pings
#         touch {output}
#         date
#         """

    #input: os.path.join(DATA_DIR, "multifast5/{run}")
# Basecalling the multi-FAST5 files.
# Currently requires dummy target, but would be nice to have this be a "real" file target. However, it seems that files are filled incrementally, hence there is no check if the whole process was successful or not if checking only for file presence.
rule guppy_gpu_basecall:
    input: rules.create_multifast5s.output
    output: os.path.join(RESULTS_DIR, "basecalled/{run}/basecalling.done")
    log: os.path.join(RESULTS_DIR, "basecalled/{run}/basecalling")
    threads: config["guppy_gpu"]["cpu_threads"]
    shell:
        """
        (date &&\
        {config[guppy_gpu][bin]} --input_path {input} --save_path $(dirname {output}) --config {config[guppy_gpu][config]} --cpu_threads_per_caller {threads} -x {config[guppy_gpu][gpu_device]} --disable_pings --compress_fastq --records_per_fastq {config[guppy_gpu][records_per_fastq]} --chunk_size {config[guppy_gpu][chunk_size]} --chunks_per_runner {config[guppy_gpu][chunks_per_runner]} --gpu_runners_per_device {config[guppy_gpu][runners_per_device]} --num_callers {config[guppy_gpu][num_callers]} && \
        touch {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule guppy_gpu_basecall_NO_MOD:
    input: rules.create_multifast5s.output
    output: os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/NO_MOD_basecalling.done")
    log: os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/NO_MOD_basecalling")
    threads: config["guppy_gpu"]["cpu_threads"]
    shell: 
        """
        (date &&\
        {config[guppy_gpu][bin]} --input_path {input} --save_path $(dirname {output}) --config {config[guppy_gpu][hac_config]} --cpu_threads_per_caller {threads} -x {config[guppy_gpu][gpu_device]} --disable_pings --compress_fastq --records_per_fastq {config[guppy_gpu][records_per_fastq]} --chunk_size {config[guppy_gpu][chunk_size]} --chunks_per_runner {config[guppy_gpu][chunks_per_runner]} --gpu_runners_per_device {config[guppy_gpu][runners_per_device]} --num_callers {config[guppy_gpu][num_callers]} && \
        touch {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Demultiplex.
# Currently requires dummy target, but would be nice to have this be a "real" file target. However, it seems that files are filled incrementally, hence there is no check if the whole process was successful or not if checking only for file presence.
rule guppy_gpu_demux:
    input: rules.guppy_gpu_basecall.output
    output: os.path.join(RESULTS_DIR, "basecalled/{run}/demux.done")
    log: os.path.join(RESULTS_DIR, "basecalled/{run}/demux")
    threads: config["guppy_barcoder"]["threads"]
    shell:
        """
        (date &&\
        {config[guppy_barcoder][bin]} --input_path $(dirname {input}) --save_path $(dirname {output}) -t {threads} --compress_fastq --records_per_fastq {config[guppy_gpu][records_per_fastq]} --trim_barcodes && \
        touch {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """
rule guppy_gpu_demux_NO_MOD:
    input: rules.guppy_gpu_basecall_NO_MOD.output
    output: os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/demux_NO_MOD.done")
    log: os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/demux")
    threads: config["guppy_barcoder"]["threads"]
    shell:
        """
        (date &&\
        {config[guppy_barcoder][bin]} --input_path $(dirname {input}) --save_path $(dirname {output}) -t {threads} --compress_fastq --records_per_fastq {config[guppy_gpu][records_per_fastq]} --trim_barcodes && \
        touch {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Merge multiple basecalled files, i.e., FASTQ files per barcode and run. Each multi-FAST5 file will produce a respective FASTQ file.
# TODO: Simplify output path to `basecalled/{run}/{barcode}.fastq.gz`. This will require to change `$(dirname {output})` to `$(dirname {output})/{wildcards.barcode}`
rule merge_individual_fastq_per_barcode:
    input: rules.guppy_gpu_demux.output
    output: os.path.join(RESULTS_DIR, "basecalled/{run}/{barcode}/{barcode}.fastq.gz")
    shell:
        """
        date
        cat $(find $(dirname {output}) -name "*.fastq.gz" | sort) > {output}
        touch {output}
        date
        """

rule merge_individual_fastq_per_barcode_NO_MOD:
    input: rules.guppy_gpu_demux_NO_MOD.output
    output: os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/{barcode}/{barcode}.fastq.gz")
    shell:
        """
        date
        cat $(find $(dirname {output}) -name "*.fastq.gz" | sort) > {output}
        touch {output}
        date
        """

# Helper function to get the input in order to create a *single* fastq file for all runs with the *same* barcode
# TODO: Adjust according to TODO-merge_individual_fastq_per_barcode
def get_all_fastq_for_barcode(wildcards):
    return expand(os.path.join(RESULTS_DIR, "basecalled/{run}/{barcode}/{barcode}.fastq.gz"), run=RUNS, barcode=wildcards.barcode)

def get_all_fastq_for_barcode_NO_MOD(wildcards):
    return expand(os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/{barcode}/{barcode}.fastq.gz"), run=RUNS, barcode=wildcards.barcode)

# Merge basecalled files (FASTQ) per barcode, i.e., combine reads from multiple runs according to their barcode.
rule get_single_fastq_per_barcode:
    input: get_all_fastq_for_barcode 
    output: os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz")
    shell:
        """
        date
        cat {input} > {output}
        date
        """
rule get_single_fastq_per_barcode_NO_MOD:
    input: get_all_fastq_for_barcode_NO_MOD
    output: os.path.join(RESULTS_DIR, "basecalled_NO_MOD/merged/{barcode}.fastq.gz")
    shell:
        """
        date
        cat {input} > {output}
        date
        """

# Get basic statistics of the ONT data; over all runs
# TODO: Simplify top of path: replace `lr` by `nanostat`
# TODO: Write more generic: have patterned `input: os.path.join(RESULTS_DIR, "basecalled/{prefix}/{barcode}.fastq.gz")` and patterned `output: os.path.join(RESULTS_DIR, "qc/nanostat/{prefix}/{barcode}/{barcode}NanoStats.txt")
# 
rule nanostat:
    input: rules.get_single_fastq_per_barcode.output
    output: os.path.join(RESULTS_DIR, "qc/lr/merged/{barcode}/{barcode}NanoStats.txt")
    log: os.path.join(RESULTS_DIR, "qc/lr/merged/{barcode}/{barcode}NanoStats")
    conda: "envs/nanostat.yaml"
    shell:
        """
        (date &&\
        NanoStat --fastq {input} --outdir $(dirname {output}) --prefix {wildcards.barcode} -n $(basename -s "" {output}) &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule nanostat_NO_MOD:
    input: rules.get_single_fastq_per_barcode_NO_MOD.output
    output: os.path.join(RESULTS_DIR, "qc/lr/merged/{barcode}/{barcode}NanoStats_NO_MOD.txt")
    log: os.path.join(RESULTS_DIR, "qc/lr/merged/{barcode}/{barcode}NanoStats_NO_MOD")
    conda: "envs/nanostat.yaml"
    shell:
        """
        (date &&\
        NanoStat --fastq {input} --outdir $(dirname {output}) --prefix {wildcards.barcode} -n $(basename -s "" {output}) &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# TODO: Deprecated if TODO-nanostat is completed
# Get basic statistics of the ONT data; per run
rule nanostat_per_run:
    input: os.path.join(RESULTS_DIR, "basecalled/{run}/{barcode}/{barcode}.fastq.gz")
    output: os.path.join(RESULTS_DIR, "qc/lr/{run}/{barcode}/{barcode}NanoStats.txt")
    conda: "envs/nanostat.yaml"
    shell:
        """
        date
        NanoStat --fastq {input} --outdir $(dirname {output}) --prefix {wildcards.barcode} -n $(basename -s "" {output})
        date
        """

rule nanostat_per_run_NO_MOD:
    input: os.path.join(RESULTS_DIR, "basecalled_NO_MOD/{run}/{barcode}/{barcode}.fastq.gz")
    output: os.path.join(RESULTS_DIR, "qc/lr/{run}/{barcode}/{barcode}NanoStats_NO_MOD.txt")
    conda: "envs/nanostat.yaml"
    shell:
        """
        date
        NanoStat --fastq {input} --outdir $(dirname {output}) --prefix {wildcards.barcode} -n $(basename -s "" {output})
        date
        """

# Download the IGC (http://www.nature.com/nbt/journal/vaop/ncurrent/full/nbt.2942.html)
# Downloading using wget-command instead of FTP-remote in snakemake as the former offers progress information
rule fetch_igc:
    output: protected(os.path.join(DB_DIR, "igc/igc.fna.gz"))
    log: os.path.join(DB_DIR, "igc/fetch_igc")
    shell:
        """
        (date &&\
        wget ftp://{IGC_URI} -O {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Download HG38 human reference genome
# Downloading using wget-command instead of FTP-remote in snakemake as the former offers progress information
rule fetch_hg38:
    output: protected(os.path.join(DB_DIR, "hg38/GCF_000001405.38_GRCh38.p12_genomic.fna.gz"))
    log: os.path.join(DB_DIR, "hg38/fetch_hg38")
    shell:
        """
        (date &&\
        wget ftp://{HG38_URI} -O {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule create_link_to_hg38:
    input: os.path.join(DB_DIR, "hg38/GCF_000001405.38_GRCh38.p12_genomic.fna")
    output: protected(os.path.join(DB_DIR, "hg38/hg38.fna"))
    shell:
        """
        date
        ln -s $(basename -s "" {input}) {output}
        date
        """

rule unpack_reference_data:
    input: os.path.join(DB_DIR, "{prefix}.fna.gz")
    output: protected(os.path.join(DB_DIR, "{prefix}.fna"))
    shell:
        """
        date
        gzip -dc {input} > {output}
        date
        """

# Build a minimap2 index to map *ONT* reads against
rule build_minimap2_ont_index:
    input: "{prefix}.fna"
    output: "{prefix}-ont.mmi"
    log: "{prefix}.minimap2_ont_index"
    conda: "envs/minimap2.yaml"
    shell:
        """
        (date &&\
        minimap2 -x map-ont -d {output} {input} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Build a minimap2 short read index in order to map short reads to a reference/long read-based contigs
# NOTE: Unclear whether this would be also the right option for mapping short reads to long reads (i.e., reads-to-reads), due to the increased error rate in long reads. For more discussions on this, see https://github.com/lh3/minimap2/issues/407#issuecomment-493823000
# TODO: Combine this rule and `build_minimap2_ont_index` using a conditional on the suffix (ont vs. sr). This makes the code more readable as it removes largely redundant `input` and `output` keywords
rule build_minimap2_sr_index:
    input: "{prefix}.fna"
    output: "{prefix}-sr.mmi"
    log: "{prefix}.minimap2_sr_index"
    conda: "envs/minimap2.yaml"
    shell:
        """
        (date &&\
        minimap2 -x sr -d {output} {input} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Build a bwa index (may be used to map short reads against it, or long reads by using the "-x ont2d" flag)
rule build_bwa_index:
    input: "{prefix}.fna"
    output: index_amb="{prefix}.amb",
       index_ann="{prefix}.ann",
       index_bwt="{prefix}.bwt",
       index_pac="{prefix}.pac",
       index_sa="{prefix}.sa"
    log: "{prefix}.bwa_index"
    conda: "envs/bwa.yaml"
    shell:
        """
        (date &&\
        bwa index {input} -p {wildcards.prefix} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Simple utility function
rule convert_fastq_to_fasta:
    input: "{prefix}.fastq.gz"
    output: "{prefix}.fna"
    shell:
        """
        date
        zcat {input} | sed -n '1~4s/^@/>/p;2~4p' > {output}
        date
        """

rule convert_fq_to_fasta:
    input: "{prefix}.fq.gz"
    output: "{prefix}.fna"
    shell:
        """
        date
        zcat {input} | sed -n '1~4s/^@/>/p;2~4p' > {output}
        date
        """

# Map long reads to a reference (IGC/HG38/other). The resulting BAM file can be use for many things, among other to estimate which long reads are "known" or "novel", or to estimate the coverage of reference sequences (regions thereof).
rule map_long_reads_to_reference_w_minimap2:
    input: 
        query=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"),
        index=os.path.join(DB_DIR, "{reference}/{reference}-ont.mmi"),
        index_fa=os.path.join(DB_DIR, "{reference}/{reference}.fna")
    output: os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}-x-{reference}.bam")
    params:
        prefix=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}-x-{reference}")
    threads: config["minimap2"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}-x-{reference}.minimap2_samtools")
    conda: "envs/minimap2.yaml"
    shell:
        """
        (date &&\
        minimap2 -a -t {threads} {input.index} {input.query} | samtools view -bT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Preprocess the short reads using fastp
rule run_fastp_on_short_reads:
    input:
        r1=os.path.join(config["short_reads_prefix"], "{sr_sample}_R1_001.fastq.gz"),
        r2=os.path.join(config["short_reads_prefix"], "{sr_sample}_R2_001.fastq.gz")
    output:
        r1=os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}_R1_001.fastp.fq.gz"),
        r2=os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}_R2_001.fastp.fq.gz"),
        html=os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}.fastp.html"),
        json=os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}.fastp.json")
    log: os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}.fastp")
    conda: "envs/fastp.yaml"
    shell:
        """
        (date &&\
        fastp -l {config[fastp][min_length]} -i {input.r1} -I {input.r2} -o {output.r1} -O {output.r2} -h {output.html} -j {output.json} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# TODO: Simplify paths by removing the `sr`.
rule run_fastqc_on_preprocessed_short_reads:
    input: os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}_{suffix}.fq.gz"),
    output:
        html=os.path.join(RESULTS_DIR, "qc/sr/fastqc/{sr_sample}/{sr_sample}_{suffix}.fastqc.html"),
        zip=os.path.join(RESULTS_DIR, "qc/sr/fastqc/{sr_sample}/{sr_sample}_{suffix}.fastqc.zip") # the suffix _fastqc.zip is necessary for multiqc to find the file. If not using multiqc, you are free to choose an arbitrary filename
    conda: "envs/fastqc.yaml"
    script: "scripts/run_fastqc.py" # Use a separate script rather than a simple call to FASTQC to overcome possible race conditions when running in parallel (https://bitbucket.org/snakemake/snakemake-wrappers/raw/7f5e01706728e32d52aa413f213e950f50bf4129/bio/fastqc/wrapper.py)

rule map_short_reads_to_reference_w_minimap2:
    input: 
        query_r1=rules.run_fastp_on_short_reads.output.r1,
        query_r2=rules.run_fastp_on_short_reads.output.r2,
        index=os.path.join(DB_DIR, "{reference}/{reference}-sr.mmi"),
        index_fa=os.path.join(DB_DIR, "{reference}/{reference}.fna")
    output: os.path.join(RESULTS_DIR, "mapping/sr/minimap2/{sr_sample}-x-{reference}.bam")
    params:
        prefix=os.path.join(RESULTS_DIR, "mapping/sr/minimap2/{sr_sample}-x-{reference}")
    threads: config["minimap2"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/sr/minimap2/{sr_sample}-x-{reference}.minimap2_samtools")
    conda: "envs/minimap2.yaml"
    shell:
        """
        (date &&\
        minimap2 -a -t {threads} {input.index} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -bT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule map_short_reads_to_reference_w_bwa_mem:
    input: 
        query_r1=rules.run_fastp_on_short_reads.output.r1,
        query_r2=rules.run_fastp_on_short_reads.output.r2,
        index_amb=os.path.join(DB_DIR, "{reference}/{reference}.amb"),
        index_ann=os.path.join(DB_DIR, "{reference}/{reference}.ann"),
        index_bwt=os.path.join(DB_DIR, "{reference}/{reference}.bwt"),
        index_pac=os.path.join(DB_DIR, "{reference}/{reference}.pac"),
        index_sa=os.path.join(DB_DIR, "{reference}/{reference}.sa"),
        index_fa=os.path.join(DB_DIR, "{reference}/{reference}.fna")
    output: os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}-x-{reference}.bam")
    params:
        index_prefix=os.path.join(DB_DIR, "{reference}/{reference}"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}-x-{reference}")
    threads: config["bwa"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}-x-{reference}.bwa_mem_samtools")
    conda: "envs/bwa.yaml"
    shell:
        """
        (date &&\
        bwa mem -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -bT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Compute the genome coverage histogram - https://bedtools.readthedocs.io/en/latest/content/tools/genomecov.html
# The resulting histogram is used to compute the expected, i.e., average, coverage of the underyling reference (which may be de novo contigs).
rule run_genomecov:
    input: os.path.join(RESULTS_DIR, "mapping/{prefix}.bam")
    output: os.path.join(RESULTS_DIR, "genomecov/{prefix}.cov.txt")
    conda: "envs/bedtools.yaml"
    shell:
        """
        date
        bedtools genomecov -ibam {input} > {output}
        date
        """

# Compute the expected, i.e., average, coverage given the coverage histogram
rule compute_avg_genomecov:
    input: os.path.join(RESULTS_DIR, "genomecov/{prefix}.cov.txt")
    output: os.path.join(RESULTS_DIR, "genomecov/{prefix}.avg_cov.txt")
    shell:
        """
        date
        cat {input} | awk -f {config[compute_avg_coverage][bin]} | tail -n+2 > {output}
        date
        """

# TODO: All above rules that build an index build it in the same path as the reference/target file
# In fact, it is redundant with `build_bwa_index`. The only advantage is that this replication allows more fine-grained control of the resources (slurm, etc.). This is, however, probably not needed as `build_bwa_index` can be quite resource-intensive, e.g., for the IGC.
# Similar to building a minimap2 index, but this time for bwa
rule build_long_read_bwa_index:
    input: rules.get_single_fastq_per_barcode.output
    output: os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}.amb"),
            os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}.ann"),
            os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}.bwt"),
            os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}.pac"),
            os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}.sa")
    params:
        prefix=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}")
    log: os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}.bwa_index")
    conda: "envs/bwa.yaml"
    shell:
        """
        (date &&\
        bwa index {input} -p {params.prefix} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """
# TODO: If TODO-build_long_read_bwa_index is completed, need to adjust the `input.index` here too
# config[bwa][long_reads_index][opts] according to https://github.com/sfu-compbio/colormap/blob/baa6802b43da8640c305032b7f7fcadd7c496a03/run_corr.sh#L45 and https://academic.oup.com/bioinformatics/article/32/17/i545/2450793#112481513 -> "2.2.3 Mapping parameters"
rule map_short_reads_to_long_reads:
    input: 
        query_r1=rules.run_fastp_on_short_reads.output.r1,
        query_r2=rules.run_fastp_on_short_reads.output.r2,
        index=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}.amb"),
        index_fa=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fna"),
    output: os.path.join(RESULTS_DIR, "mapping/lr/merged/{sr_sample}-x-{barcode}.bam")
    threads: config["bwa"]["threads"] 
    params:
        index_prefix=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/lr/merged/{sr_sample}-x-{barcode}.")
    log: os.path.join(RESULTS_DIR, "mapping/lr/merged/{sr_sample}-x-{barcode}.bwa_mem_samtools")
    conda: "envs/bwa.yaml"
    shell:
        """
        (date &&\
        bwa mem {config[bwa][long_reads_index][opts]} -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -S -b -u -T {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# TODO: Add extra rules for other assemblers instead of a generic rule as this offers then greater flexibility, e.g., in terms of resource-reservation.
rule assemble_long_reads_w_flye:
    input: rules.get_single_fastq_per_barcode.output
    output: os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/{barcode}/assembly.fna")
    threads: config["flye"]["threads"]
    log: os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/{barcode}/assembly.flye")
    conda: "envs/flye.yaml"
    shell:
        """
        (date &&\
        flye --nano-raw {input} --meta --out-dir $(dirname {output}) --genome-size {config[flye][genome_size]} --threads {threads} &&\
        ln -s assembly.fasta {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# "r941_min_high" is the MinION model, high accuarcy
rule polish_long_read_contigs_w_medaka:
    input: 
        basecalls=rules.get_single_fastq_per_barcode.output,
        contigs=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.fna")
    output: os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/medaka/round_{round}/lr/merged/{barcode}/assembly_polished.fna")
    threads: config["medaka"]["threads"]
    log: os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/medaka/round_{round}/lr/merged/{barcode}/assembly_polished.medaka")
    conda: "envs/medaka.yaml"
    shell:
        """
        (date &&\
        medaka_consensus -i {input.basecalls} -d {input.contigs} -o $(dirname {output}) -t {threads} -m r941_min_high &&\
        ln -s consensus.fasta {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Having the "initial" explicitly in the rule name is not ideal. However, the path to the original contigs is somewhat variable and hence a discriminative rule name was chosen.
# As piping is no longer supported by racon since v. 1.0.0 (https://github.com/isovic/racon/issues/46#issuecomment-519380742) and the file extensions are checked for all files including the contigs), some temporary files have to be created and removed afterwards.
rule polish_long_read_contigs_w_racon_initial_round:
    input: 
        basecalls=rules.get_single_fastq_per_barcode.output,
        contigs=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.fna"),
        lr_to_lr_contigs_mapping=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam"),
    output: os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/racon/round_{round}/lr/merged/{barcode}/assembly_polished.fna")
    params:
        tmp_sam=os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/racon/round_{round}/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.tmp.sam"),
        tmp_contigs=os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/racon/round_{round}/lr/merged/{barcode}/assembly_polished.tmp_draft.fa")
    threads: config["racon"]["threads"]
    log: os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/racon/round_{round}/lr/merged/{barcode}/assembly_polished.racon.log")
    conda: "envs/racon.yaml"
    shell:
        """
        (date &&\
        samtools view -h {input.lr_to_lr_contigs_mapping} > {params.tmp_sam} &&\
        cp {input.contigs} {params.tmp_contigs} &&\
        racon -t {threads} {input.basecalls} {params.tmp_sam} {params.tmp_contigs} > {output} &&\
        rm {params.tmp_sam} {params.tmp_contigs} &&\
        date) 2> >(tee {log}) > >(tee {log})
        """

#polished_contigs=os.path.join(RESULTS_DIR, "assembly/operams/lr_{barcode}-sr_{sr_sample}/contigs.polished.fna")
#ln -s contigs.polished.fasta {output.polished_contigs} &&\
#rule assemble_long_and_short_reads_w_operams:
#    input: 
#        lr=rules.get_single_fastq_per_barcode.output,
#        sr_r1=rules.run_fastp_on_short_reads.output.r1,
#        sr_r2=rules.run_fastp_on_short_reads.output.r2,
#        sr_contigs=os.path.join(RESULTS_DIR, "assembly/megahit/{sr_sample}/final.contigs.fna")
#    output: 
#        contigs=os.path.join(RESULTS_DIR, "assembly/operams/lr_{barcode}-sr_{sr_sample}/contigs.fna")
#    params:
#        tmp_lr=os.path.join(RESULTS_DIR, "assembly/operams/lr_{barcode}-sr_{sr_sample}/tmp_lr_{barcode}.fna")
#    threads: config["operams"]["threads"]
#    log: os.path.join(RESULTS_DIR, "assembly/operams/lr_{barcode}-sr_{sr_sample}/opera_ms.log")
#    shell:
#        """
#        (date &&\
#        zcat {input.lr} > {params.tmp_lr} &&\
#        {config[operams][bin]} --long-read-mapper minimap2 --num-processors {threads} --short-read1 {input.sr_r1} --short-read2 {input.sr_r2} --long-read {params.tmp_lr} --contig-file {input.sr_contigs} --out-dir $(dirname {output.contigs}) &&\
#        ln -s contigs.fasta {output.contigs} &&\
#        rm {params.tmp_lr} &&\
#        date) 2> >(tee {log}) > >(tee {log})
#        """

#zcat {input.lr} > {params.tmp_lr} &&\
rule assemble_long_and_short_reads_w_metaspades:
    input: 
        lr=rules.get_single_fastq_per_barcode.output,
        sr_r1=rules.run_fastp_on_short_reads.output.r1,
        sr_r2=rules.run_fastp_on_short_reads.output.r2,
        sr_contigs=os.path.join(RESULTS_DIR, "assembly/megahit/{sr_sample}/final.contigs.fna")
    output: 
        contigs=os.path.join(RESULTS_DIR, "assembly/metaspades_hybrid/lr_{barcode}-sr_{sr_sample}/contigs.fna")
    params:
        tmp_lr=os.path.join(RESULTS_DIR, "assembly/metaspades_hybrid/lr_{barcode}-sr_{sr_sample}/tmp_lr_{barcode}.fna")
    threads: config["metaspades"]["threads"]
    log: os.path.join(RESULTS_DIR, "assembly/metaspades_hybrid/lr_{barcode}-sr_{sr_sample}/metaspades.log")
    conda: "envs/spades.yaml"
    shell:
        """
        (date &&\
        spades.py --meta -o $(dirname {output.contigs}) -k 21,33,55,77 -t {threads} -1 {input.sr_r1} -2 {input.sr_r2} --nanopore {input.lr} &&\
        ln -s contigs.fasta {output.contigs} &&\
        date) &> >(tee {log})
        """

rule map_short_reads_to_hybrid_contigs_w_bwa_mem:
    input: 
        query_r1=rules.run_fastp_on_short_reads.output.r1,
        query_r2=rules.run_fastp_on_short_reads.output.r2,
        index_amb=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.amb"),
        index_ann=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.ann"),
        index_bwt=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.bwt"),
        index_pac=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.pac"),
        index_sa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.sa"),
        index_fa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.fna")
    output: os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.bam")
    params:
        index_prefix=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs")
    log: os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.bwa_mem_samtools")
    conda: "envs/bwa.yaml"
    threads: config["bwa"]["threads"]
    shell:
        """
        (date &&\
        bwa mem -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule map_long_reads_to_hybrid_contigs_w_bwa_mem:
    input: 
        query_lr=rules.get_single_fastq_per_barcode.output,
        index_amb=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.amb"),
        index_ann=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.ann"),
        index_bwt=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.bwt"),
        index_pac=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.pac"),
        index_sa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.sa"),
        index_fa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.fna")
    output: os.path.join(RESULTS_DIR, "mapping/lr/bwa_mem/{barcode}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.bam")
    params:
        index_prefix=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/lr/bwa_mem/{barcode}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs")
    conda: "envs/bwa.yaml"
    threads: config["bwa"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/lr/bwa_mem/{barcode}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.bwa_mem_samtools")
    shell:
        """
        (date &&\
        bwa mem -x ont2d -t {threads} {params.index_prefix} {input.query_lr} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# rule build_minimap2_long_read_contigs_ont_index:
#     input: os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/{prefix}.fna")
#     output: os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/{prefix}-ont.mmi")
#     conda: "envs/minimap2.yaml"
#     shell:
#         """
#         date
#         minimap2 -x map-ont -d {output} {input}
#         date
#         """

rule map_long_reads_to_long_read_contigs:
    input: 
        query=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"),
        index=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly-ont.mmi"),
        index_fa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.fna")
    output: os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam")
    params:
        prefix=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x_{barcode}_contigs")
    conda: "envs/minimap2.yaml"
    threads: config["minimap2"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.minimap2_samtools")
    shell:
        """
        (date &&\
        minimap2 -a -t {threads} {input.index} {input.query} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule map_long_reads_to_long_read_polished_contigs:
    input: 
        query=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"),
        index=os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/{polisher}/round_{round}/lr/merged/{barcode}/assembly_polished-ont.mmi"),
        index_fa=os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/{polisher}/round_{round}/lr/merged/{barcode}/assembly_polished.fna")
    output: os.path.join(RESULTS_DIR, "mapping/polished/{polisher}/round_{round}/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_polished_contigs.bam")
    params:
        prefix=os.path.join(RESULTS_DIR, "mapping/polished/{polisher}/round_{round}/lr/merged/{barcode}/{barcode}_reads-x_{barcode}_polished_contigs")
    conda: "envs/minimap2.yaml"
    threads: config["minimap2"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/polished/{polisher}/round_{round}/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_polished_contigs.minimap2_samtools")
    shell:
        """
        (date &&\
        minimap2 -a -t {threads} {input.index} {input.query} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# rule build_minimap2_long_read_contigs_sr_index:
#     input: os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/{prefix}.fna")
#     output: os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/{prefix}-sr.mmi")
#     conda: "envs/minimap2.yaml"
#     shell:
#         """
#         date
#         minimap2 -x sr -d {output} {input}
#         date
#         """

rule map_short_reads_to_long_read_contigs_w_minimap2:
    input: 
        query_r1=rules.run_fastp_on_short_reads.output.r1,
        query_r2=rules.run_fastp_on_short_reads.output.r2,
        index=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly-sr.mmi"),
        index_fa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.fna")
    output: os.path.join(RESULTS_DIR, "mapping/sr/minimap2/{sr_sample}_reads-x-{barcode}-{assembler}_contigs.bam")
    params:
        prefix=os.path.join(RESULTS_DIR, "mapping/sr/minimap2/{sr_sample}_reads-x_{barcode}_contigs")
    conda: "envs/minimap2.yaml"
    threads: config["minimap2"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/sr/minimap2/{sr_sample}_reads-x-{barcode}-{assembler}_contigs.minimap2_samtools")
    shell:
        """
        (date &&\
        minimap2 -a -t {threads} {input.index} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.prefix} > {output}
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Also map the short reads with bwa mem as it might be more sensitive (https://github.com/lh3/minimap2/issues/214#issuecomment-413492006)
rule map_short_reads_to_long_read_contigs_w_bwa_mem:
    input: 
        query_r1=rules.run_fastp_on_short_reads.output.r1,
        query_r2=rules.run_fastp_on_short_reads.output.r2,
        index_amb=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.amb"),
        index_ann=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.ann"),
        index_bwt=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.bwt"),
        index_pac=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.pac"),
        index_sa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.sa"),
        index_fa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.fna")
    output: os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x-{barcode}-{assembler}_contigs.bam")
    params:
        index_prefix=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x_{barcode}_contigs")
    conda: "envs/bwa.yaml"
    threads: config["bwa"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x-{barcode}-{assembler}_contigs.bwa_mem_samtools")
    shell:
        """
        (date &&\
        bwa mem -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule assemble_short_reads_w_megahit:
    input: 
        sr_r1=rules.run_fastp_on_short_reads.output.r1,
        sr_r2=rules.run_fastp_on_short_reads.output.r2
    output: os.path.join(RESULTS_DIR, "assembly/megahit/{sr_sample}/final.contigs.fna")
    threads: config["megahit"]["threads"]
    conda: "envs/megahit.yaml"
    params:
        tmp_dir=os.path.join(RESULTS_DIR, "assembly/megahit/{sr_sample}_tmp")
    log: os.path.join(RESULTS_DIR, "assembly/megahit/{sr_sample}/final.contigs.megahit")
    shell:
        """
        (date &&\
        megahit -1 {input.sr_r1} -2 {input.sr_r2} -t {threads} -o {params.tmp_dir} &&\
        mv {params.tmp_dir}/* $(dirname {output}) &&\
        rmdir {params.tmp_dir} &&\
        ln -s final.contigs.fa {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# N.B. A similar function exists above and code duplication is generally discouraged.
rule map_short_reads_to_short_read_contigs_w_bwa_mem:
    input: 
        query_r1=rules.run_fastp_on_short_reads.output.r1,
        query_r2=rules.run_fastp_on_short_reads.output.r2,
        index_amb=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.amb"),
        index_ann=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.ann"),
        index_bwt=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.bwt"),
        index_pac=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.pac"),
        index_sa=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.sa"),
        index_fa=rules.assemble_short_reads_w_megahit.output
    output: os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x-{sr_sample}-{assembler}_contigs.bam")
    params:
        index_prefix=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x-{sr_sample}-{assembler}_contigs")
    conda: "envs/bwa.yaml"
    threads: config["bwa"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/{sr_sample}_reads-x-{sr_sample}-{assembler}_contigs.bwa_mem_samtools")
    shell:
        """
        (date &&\
        bwa mem -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule map_long_reads_to_short_read_contigs_w_bwa_mem:
    input: 
        query_lr=rules.get_single_fastq_per_barcode.output,
        index_amb=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.amb"),
        index_ann=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.ann"),
        index_bwt=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.bwt"),
        index_pac=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.pac"),
        index_sa=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.sa"),
        index_fa=rules.assemble_short_reads_w_megahit.output
    output: os.path.join(RESULTS_DIR, "mapping/lr/bwa_mem/{barcode}_reads-x-{sr_sample}-{assembler}_contigs.bam")
    params:
        index_prefix=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/lr/bwa_mem/{barcode}_reads-x-{sr_sample}-{assembler}_contigs")
    conda: "envs/bwa.yaml"
    threads: config["bwa"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/lr/bwa_mem/{barcode}_reads-x-{sr_sample}-{assembler}_contigs.bwa_mem_samtools")
    shell:
        """
        (date &&\
        bwa mem -x ont2d -t {threads} {params.index_prefix} {input.query_lr} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """
#npa, npc, .npl, and .npo
rule run_nonpareil_on_short_reads:
    input: os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}_R{orientation}_{suffix}.fastp.fq.gz"),
    output: 
        os.path.join(RESULTS_DIR, "annotation/nonpareil/{sr_sample}_R{orientation}_{suffix}.fastp.npa"),
        os.path.join(RESULTS_DIR, "annotation/nonpareil/{sr_sample}_R{orientation}_{suffix}.fastp.npc"),
        os.path.join(RESULTS_DIR, "annotation/nonpareil/{sr_sample}_R{orientation}_{suffix}.fastp.npl"),
        os.path.join(RESULTS_DIR, "annotation/nonpareil/{sr_sample}_R{orientation}_{suffix}.fastp.npo")
    params:
        out_prefix=os.path.join(RESULTS_DIR, "annotation/nonpareil/{sr_sample}_R{orientation}_{suffix}.fastp"),
        tmp_fasta=os.path.join(RESULTS_DIR, "annotation/nonpareil/{sr_sample}_R{orientation}_{suffix}.tmp.fna")
    threads: config["nonpareil"]["threads"]
    log: os.path.join(RESULTS_DIR, "annotation/nonpareil/{sr_sample}_R{orientation}_{suffix}.fastp.nonpareil")
    conda: "envs/nonpareil.yaml"
    shell:
        """
        (date &&\
        zcat {input} | sed -n '1~4s/^@/>/p;2~4p' > {params.tmp_fasta} &&\
        nonpareil -s {params.tmp_fasta} -R {config[nonpareil][memory]} -t {threads} -T kmer -f fasta -b {params.out_prefix} &&\
        rm {params.tmp_fasta} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# TODO: Cleanup how the input files are defined. Currently explicitly specified, which is discouraged to do.
rule create_nanopolish_index:
    input:
        multifast5_dir_01=os.path.join(DATA_DIR, "multifast5/20181108_0827_test"),
        multifast5_dir_02=os.path.join(DATA_DIR, "multifast5/20181107_0906_same"),
        multifast5_dir_03=os.path.join(DATA_DIR, "multifast5/20181106_1450_noselection_sizeselection"),
        reads=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz")
    output:
        index=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index"),
        fai=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index.fai"),
        index_gzi=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index.gzi"),
        readdb=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index.readdb")
    log: os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.nanopolish_index")
    conda: "envs/nanopolish.yaml"
    shell:
        """
        (date &&\
        nanopolish index -d {input.multifast5_dir_01}/ -d {input.multifast5_dir_02}/ -d {input.multifast5_dir_03}/ {input.reads} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule samtools_index:
    input: "{prefix}.bam"
    output: "{prefix}.bam.bai"
    conda: "envs/bwa.yaml"
    shell:
        """
        date
        samtools index {input} {output}
        date
        """

# # Legacy version before snakemake's `checkpoint` was used
# rule polish_long_read_contigs_w_nanopolish:
#     input:
#         lr_to_lr_contigs_mapping=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam"),
#         lr_to_lr_contigs_mapping_index=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam.bai"),
#         reads=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"),
#         draft_contigs=os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/{barcode}/assembly.fna"),
#         nanopolish_index=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index"),
#         nanopolish_fai=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index.fai"),
#         nanopolish_index_gzi=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index.gzi"),
#         nanopolish_readdb=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz.index.readdb")
#     output: os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/nanopolish/round_{round}/lr/merged/{barcode}/assembly_polished.fna")
#     log: os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/nanopolish/round_{round}/lr/merged/{barcode}/assembly_polished.nanopolish")
#     threads: 4
#     conda: "envs/nanopolish.yaml"
#     shell:
#         """
#         (date &&\
#         nanopolish_makerange.py {input.draft_contigs} | parallel --results $(dirname {output}) -P 7 \
#         nanopolish variants --consensus -o $(dirname {output})/polished.{{1}}.vcf -w {{1}} -r {input.reads} -b {input.lr_to_lr_contigs_mapping} -g {input.draft_contigs} -t {threads} --min-candidate-frequency 0.1 &&\
#         nanopolish vcf2fasta -g {input.draft_contigs} $(dirname {output})/polished.*.vcf > {output} &&\
# 
#         date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
#         """

# Define the individual windows that nanopolish will parallelize over
checkpoint nanopolish_makerange:
    input: os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/{barcode}/assembly.fna")
    output: directory(os.path.join(RESULTS_DIR, "assembly/flye/polished/nanopolish/round_{round}/lr/merged/{barcode}/windows"))
    conda: "envs/nanopolish.yaml"
    shell:
        """
        mkdir {output}
        nanopolish_makerange.py {input} > {output}/ranges.txt
        for i in $(head -n 20 {output}/ranges.txt); do echo $i > {output}/"$i".coord; done
        date
        """

# Perform variant calling on single window
rule nanopolish_variant_call:
    input: 
        reads=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"),
        lr_to_lr_contigs_mapping=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam"),
        lr_to_lr_contigs_mapping_index=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam.bai"),
        draft_contigs=os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/{barcode}/assembly.fna"),
        window=os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/nanopolish/round_{round}/lr/merged/{barcode}/windows/{window}.coord")
    output: os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/nanopolish/round_{round}/lr/merged/{barcode}/vcfs/{window}.vcf")
    threads: 4
    conda: "envs/nanopolish.yaml"
    shell:
        """
        nanopolish variants --consensus -o {output} -w {wildcards.window} -r {input.reads} -b {input.lr_to_lr_contigs_mapping} -g {input.draft_contigs} -t {threads} --min-candidate-frequency 0.1
        """

# Gather all the individual VCF files that were produced.
# The "overview" comes from the rule that created the "split-up", i.e., from the `nanopolish_makerange` rule.
# But the actual files that are aggregated here are those that should eventually be created and used, i.e., the VCF files
# N.B. snakemake threw an error when the `window=` line below was wrong, but the error was supposedly coming from `nanopolish_makerange`. So much about tracking errors, which is not always straightforward as one might wish.
def aggregate_nanopolish_vcf(wildcards):
    checkpoint_output = checkpoints.nanopolish_makerange.get(**wildcards).output[0]
    return expand(os.path.join(RESULTS_DIR, "assembly/flye/polished/nanopolish/round_{round}/lr/merged/{barcode}/vcfs/{window}.vcf"),
            round=wildcards.round,
            barcode=wildcards.barcode,
            window=glob_wildcards(os.path.join(checkpoint_output, "{window}.coord")).window)

# Combine the information in the individual VCF files to produce a polished assembly
rule polish_long_read_contigs_w_nanopolish:
    input: 
        vcf_calls=aggregate_nanopolish_vcf,
        draft_contigs=os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/{barcode}/assembly.fna"),
    output: os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/nanopolish/round_{round}/lr/merged/{barcode}/assembly_polished.fna")
    log: os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/nanopolish/round_{round}/lr/merged/{barcode}/assembly_polished.nanopolish")
    conda: "envs/nanopolish.yaml"
    shell:
        """
        (date &&\
        nanopolish vcf2fasta -g {input.draft_contigs} {input.vcf_calls} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule nanopolish_methylation_call:
    input: 
        reads=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"),
        lr_to_lr_contigs_mapping=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam"),
        lr_to_lr_contigs_mapping_index=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}/{barcode}_reads-x-{barcode}-{assembler}_contigs.bam.bai"),
        draft_contigs=os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/{barcode}/assembly.fna"),
        window=os.path.join(RESULTS_DIR, "assembly/{assembler}/polished/nanopolish/round_{round}/lr/merged/{barcode}/windows/{window}.coord")
    output: os.path.join(RESULTS_DIR, "annotation/methylation/{assembler}/nanopolish/round_{round}/lr/merged/{barcode}/methylation_calls/{window}.methylation_calls.tsv")
    threads: 4
    conda: "envs/nanopolish.yaml"
    shell:
        """
        nanopolish call-methylation -t {threads} -r {input.reads} -b {input.lr_to_lr_contigs_mapping} -g {input.draft_contigs} -w {wildcards.window} > {output}
        """

def aggregate_nanopolish_methylation(wildcards):
    checkpoint_output = checkpoints.nanopolish_makerange.get(**wildcards).output[0]
    return expand(os.path.join(RESULTS_DIR, "annotation/methylation/{assembler}/nanopolish/round_{round}/lr/merged/{barcode}/methylation_calls/{window}.methylation_calls.tsv"),
            assembler=wildcards.assembler,
            barcode=wildcards.barcode,
            round=wildcards.round,
            window=glob_wildcards(os.path.join(checkpoint_output, "{window}.coord")).window)

# Combine the information in the individual methylation call files to get the overall methylation state
rule merge_methylation_calls_from_long_read_contigs_w_nanopolish:
    input: aggregate_nanopolish_methylation,
    output: os.path.join(RESULTS_DIR, "annotation/methylation/{assembler}/nanopolish/round_{round}/lr/merged/{barcode}/methylation_frequencies.tsv")
    log: os.path.join(RESULTS_DIR, "annotation/methylation/{assembler}/nanopolish/round_{round}/lr/merged/{barcode}/methylation_frequencies")
    shell:
        """
        (date &&\
        scripts/calculate_methylation_frequency.py {input} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

# Works until the final shred and polish step, then crashes with a Key error
# rule run_rebaler_on_short_read_contigs:
#     input:
#         lr=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"),
#         sr_contigs=rules.assemble_short_reads_w_megahit.output
#     output: os.path.join(RESULTS_DIR, "assembly/rebaler/lr/merged/{barcode}/lr_{barcode}-sr_contigs_{sr_sample}.fna")
#     threads: config["rebaler"]["threads"]
#     params:
#         tmp_sr_contigs=os.path.join(RESULTS_DIR, "assembly/rebaler/lr/merged/{barcode}/tmp_sr_contigs.fna")
#     conda: "envs/rebaler.yaml"
#     shell:
#         """
#         date
#         awk 'BEGIN{{counter=0}} FNR==NR{{if($1 ~ /^>/) {{a[$1]=counter; counter=counter+1}};next}}{{if($1 ~/^>/) print ">"a[$1]; else print}}' {input.sr_contigs} {input.sr_contigs} > {params.tmp_sr_contigs} # Rename the contigs from 0 to num. contigs to ensure rebaler has no "Key error"
#         rebaler -t {threads} {params.tmp_sr_contigs} {input.lr} > {output}
#         rm {params.tmp_sr_contigs}
#         date
#         """

# For info on "--longs-reads" option see https://github.com/bbuchfink/diamond/releases/tag/v0.9.23
rule run_diamond_on_long_reads:
    input:
        reads=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"),
        db=config["diamond"]["db"]
    output: os.path.join(RESULTS_DIR, "annotation/diamond/lr/merged/{barcode}.diamond.daa")
    threads: config["diamond"]["threads"]
    conda: "envs/diamond.yaml"
    shell:
        """
        date
        diamond blastx -q {input.reads} --db {input.db} -p {threads} --long-reads --sensitive --outfmt 100 --out {output}
        date
        """

rule call_genes_w_prodigal:
    input: os.path.join(RESULTS_DIR, "assembly/{prefix}.fna")
    output: os.path.join(RESULTS_DIR, "annotation/proteins/{prefix}.faa")
    log: os.path.join(RESULTS_DIR, "annotation/proteins/{prefix}.prodigal.log")
    conda: "envs/prodigal.yaml"
    shell: 
        """
        (date &&\
        prodigal -a {output} -p meta -i {input} &&\
        date) &> >(tee {log})
        """

# Write to DAA format to be able to use `diamond view` later to define custom formats.
# N.B. `diamond view` seems not to support compressed input AND seems to require the *prefix* of the DAA file.
# TODO: This could maybe be parametrized by the respectively used DB
rule run_diamond_on_proteins:
    input: 
        proteins=rules.call_genes_w_prodigal.output,
        db=config["diamond"]["db"]
    output: os.path.join(RESULTS_DIR, "annotation/diamond/{prefix}.daa")
    params:
        outfmt="100"
    log: os.path.join(RESULTS_DIR, "annotation/diamond/{prefix}.diamond.log")
    conda: "envs/diamond.yaml"
    threads: config["diamond"]["threads"]
    shell: 
        """
        (date &&\
        diamond blastp -q {input.proteins} --db {input.db} -p {threads} --outfmt {params.outfmt} --out {output} &&\
        date) &> >(tee {log})
        """

# Not sure why, but with the conda version of diamond (and `diamond view`) it was necessary to specify the entire column layout (outfmt) explicitly
rule format_diamond_output_for_ideel:
    input: os.path.join(RESULTS_DIR, "annotation/diamond/{prefix}.daa")
    output: os.path.join(RESULTS_DIR, "annotation/diamond/{prefix}.tsv")
    params:
        outfmt="6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore qlen slen"
    log: os.path.join(RESULTS_DIR, "annotation/diamond/{prefix}.diamond_view.log")
    conda: "envs/diamond.yaml"
    shell: 
        """
        (date &&\
        diamond view --daa $(echo {input} | sed 's/\.daa$//') --max-target-seqs 1 --outfmt {params.outfmt} --out {output} &&\
        date) &> >(tee {log})
        """

## MMSEQ2 analyses
rule mmseq2_database:
    input: 
        int1=expand(os.path.join(RESULTS_DIR, "annotation/proteins/{assembler}/lr/merged/{barcode}/assembly.faa"), assembler="flye", barcode=BARCODES),
        int2=expand(os.path.join(RESULTS_DIR, "annotation/proteins/{assembler}/{sample}/final.contigs.faa"), assembler="megahit", sample=SAMPLES),
        int3=expand(os.path.join(RESULTS_DIR, "annotation/proteins/{assembler}/lr_{barcode}-sr_{sample}/contigs.faa"), assembler="metaspades_hybrid", barcode=BARCODES, sample=SAMPLES)
#        expand(os.path.join(RESULTS_DIR, "annotation/proteins/{assembler}/{prefix}.faa"), assembler=["flye", "megahit", "metaspades_hybrid"])
    output: 
        dat1=expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/{assembler}_db"), assembler="flye"),
        dat2=expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/{assembler}_db"), assembler="megahit"),
        dat3=expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/{assembler}_db"), assembler="metaspades_hybrid")
    log: os.path.join(RESULTS_DIR, "annotation/mmseq2/database.mmseq2.log")
    conda: "cd-hit.yaml"
    shell:
        """
        (date &&\
        mmseqs createdb {input.int1} {output.dat1} &&\
        mmseqs createdb {input.int2} {output.dat2} &&\
        mmseqs createdb {input.int3} {output.dat3} &&\
        date) &> >(tee {log})
        """

rule mmseq2_compare:
    input:
        mm1=expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/{assembler}_db"), assembler="flye"),
        mm2=expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/{assembler}_db"), assembler="megahit"),
        mm3=expand(os.path.join(RESULTS_DIR, "annotation/mmseq2/{assembler}_db"), assembler="metaspades_hybrid")
    output: 
        mo1=os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_megahit_rbh"),
        mo2=os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_metaspades_hybrid_rbh"),
        mo3=os.path.join(RESULTS_DIR, "annotation/mmseq2/megahit_metaspades_hybrid_rbh")
    log: os.path.join(RESULTS_DIR, "annotation/mmseq2/compare.mmseq2.log")
    threads: config["mmseq2"]["threads"]
    conda: "cd-hit.yaml"
    shell:
        """
        (date &&\
        mmseqs rbh {input.mm1} {input.mm2} {output.mo1} --min-seq-id 0.9 mmseq2_tmp --threads {threads} &&\
        mmseqs rbh {input.mm1} {input.mm3} {output.mo2} --min-seq-id 0.9 mmseq2_tmp --threads {threads} &&\
        mmseqs rbh {input.mm2} {input.mm3} {output.mo3} --min-seq-id 0.9 mmseq2_tmp --threads {threads} &&\
        date) &> >(tee {log})
        """

rule mmseq2_m8_format:
    input:
        fo1=os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_db"),
        fo2=os.path.join(RESULTS_DIR, "annotation/mmseq2/megahit_db"),
        fo3=os.path.join(RESULTS_DIR, "annotation/mmseq2/metaspades_hybrid_db"),
        fo4=os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_megahit_rbh"),
        fo5=os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_metaspades_hybrid_rbh"),
        fo6=os.path.join(RESULTS_DIR, "annotation/mmseq2/megahit_metaspades_hybrid_rbh")
    output: 
        form1=os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_megahit.m8"),
        form2=os.path.join(RESULTS_DIR, "annotation/mmseq2/flye_metaspades_hybrid.m8"),
        form3=os.path.join(RESULTS_DIR, "annotation/mmseq2/megahit_metaspades_hybrid.m8")
    log: os.path.join(RESULTS_DIR, "annotation/mmseq2/convertalis.mmseq2.log")
    conda: "cd-hit.yaml"
    shell:
        """
        (date &&\
        mmseqs convertalis {input.fo1} {input.fo2} {input.fo4} {output.form1}  &&\
        mmseqs convertalis {input.fo1} {input.fo3} {input.fo5} {output.form2} &&\
        mmseqs convertalis {input.fo2} {input.fo3} {input.fo6} {output.form3} &&\
        date) &> >(tee {log})
        """

rule prepare_plot_files:
    output:
        touch(os.path.join(RESULTS_DIR, "annotation/mmseq2/plot_files_ready.done"))
    log: os.path.join(RESULTS_DIR, "annotation/mmseq2/files_ready.mmseq2.log")
    shell:
        """
        (date &&\
        ./prepare_plot_files.sh &&\
        date) &> >(tee {log})
        """

rule mmseq2_plots:
    input:
        plot1=os.path.join(RESULTS_DIR, "annotation/mmseq2/overlap_sizes.txt"),
        plot2=os.path.join(RESULTS_DIR, "annotation/mmseq2/total.txt")
    output: 
        touch(os.path.join(RESULTS_DIR, "annotation/mmseq2/upset_plots.done"))
    log: os.path.join(RESULTS_DIR, "annotation/mmseq2/plot.mmseq2.log") 
    conda: "renv.yaml"
    shell:
        """
        (date &&\
        Rscript mmseq_plots.R {input.plot1} {input.plot2} &&\
        date) &> >(tee {log})
        """

# rule mmseq_compare:
#     input:
#         f"{RESULTS_DIR}/annotation/mmseq2/{{assembler1}}_db",
#         f"{RESULTS_DIR}/annotation/mmseq2/{{assembler2}}_db"
#     output: 
#         f"{RESULTS_DIR}/annotation/mmseq2/{{assembler1}}__{{assembler2}}_rbh"
#     conda:
#         "cd-hit.yaml"
#     shell:
#         "mmseqs rbh {input[0]} {input[1]} {output} --min-seq-id 0.9 mmseq2_tmp --threads 12"

###########
## metaT ##
###########
# Preprocess the metaT reads using fastp
rule run_fastp_on_metaT_reads:
    input:
        r1=os.path.join(config["metaT_prefix"], "{metaT_sample}_R1_001.fastq.gz"),
        r2=os.path.join(config["metaT_prefix"], "{metaT_sample}_R2_001.fastq.gz")
    output:
        r1=os.path.join(RESULTS_DIR, "preprocessing/metaT/{metaT_sample}_R1_001.fastp.fq.gz"),
        r2=os.path.join(RESULTS_DIR, "preprocessing/metaT/{metaT_sample}_R2_001.fastp.fq.gz"),
        html=os.path.join(RESULTS_DIR, "preprocessing/metaT/{metaT_sample}.fastp.html"),
        json=os.path.join(RESULTS_DIR, "preprocessing/metaT/{metaT_sample}.fastp.json")
    log: os.path.join(RESULTS_DIR, "preprocessing/metaT/{metaT_sample}.fastp")
    conda: "envs/fastp.yaml"
    shell:
        """
        (date &&\
        fastp -l {config[fastp][min_length]} -i {input.r1} -I {input.r2} -o {output.r1} -O {output.r2} -h {output.html} -j {output.json} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule run_fastqc_on_preprocessed_metaT_reads:
    input: os.path.join(RESULTS_DIR, "preprocessing/metaT/{metaT_sample}_{suffix}.fq.gz"),
    output:
        html=os.path.join(RESULTS_DIR, "qc/metaT/fastqc/{metaT_sample}/{metaT_sample}_{suffix}.fastqc.html"),
        zip=os.path.join(RESULTS_DIR, "qc/metaT/fastqc/{metaT_sample}/{metaT_sample}_{suffix}.fastqc.zip") # the suffix _fastqc.zip is necessary for multiqc to find the file. If not using multiqc, you are free to choose an arbitrary filename
    conda: "envs/fastqc.yaml"
    script: "scripts/run_fastqc.py" # Use a separate script rather than a simple call to FASTQC to overcome possible race conditions when running in parallel (https://bitbucket.org/snakemake/snakemake-wrappers/raw/7f5e01706728e32d52aa413f213e950f50bf4129/bio/fastqc/wrapper.py)

# rule map_metaT_reads_to_long_reads:
#     input:
#         query_r1=rules.run_fastp_on_metaT_reads.output.r1,
#         query_r2=rules.run_fastp_on_metaT_reads.output.r2,
#         index=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}.amb"),
#         index_fa=os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fna"),
#     output: os.path.join(RESULTS_DIR, "mapping/metaT/lr/{metaT_sample}-x-{barcode}.bam")
#     threads: config["bwa"]["threads"]
#     params:
#         index_prefix=os.path.join(RESULTS_DIR, "mapping/lr/merged/{barcode}"),
#         out_prefix=os.path.join(RESULTS_DIR, "mapping/metaT/lr/{metaT_sample}-x-{barcode}.")
#     log: os.path.join(RESULTS_DIR, "mapping/metaT/lr/{metaT_sample}-x-{barcode}.bwa_mem_samtools")
#     conda: "envs/bwa.yaml"
#     shell:
#         """
#         (date &&\
#         bwa mem {config[bwa][long_reads_index][opts]} -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -S -b -u -T {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
#         date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
#         """

rule map_metaT_reads_to_long_read_contigs_w_bwa_mem:
    input:
        query_r1=rules.run_fastp_on_metaT_reads.output.r1,
        query_r2=rules.run_fastp_on_metaT_reads.output.r2,
        index_amb=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.amb"),
        index_ann=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.ann"),
        index_bwt=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.bwt"),
        index_pac=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.pac"),
        index_sa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.sa"),
        index_fa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly.fna")
    output: os.path.join(RESULTS_DIR, "mapping/metaT/lr/{metaT_sample}_reads-x-{barcode}-{assembler}_contigs.bam")
    params:
        index_prefix=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr/merged/{barcode}/assembly"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/metaT/lr/{metaT_sample}_reads-x_{barcode}_contigs")
    conda: "envs/bwa.yaml"
    threads: config["bwa"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/metaT/lr/{metaT_sample}_reads-x-{barcode}-{assembler}_contigs.bwa_mem_samtools")
    shell:
        """
        (date &&\
        bwa mem -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule map_metaT_reads_to_short_read_contigs_w_bwa_mem:
    input:
        query_r1=rules.run_fastp_on_metaT_reads.output.r1,
        query_r2=rules.run_fastp_on_metaT_reads.output.r2,
        index_amb=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.amb"),
        index_ann=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.ann"),
        index_bwt=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.bwt"),
        index_pac=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.pac"),
        index_sa=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs.sa"),
        index_fa=rules.assemble_short_reads_w_megahit.output
    output: os.path.join(RESULTS_DIR, "mapping/metaT/sr/{metaT_sample}_reads-x-{sr_sample}-{assembler}_contigs.bam")
    params:
        index_prefix=os.path.join(RESULTS_DIR, "assembly/{assembler}/{sr_sample}/final.contigs"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/metaT/sr/{metaT_sample}_reads-x-{sr_sample}-{assembler}_contigs")
    conda: "envs/bwa.yaml"
    threads: config["bwa"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/metaT/sr/{metaT_sample}_reads-x-{sr_sample}-{assembler}_contigs.bwa_mem_samtools")
    shell:
        """
        (date &&\
        bwa mem -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

rule map_metaT_reads_to_hybrid_contigs_w_bwa_mem:
    input:
        query_r1=rules.run_fastp_on_metaT_reads.output.r1,
        query_r2=rules.run_fastp_on_metaT_reads.output.r2,
        index_amb=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.amb"),
        index_ann=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.ann"),
        index_bwt=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.bwt"),
        index_pac=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.pac"),
        index_sa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.sa"),
        index_fa=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs.fna")
    output: os.path.join(RESULTS_DIR, "mapping/metaT/sr/{metaT_sample}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.bam")
    params:
        index_prefix=os.path.join(RESULTS_DIR, "assembly/{assembler}/lr_{barcode}-sr_{sr_sample}/contigs"),
        out_prefix=os.path.join(RESULTS_DIR, "mapping/metaT/sr/{metaT_sample}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs")
    log: os.path.join(RESULTS_DIR, "mapping/metaT/sr/{metaT_sample}_reads-x-lr_{barcode}_sr_{sr_sample}-{assembler}_contigs.bwa_mem_samtools")
    conda: "envs/bwa.yaml"
    threads: config["bwa"]["threads"]
    shell:
        """
        (date &&\
        bwa mem -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
        date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
        """

#######
# End of Snakefile as of 2020-04-20
#######


##############################
# Binning and Hybrid mapping #
##############################
#######################
## Mapping to hybrid ##
#######################
rule build_hybrid_bwa_index:
    input: expand(os.path.join(RESULTS_DIR, "assembly/{hybrid_assembler}/lr_{barcode}-sr_{sr_sample}/contigs.fasta"), sr_sample="NEB2_MG_S17", barcode="barcode07", hybrid_assembler=HYBRID_ASSEMBLER)
    output:
#           done=touch("bwa_index.done")
           index_amb=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{prefix}.amb"),
           index_ann=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{prefix}.ann"),
           index_bwt=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{prefix}.bwt"),
           index_pac=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{prefix}.pac"),
           index_sa=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{prefix}.sa")
    params:
           prefix=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{prefix}")
    log: os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{prefix}.bwa_index")
    conda: "envs/bwa.yaml"
    shell: """
           (date &&\
           bwa index {input} -p {params.prefix} &&\
           date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
           """

rule hybrid_map_short_reads_to_hybrid_contigs_w_bwa_mem:
    input:
          query_r1=expand(os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}_R1_001.fastp.fq.gz"), sr_sample="NEB2_MG_S17"),
          query_r2=expand(os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}_R2_001.fastp.fq.gz"), sr_sample="NEB2_MG_S17"),
          index_amb=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.amb"),
          index_ann=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.ann"),
          index_bwt=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.bwt"),
          index_pac=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.pac"),
          index_sa=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.sa"),
          index_fa=expand(os.path.join(RESULTS_DIR, "assembly/{hybrid_assembler}/lr_{barcode}-sr_{sr_sample}/contigs.fasta"), sr_sample="NEB2_MG_S17", barcode="barcode07", hybrid_assembler=HYBRID_ASSEMBLER)
    output:
          os.path.join(RESULTS_DIR, "mapping/bwa_sr_{hybrid_assembler}/bwa_sr_{hybrid_assembler}.bam")
    params:
          index_prefix=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades"),
          out_prefix=os.path.join(RESULTS_DIR, "mapping/bwa_sr_{hybrid_assembler}/bwa_sr_{hybrid_assembler}")
    log: os.path.join(RESULTS_DIR, "mapping/bwa_sr_{hybrid_assembler}/sr_contigs.bwa_mem_samtools")
    conda: "envs/bwa.yaml"
    threads:
          config["bwa"]["threads"]
#    wildcard_constraints:
#          sample='\w+'
    shell: """
           (date &&\
           bwa mem -t {threads} {params.index_prefix} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
           date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
           """

rule hybrid_map_long_reads_to_hybrid_contigs_w_bwa_mem:
    input:
          query_lr=expand(os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"), barcode=BARCODES),
          index_amb=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.amb"),
          index_ann=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.ann"),
          index_bwt=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.bwt"),
          index_pac=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.pac"),
          index_sa=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades.sa"),
          index_fa=expand(os.path.join(RESULTS_DIR, "assembly/{hybrid_assembler}/lr_{barcode}-sr_{sr_sample}/contigs.fasta"), sr_sample="NEB2_MG_S17", barcode="barcode07", hybrid_assembler=HYBRID_ASSEMBLER)
    output:
          os.path.join(RESULTS_DIR, "mapping/bwa_lr_{hybrid_assembler}/bwa_lr_{hybrid_assembler}.bam")
    params:
          index_prefix=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/metaspades"),
          out_prefix=os.path.join(RESULTS_DIR, "mapping/bwa_lr_{hybrid_assembler}/bwa_lr_{hybrid_assembler}"),
    log:  os.path.join(RESULTS_DIR, "mapping/bwa_lr_{hybrid_assembler}/lr_contigs.bwa_mem_samtools")
    conda: "envs/bwa.yaml"
    threads: config["bwa"]["threads"]
#    wildcard_constraints:
#          sample='\w+'
    shell: """
          (date &&\
          bwa mem -x ont2d -t {threads} {params.index_prefix} {input.query_lr} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.out_prefix} > {output} &&\
          date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
          """

rule hybrid_build_minimap2_long_read_contigs_ont_index:
    input: expand(os.path.join(RESULTS_DIR, "assembly/{hybrid_assembler}/lr_{barcode}-sr_{sr_sample}/contigs.fasta"), sr_sample="NEB2_MG_S17", barcode="barcode07", hybrid_assembler=HYBRID_ASSEMBLER)
    output: os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{hybrid_assembler}.mmi")
    conda: "envs/minimap2.yaml"
    shell:
           """
            date
            minimap2 -x map-ont -d {output} {input}
            date
            """

rule hybrid_map_long_reads_to_hybrid_contigs_w_minimap2:
    input:
          query=expand(os.path.join(RESULTS_DIR, "basecalled/merged/{barcode}.fastq.gz"), barcode="barcode07"),
          index=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{hybrid_assembler}.mmi"),
          index_fa=expand(os.path.join(RESULTS_DIR, "assembly/{hybrid_assembler}/lr_{barcode}-sr_{sr_sample}/contigs.fasta"), sr_sample="NEB2_MG_S17", barcode="barcode07", hybrid_assembler=HYBRID_ASSEMBLER)
    output: os.path.join(RESULTS_DIR, "mapping/mmi_lr_{hybrid_assembler}/mmi_lr_{hybrid_assembler}.bam")
    params:
          prefix=os.path.join(RESULTS_DIR, "mapping/mmi_lr_{hybrid_assembler}/mmi_lr_{hybrid_assembler}")
    conda: "envs/minimap2.yaml"
    threads: config["minimap2"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/mmi_lr_{hybrid_assembler}/metaspades.minimap2_samtools")
    shell: """
            (date &&\
            minimap2 -a -t {threads} {input.index} {input.query} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.prefix} > {output} &&\
            date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
            """


rule hybrid_map_short_reads_to_hybrid_contigs_w_minimap2:
    input:
          query_r1=expand(os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}_R1_001.fastp.fq.gz"), sr_sample="NEB2_MG_S17"),
          query_r2=expand(os.path.join(RESULTS_DIR, "preprocessing/sr/{sr_sample}_R2_001.fastp.fq.gz"), sr_sample="NEB2_MG_S17"),
          index=os.path.join(RESULTS_DIR, "mapping/{hybrid_assembler}/{hybrid_assembler}.mmi"),
          index_fa=expand(os.path.join(RESULTS_DIR, "assembly/{hybrid_assembler}/lr_{barcode}-sr_{sr_sample}/contigs.fasta"), sr_sample="NEB2_MG_S17", barcode="barcode07", hybrid_assembler=HYBRID_ASSEMBLER)
    output: os.path.join(RESULTS_DIR, "mapping/mmi_sr_{hybrid_assembler}/mmi_sr_{hybrid_assembler}.bam")
    params:
          prefix=os.path.join(RESULTS_DIR, "mapping/mmi_sr_{hybrid_assembler}/mmi_sr_{hybrid_assembler}")
    conda: "envs/minimap2.yaml"
    threads: config["minimap2"]["threads"]
    log: os.path.join(RESULTS_DIR, "mapping/mmi_sr_{hybrid_assembler}/metaspades.minimap2_samtools")
    shell:
           """
            (date &&\
           minimap2 -a -t {threads} {input.index} {input.query_r1} {input.query_r2} | samtools view -@ {config[samtools][view][threads]} -SbT {input.index_fa} | samtools sort -@ {config[samtools][sort][threads]} -m {config[samtools][sort][chunk_size]} -T {params.prefix} > {output}
           date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
            """
 
rule hybrid_merge_bam_files_for_metaspades:
    input:
          bam_sr=os.path.join(RESULTS_DIR, "mapping/{mapper}_sr_{hybrid_assembler}/{mapper}_sr_{hybrid_assembler}.bam"),
          bam_lr=os.path.join(RESULTS_DIR, "mapping/{mapper}_lr_{hybrid_assembler}/{mapper}_lr_{hybrid_assembler}.bam")
    output: os.path.join(RESULTS_DIR, "mapping/{mapper}_merged_{hybrid_assembler}/{mapper}_merged_{hybrid_assembler}.bam")
    conda: "/home/users/sbusi/apps/environments/binning.yml"
    shell:
          """
          (date &&\
          samtools merge {output} {input.bam_sr} {input.bam_lr} &&\
          date)
          """


###########
# Binning #
###########
# Need to edit the first rule and find a way to incorporate wildcards to make it less UGLY. Also, for some reason, symlinks did not seem to work, so switching to hardlinks

rule prepare_assembly_files:
    input:
         ass1=os.path.join(RESULTS_DIR, "assembly/flye/lr/merged/barcode07/assembly.fasta"),
         ass2=os.path.join(RESULTS_DIR, "assembly/megahit/NEB2_MG_S17/final.contigs.fa"),
         ass3=os.path.join(RESULTS_DIR, "assembly/metaspades_hybrid/lr_barcode07-sr_NEB2_MG_S17/contigs.fasta"),
         bam1=os.path.join(RESULTS_DIR, "mapping/sr/bwa_mem/NEB2_MG_S17_reads-x-NEB2_MG_S17-megahit_contigs.bam"),
         bam2=os.path.join(RESULTS_DIR, "mapping/lr/merged/barcode07/barcode07_reads-x-barcode07-flye_contigs.bam")
    output:
         fa1=os.path.join(RESULTS_DIR, "assembly/flye.fa"), 
         fa2=os.path.join(RESULTS_DIR, "assembly/megahit.fa"), 
         fa3=os.path.join(RESULTS_DIR, "assembly/metaspades_hybrid.fa"),
#         fa4=expand(os.path.join(RESULTS_DIR, "assembly/{mapper}_{reads}_{hybrid_assembler}.fa"), mapper=["bwa", "mmi"], reads=["sr", "lr", "merged"], hybrid_assembler=HYBRID_ASSEMBLER),
         fa4=os.path.join(RESULTS_DIR, "assembly/bwa_sr_metaspades_hybrid.fa"),
         fa5=os.path.join(RESULTS_DIR, "assembly/bwa_lr_metaspades_hybrid.fa"),
         fa6=os.path.join(RESULTS_DIR, "assembly/bwa_merged_metaspades_hybrid.fa"),
         fa7=os.path.join(RESULTS_DIR, "assembly/mmi_sr_metaspades_hybrid.fa"),
         fa8=os.path.join(RESULTS_DIR, "assembly/mmi_lr_metaspades_hybrid.fa"),
         fa9=os.path.join(RESULTS_DIR, "assembly/mmi_merged_metaspades_hybrid.fa"),
         bout1=os.path.join(RESULTS_DIR, "mapping/megahit/megahit.bam"),
         bout2=os.path.join(RESULTS_DIR, "mapping/flye/flye.bam")
    shell:
#         """
#         (date &&\
#         ln -s {input.ass1} {output.fa1} &&\
#         ln -s {input.ass2} {output.fa2} &&\
#         ln -s {input.ass3} {output.fa3} &&\
#         ln -s {input.ass3} {output.fa4} &&\
#         ln -s {input.ass3} {output.fa5} &&\
#         ln -s {input.ass3} {output.fa6} &&\
#         ln -s {input.ass3} {output.fa7} &&\
#         ln -s {input.ass3} {output.fa8} &&\
#         ln -s {input.ass3} {output.fa9} &&\
#         ln -s {input.bam1} {output.bout1} &&\
#         ln -s {input.bam2} {output.bout2} &&\
#         date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
#         """
         """
         (date &&\
         ln {input.ass1} {output.fa1}
         ln {input.ass2} {output.fa2}
         ln {input.ass3} {output.fa3}
         ln {input.ass3} {output.fa4}
         ln {input.ass3} {output.fa5}
         ln {input.ass3} {output.fa6}
         ln {input.ass3} {output.fa7}
         ln {input.ass3} {output.fa8}
         ln {input.ass3} {output.fa9}
         ln {input.bam1} {output.bout1}
         ln {input.bam2} {output.bout2} 
         date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
         """
   
#####################################
########### MAXBIN2 #################
#####################################
rule maxbin2:
    input:
          bbref=os.path.join(RESULTS_DIR, "assembly/{sample}.fa"),
          bbam=os.path.join(RESULTS_DIR, "mapping/{sample}/{sample}.bam")
    threads:config["threads"]
    conda:"/home/users/sbusi/apps/environments/binning.yml"
    output:
           map=os.path.join(RESULTS_DIR, "Binning/{sample}/{sample}.sam"),
           coverage=os.path.join(RESULTS_DIR, "Binning/{sample}/{sample}_cov.txt"),
           abund=os.path.join(RESULTS_DIR, "Binning/{sample}/{sample}_abundance.txt"),
           file=os.path.join(RESULTS_DIR, "Binning/{sample}/maxbin_output/maxbin_output.001.fasta")
#        mx=directory(os.path.join(RESULTS_DIR, "Binning/{sample}/maxbin_output/"))
    shell:
           """
            (date &&\
            samtools view -h -o {output.map} {input.bbam} &&\
            pileup.sh in={output.map} out={output.coverage} &&\
            awk '{{print $1"\\t"$2}}' {output.coverage} | grep -v '^#' > {output.abund} &&\
            run_MaxBin.pl -thread {threads} -contig {input.bbref} -out results/Binning/{wildcards.sample}/maxbin_output/maxbin_output -abund {output.abund} &&\
            touch maxbin.done &&\
            date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
            """


#####################################
########### METABAT2 ################
#####################################

rule depth_files:
    input:
          dep1=os.path.join(RESULTS_DIR, "mapping/{sample}/{sample}.bam")
    threads:config["threads"]
    conda:"/home/users/sbusi/apps/environments/binning.yml"
    output:
          depout=os.path.join(RESULTS_DIR, "Binning/{sample}/{sample}_depth.txt")
    shell:
           """
            (date &&\
            jgi_summarize_bam_contig_depths --outputDepth {output.depout} {input.dep1} &&\
            touch depth_file.done &&\
            date) 2> >(tee {log}.stderr) > >(tee {log}.stdout) 
            """

rule metabat2:
    input:
           met1=os.path.join(RESULTS_DIR, "assembly/{sample}.fa"),
           met2=os.path.join(RESULTS_DIR, "Binning/{sample}/{sample}_depth.txt")
    threads:config["threads"]
    conda:"/home/users/sbusi/apps/environments/binning.yml"
    output:
           outdir=directory(os.path.join(RESULTS_DIR, "Binning/{sample}/metabat_output/"))
    shell:
           """
            (date &&\
            metabat2 -i {input.met1} -a {input.met2} -o results/Binning/{wildcards.sample}/metabat_output/{wildcards.sample}.metabat -t {threads} -m 1500 -v --unbinned &&\
            touch metabat2.done &&\
            date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
            """

#####################################
########### DAS_Tool ################
#####################################
rule scaffold_list:
    input:
        sca1=os.path.join(RESULTS_DIR, "Binning/{sample}/metabat_output/"),
#       sca2=os.path.join(RESULTS_DIR, "Binning/{sample}/maxbin_output/")
        file=os.path.join(RESULTS_DIR, "Binning/{sample}/maxbin_output/maxbin_output.001.fasta")
    threads:config["threads"]
#    conda:"/home/users/sbusi/apps/environments/base.yml"
    conda: "dastool.yaml"
    output:
            scout1=os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_metabat.scaffolds2bin.tsv"),
            scout2=os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_maxbin.scaffolds2bin.tsv")
    shell:
            """
            (date &&\
            export Rscript={config[Rscript]} &&\
            Fasta_to_Scaffolds2Bin.sh -i {input.sca1} -e fa > {output.scout1} &&\
            Fasta_to_Scaffolds2Bin.sh -i results/Binning/{wildcards.sample}/maxbin_output -e fasta > {output.scout2} &&\
            touch scaffold_list.done &&\
            date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
            """

rule DAS_Tool:
    input:
            da1=os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_metabat.scaffolds2bin.tsv"),
            da2=os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_maxbin.scaffolds2bin.tsv"),
            da3=os.path.join(RESULTS_DIR, "assembly/{sample}.fa"),
            db=config["dastool_database"]
    threads:config["threads"]
#    conda:"/home/users/sbusi/apps/environments/base.yml"
    conda: "dastool.yaml"
    params:
            basename=os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}")
    output:
            daout=directory(os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_DASTool_bins")),
            dafile=os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_proteins.faa"),
            damfile=touch(os.path.join(RESULTS_DIR, "Binning/{sample}_das_tool.done"))
    shell:
            """
            (date &&\
            export Rscript={config[Rscript]} &&\
            DAS_Tool -i {input.da1},{input.da2} -c {input.da3} -o {params.basename} --search_engine diamond -l maxbin2,metabat2 --write_bins 1 --write_bin_evals 1 --threads {threads} --db_directory {input.db} --create_plots 1 &&\
            touch {output.damfile} &&\
            date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
            """

#####################################
######## GTDBTK & CHECKM ############
#####################################

rule gtdbtk:
    input:
           gen=os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_DASTool_bins")
    threads:config["threads"]
    conda:"gtdbtk.yaml"
    params:
           outdir=os.path.join(RESULTS_DIR, "Binning/gtdbtk_output/{sample}")
    output:
#           gtdb=directory(os.path.join(RESULTS_DIR, "Binning/gtdbtk_output/{sample}")),
           gtdbfile=os.path.join(RESULTS_DIR, "Binning/gtdbtk_output/{sample}/gtdbtk.bac120.summary.tsv")
#    wildcard_constraints:
 #          sample='\w+'
    shell:
            """
            (date &&\
            export GTDBTK_DATA_PATH={config[GTDBTK][DATA]} &&\
            gtdbtk classify_wf --cpus {threads} -x fa --genome_dir {input.gen} --out_dir {params.outdir} &&\
            touch gtdbtk.done &&\
            date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
            """

rule checkm:
    input:
           chkm=os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_DASTool_bins")
    threads:config["threads"]
    conda:"checkm.yml"
    output:
           chkmdir=directory(os.path.join(RESULTS_DIR, "Binning/checkm_output/{sample}/lineage.ms")),
           chkout=os.path.join(RESULTS_DIR, "Binning/checkm_output/{sample}_output.txt")
    shell:
            """
            (date &&\
            checkm lineage_wf -r -t {threads} -x fa {input.chkm} {output.chkmdir}  | tee {output.chkout} &&\
            touch checkm.done &&\
            date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
            """

rule plot_checkm:
    input:
           pl1=os.path.join(RESULTS_DIR, "Binning/checkm_output/{sample}"),
           pl2=os.path.join(RESULTS_DIR, "Binning/{sample}/dastool_output/{sample}_DASTool_bins")
    threads:config["threads"]
    conda:"/home/users/sbusi/apps/environments/checkm.yml"
    output:
           plt=directory(os.path.join(RESULTS_DIR, "Binning/checkm_output/{sample}/plots"))
    shell:
         """   
         date
         checkm bin_qa_plot --image_type png --width 8 --dpi 200 --font_size 8 -x fa -q {input.pl1} {input.pl2} {output.plt}
         touch plot_checkm.done 
         date) 2> >(tee {log}.stderr) > >(tee {log}.stdout)
         """


#######
# End of Snakefile; 2020-04-22
#######