# Python standard library
from os.path import join
from os import listdir
import os, sys, json

# 3rd party imports from pypi
from snakemake.workflow import workflow as wf_api

# Local imports
from scripts.common import (
    provided, 
    references,
    allocated,
    str_bool
)

# Global workflow variables
configfile: 'config.json' 
samples  = config['samples']              # Base name of input samples
workpath = config['project']['workpath']  # Pipeline's output directory
filetype = config['project']['filetype']  # 'paired-end' or 'single-end' (not supported)
tmpdir = config['options']['tmp_dir']     # Temporary directory
nends = config['project']['nends']        # paired-end of a sample: k=SE, values=PE

# Separate Single and Paired-end Samples
se_samples = list(config['project']['nends'].keys())
pe_samples = [s for s in config['project']['nends'].values() if s]

# Analysis Options 
do_aggregate = str_bool(                  # Aggregate results into single report,
    config['options']['aggregate']        # default: False
)

# Read in resource information,
# containing information about 
# threads, mem, walltimes, etc.
# TODO: Add handler for when the
# mode is set to local.
with open(join('config', 'cluster.json')) as fh:
    cluster = json.load(fh)

# Imported rules
include: join("rules", "common.smk")
include: join("rules", "paired-end.smk")

# Targets of the pipeline to build the DAG,
# assumes the user provided paired-end data.
rule all:
    input:
        # FastQValidator, only runs if Snakemake is 
        # invoked with the --use-singularity option
        # @imported from rules/paired-end.smk
        expand(
            join(workpath,"{name}","rawQC","{name}.validated.R1.fastq.log"), 
            name=provided(samples, wf_api.use_singularity)
        ),
        # Only generated if sample is paired-end 
        expand(
            join(workpath,"{name}","rawQC","{name}.validated.R2.fastq.log"), 
            name=provided(pe_samples, wf_api.use_singularity)
        ),
        # Extract Flowcell Lane information
        # @imported from rules/common.smk
        expand(
            join(workpath,"{name}","rawQC","{name}.fastq.info.txt"), 
            name=samples
        ),
        # FastQC (before and after trimming)
        # @imported from rules/paired-end.smk
        expand(
            join(workpath,"{name}","rawQC","{name}.R1_fastqc.zip"), 
            name=samples
        ),
        expand(
            join(workpath,"{name}","QC","{name}.R1.trim_fastqc.zip"), 
            name=samples
        ),
        # Cutadapt, remove adapter sequences
        # @imported from rules/paired-end.smk
        expand(
            join(workpath,"{name}","trim","{name}.R1.trim.fastq"), 
            name=samples
        ),
        # Remove Host contamination
        # @imported from rules/paired-end.smk
        expand(
            join(workpath,"{name}","trim","{name}.R1.trim.host_removed.fastq.gz"), 
            name=samples
        ),
        # Kraken2, viral taxonomic classification
        # @imported from rules/paired-end.smk
        expand(
            join(workpath,"{name}","kraken2","{name}.reads.krona.html"), 
            name=samples
        ),
        # Metaspades, viral assembly, annotation, and alignment to contigs,
        # conditional runs with pair-end samples,
        # metaspades does NOT support single-end,
        # https://github.com/ablab/spades/discussions/1026
        # @imported from rules/paired-end.smk
        expand(
            join(workpath,"{name}","output","{name}.metaspades.bam"), 
            name=pe_samples
        ),
        # Megahit, viral assembly, annotation, and alignment to contigs
        # @imported from rules/paired-end.smk
        expand(
            join(workpath,"{name}","output","{name}.megahit.bam"), 
            name=samples
        ),
        # Krona, taxonomic classification report
        # @imported from rules/paired-end.smk
        expand(
            join(workpath,"{name}","output","{name}.contig.classification.html"),
            name=samples
        ),
        # Metaquast, quality report of assembly
        expand(
            join(workpath,"{name}","temp","{name}.metaquast.fa"),
            name=samples
        ),
        expand(
            join(workpath,"{name}","output","{name}_metaquast","report.html"),
            name=samples
        ),
        # Krona, project-level multi-sample report
        # @imported from rules/paired-end.smk
        expand(
            join(workpath,"Project","temp","{name}.megahit.contigs.CAT.krona.meta"),
            name=provided(samples, do_aggregate),
        ),
        provided(
            [join(workpath,"Project","Project.contig.classification.html")],
            do_aggregate
        )