# 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
from snakemake.utils import R

import subprocess,shlex

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

# Global workflow variables
configfile: 'config.json'                        # RENEE config file
samples  = config['project']['groups']['rsamps'] # sample basenames
pfamily  = config['project']['pfamily']          # rnaseq
workpath = config['project']['workpath']         # Parent Output directory
tmpdir   = config['options']['tmp_dir']          # Temporary directory
version  = config['project']['version']         # Workflow version
# 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)

# Check for SE or PE FastQ files:
convert = {1: False, 2: True}                     # 1 = SE, 2 = PE, -1 = Unknown
try:
    paired_end = convert[config['project']['nends']]  # True if PE else false
except KeyError:
    # Catching case when value is -1 or unknown
    sys.exit("Fatal: Raw data could not be classified as single-end or paired-end data!")

# Child output directory aliases
trim_dir='trim'
star_dir='STAR_files'
bams_dir='bams'
log_dir='logfiles'
rseqc_dir='RSeQC'
kraken_dir='kraken'
preseq_dir='preseq'
degall_dir='DEG_ALL'
dtypes=['RSEM_genes']


def get_nidap_folder_input_files(wildcards):
    nidap_files=[]
    nidap_files.append(join(workpath,"Reports","RNA_Report.html"))
    nidap_files.append(join(workpath,"Reports","multiqc_report.html"))
    nidap_files.append(join(workpath,degall_dir,"RSEM.genes.FPKM.all_samples.txt"))
    nidap_files.append(join(workpath,degall_dir,"RSEM.isoforms.FPKM.all_samples.txt"))
    return nidap_files

def get_nidap_folder_ouptut_files(wildcards):
    nidap_files=[]
    if config['options']['create_nidap_folder'] == "True":
        nidap_files.append(join(workpath,"NIDAP","RNA_Report.html"))
        nidap_files.append(join(workpath,"NIDAP","multiqc_report.html"))
        nidap_files.append(join(workpath,"NIDAP","RSEM.genes.FPKM.all_samples.txt"))
        nidap_files.append(join(workpath,"NIDAP","RSEM.isoforms.FPKM.all_samples.txt"))
    return nidap_files


if paired_end:
    # Paired-end Target files
    rule all:
        params:
            batch='--time=168:00:00',
        input:
            # FastQValidator
            expand(join(workpath,"rawQC","{name}.validated.R1.fastq.log"), name=provided(samples, wf_api.use_singularity)),
            expand(join(workpath,"rawQC","{name}.validated.R2.fastq.log"), name=provided(samples, wf_api.use_singularity)),

            # Flowcell Lane information
            expand(join(workpath,"rawQC","{name}.fastq.info.txt"),name=samples),

            # FastQC (before and after trimming)
            expand(join(workpath,"rawQC","{name}.R1_fastqc.zip"), name=samples),
            expand(join(workpath,"rawQC","{name}.R2_fastqc.zip"), name=samples),
            expand(join(workpath,"QC","{name}.R1.trim_fastqc.zip"), name=samples),
            expand(join(workpath,"QC","{name}.R2.trim_fastqc.zip"), name=samples),

            # BBtools Insert Size
            expand(join(workpath,"QC","{name}_insert_sizes.txt"), name=samples),

            # FastScreen (Using two sets of reference databases)
            expand(join(workpath,"FQscreen","{name}.R1.trim_screen.txt"),name=samples),
            expand(join(workpath,"FQscreen","{name}.R1.trim_screen.png"),name=samples),
            expand(join(workpath,"FQscreen","{name}.R2.trim_screen.txt"),name=samples),
            expand(join(workpath,"FQscreen","{name}.R2.trim_screen.png"),name=samples),
            expand(join(workpath,"FQscreen2","{name}.R1_2.trim_screen.txt"),name=samples),
            expand(join(workpath,"FQscreen2","{name}.R1_2.trim_screen.png"),name=samples),
            expand(join(workpath,"FQscreen2","{name}.R2_2.trim_screen.txt"),name=samples),
            expand(join(workpath,"FQscreen2","{name}.R2_2.trim_screen.png"),name=samples),

            # Kraken + Krona
            expand(join(workpath,kraken_dir,"{name}.trim.kraken_bacteria.taxa.txt"),name=samples),
            expand(join(workpath,kraken_dir,"{name}.trim.kraken_bacteria.krona.html"),name=samples),

            # STAR
            expand(join(workpath,bams_dir,"{name}.p2.Aligned.toTranscriptome.out.bam"),name=samples),

            # Arriba
            expand(join(workpath,"fusions","{name}_fusions.tsv"),name=provided(samples, references(config, pfamily, ['FUSIONBLACKLIST', 'FUSIONCYTOBAND', 'FUSIONPROTDOMAIN']))),
            expand(join(workpath,"fusions","{name}_fusions.arriba.pdf"),name=provided(samples, references(config, pfamily, ['FUSIONBLACKLIST', 'FUSIONCYTOBAND', 'FUSIONPROTDOMAIN']))),

            # Bam to stranded bigwigs
            expand(join(workpath,bams_dir,"{name}.star_rg_added.sorted.dmark.bam"),name=samples),
            expand(join(workpath,bams_dir,"{name}.fwd.bw"),name=samples),
            expand(join(workpath,bams_dir,"{name}.rev.bw"),name=samples),

            # Picard
            expand(join(workpath,log_dir,"{name}.RnaSeqMetrics.txt"),name=samples),
            expand(join(workpath,log_dir,"{name}.star.duplic"),name=samples),

            # Preseq
            expand(join(workpath,preseq_dir,"{name}.ccurve"),name=samples),

            # QualiMap (bamQC and counts)
            expand(join(workpath,"QualiMap","{name}","qualimapReport.html"),name=samples),

            # RSeQC
            expand(join(workpath,rseqc_dir,"{name}.strand.info"),name=samples),
            expand(join(workpath,rseqc_dir,"{name}.Rdist.info"),name=samples),
            expand(join(workpath,rseqc_dir,"{name}.inner_distance_freq.txt"),name=samples),
            expand(join(workpath,rseqc_dir,"{name}.star_rg_added.sorted.dmark.summary.txt"),name=samples),
            join(workpath,degall_dir,"combined_TIN.tsv"),

            # RSEM merge and counts
            expand(join(workpath,degall_dir,"{name}.RSEM.genes.results"),name=samples),
            expand(join(workpath,degall_dir,"{name}.RSEM.isoforms.results"),name=samples),
            join(workpath,degall_dir,"RSEM.genes.FPKM.all_samples.txt"),
            join(workpath,degall_dir,"RSEM.isoforms.FPKM.all_samples.txt"),
            join(workpath, degall_dir, "RSEM.genes.expected_counts.all_samples.matrix"),
            join(workpath, degall_dir, "RSEM.isoforms.expected_counts.all_samples.matrix"),
            # MultiQC
            join(workpath,"Reports","multiqc_report.html"),

            # RNA QC Report
            join(workpath,"Reports","RNA_Report.html"),

            #NIDAP stuff
            get_nidap_folder_ouptut_files


elif not paired_end:
    # Single-end Target files
    rule all:
        params: batch='--time=168:00:00'
        input:
            # FastQValidator
            expand(join(workpath,"rawQC","{name}.validated.R1.fastq.log"), name=provided(samples, wf_api.use_singularity)),

            # Flowcell Lane information
            expand(join(workpath,"rawQC","{name}.fastq.info.txt"),name=samples),

            # FastQC (before and after trimming)
            expand(join(workpath,"rawQC","{name}.R1_fastqc.zip"), name=samples),
            expand(join(workpath,"QC","{name}.R1.trim_fastqc.zip"), name=samples),

            # FastQ Screen
            expand(join(workpath,"FQscreen","{name}.R1.trim_screen.txt"),name=samples),
            expand(join(workpath,"FQscreen","{name}.R1.trim_screen.png"),name=samples),
            expand(join(workpath,"FQscreen2","{name}.R1_2.trim_screen.txt"),name=samples),
            expand(join(workpath,"FQscreen2","{name}.R1_2.trim_screen.png"),name=samples),

            # Kraken + Krona
            expand(join(workpath,kraken_dir,"{name}.trim.kraken_bacteria.taxa.txt"),name=samples),
            expand(join(workpath,kraken_dir,"{name}.trim.kraken_bacteria.krona.html"),name=samples),

            # STAR
            expand(join(workpath,bams_dir,"{name}.p2.Aligned.toTranscriptome.out.bam"),name=samples),

            # Picard
            expand(join(workpath,log_dir,"{name}.RnaSeqMetrics.txt"),name=samples),

            # QualiMap
            expand(join(workpath,"QualiMap","{name}","qualimapReport.html"),name=samples),

            # Bam to stranded bigwigs
            expand(join(workpath,bams_dir,"{name}.fwd.bw"),name=samples),
            expand(join(workpath,bams_dir,"{name}.rev.bw"),name=samples),

            # RSeQC
            expand(join(workpath,rseqc_dir,"{name}.star_rg_added.sorted.dmark.summary.txt"),name=samples),
            join(workpath,degall_dir,"combined_TIN.tsv"),

            # RSEM
            expand(join(workpath,degall_dir,"{name}.RSEM.genes.results"),name=samples),
            expand(join(workpath,degall_dir,"{name}.RSEM.isoforms.results"),name=samples),
            join(workpath,degall_dir,"RSEM.genes.FPKM.all_samples.txt"),
            join(workpath,degall_dir,"RSEM.isoforms.FPKM.all_samples.txt"),
            join(workpath, degall_dir, "RSEM.genes.expected_counts.all_samples.matrix"),
            join(workpath, degall_dir, "RSEM.isoforms.expected_counts.all_samples.matrix"),

            # MultiQC
            join(workpath,"Reports","multiqc_report.html"),

            #NIDAP stuff
            get_nidap_folder_ouptut_files

# Rules common to single-end or paired-end data
include: "rules/common.smk"
# Include single-end and paired-specific rules
if paired_end:
    # rules common to paired-end data
    include: "rules/paired-end.smk"
elif not paired_end:
    # rules common to paired-end data
    include: "rules/single-end.smk"
include: "rules/group-info.smk"

if config['options']['create_nidap_folder']:
    include: "rules/nidap.smk"

on_complete = f"""
for cmd in spooker run_jobby_on_snakemake_log; do
    if ! command -v $cmd 2>&1 >/dev/null; then
        export PATH="$PATH:{config['bin']['rnaseq']['ccbr_tools_path']}"
    fi
done
jobby logfiles/snakemake.log --tsv | tee logfiles/snakemake.log.jobby.tsv | cut -f2,3,18 > logfiles/snakemake.log.jobby.short.tsv
spooker \
    --outdir {workpath} \
    --name RENEE \
    --version {version} \
    --path {config["project"]["pipelinehome"]} \
    > logfiles/spooker.log 2>&1
"""

onsuccess:
    print("OnSuccess")
    shell("printenv")
    shell("module list")
    print(on_complete)
    shell(on_complete)

onerror:
    print("OnError")
    shell("printenv")
    shell("module list")
    print(on_complete)
    shell(on_complete)
