import os
import snakePipes.common_functions as cf


### snakemake_workflows initialization ########################################
maindir = os.path.dirname(os.path.dirname(workflow.basedir))

# load conda ENVs (path is relative to "shared/rules" directory)
globals().update(cf.set_env_yamls())

# load config file
globals().update(cf.load_configfile(workflow.overwrite_configfiles[0], config["verbose"]))
# load organism-specific data, i.e. genome indices, annotation, etc.
globals().update(cf.load_organism_data(genome, maindir, config["verbose"]))
# return the pipeline version in the log
cf.get_version()

# do workflow specific stuff now
include: os.path.join(workflow.basedir, "internals.snakefile")

### include modules of other snakefiles ########################################
################################################################################

# deeptools cmds
include: os.path.join(maindir, "shared", "tools" , "deeptools_cmds.snakefile")

## bamCoverage_RPKM
include: os.path.join(maindir, "shared", "rules", "deepTools_RNA.snakefile")


if not fromBAM:
    ## FASTQ: either downsample FASTQ files or create symlinks to input files
    include: os.path.join(maindir, "shared", "rules", "FASTQ.snakefile")
    include: os.path.join(maindir, "shared", "rules", "umi_tools.snakefile")

    ## FastQC
    if fastqc:
        include: os.path.join(maindir, "shared", "rules", "FastQC.snakefile")

    ## Trim
    if trim:
        include: os.path.join(maindir, "shared", "rules", "trimming.snakefile")

else:
    fastqc=False
    trim=False
    downsample=None
    include: os.path.join(maindir, "shared", "rules", "LinkBam.snakefile")

## Align and run TEcounts
include: os.path.join(maindir, "shared", "rules", "tecounts.snakefile")

## MultiQC
include: os.path.join(maindir, "shared", "rules", "multiQC.snakefile")

### conditional/optional rules #################################################
################################################################################
def run_FastQC(fastqc):
    readsUse = reads
    if not pairedEnd:
        readsUse = [reads[0]]
    if fastqc:
        return( expand("FastQC/{sample}{read}_fastqc.html", sample = samples, read = readsUse) )
    else:
        return([])


def run_Trimming(trim, fastqc):
    readsUse = reads
    if not pairedEnd:
        readsUse = [reads[0]]
    if trim and fastqc:
        return( expand(fastq_dir+"/{sample}{read}.fastq.gz", sample = samples, read = readsUse) +
                expand("FastQC_trimmed/{sample}{read}_fastqc.html", sample = samples, read = readsUse) )
    elif trim:
        return( expand(fastq_dir+"/{sample}{read}.fastq.gz", sample = samples, read = readsUse) )
    else:
        return([])


def run_deepTools_qc():
    if "deepTools_qc" in mode:
        file_list = [
        expand("bamCoverage/{sample}.RPKM.bw", sample = samples),
        expand("bamCoverage/{sample}.scaleFactors.bw",sample = samples),
        expand("bamCoverage/{sample}.coverage.bw", sample = samples),
        expand("bamCoverage/{sample}.uniqueMappings.fwd.bw", sample = samples),
        expand("bamCoverage/{sample}.uniqueMappings.rev.bw", sample = samples),
        "deepTools_qc/plotEnrichment/plotEnrichment.tsv",
        expand("deepTools_qc/estimateReadFiltering/{sample}_filtering_estimation.txt",sample = samples)]
        if pairedEnd:
            file_list.append("deepTools_qc/bamPEFragmentSize/fragmentSize.metric.tsv")
        if len(samples)>1:
            file_list.append( ["deepTools_qc/multiBigwigSummary/coverage.bed.npz",
                              "deepTools_qc/plotCorrelation/correlation.pearson.bed_coverage.tsv",
                              "deepTools_qc/plotCorrelation/correlation.spearman.bed_coverage.tsv",
                              "deepTools_qc/plotPCA/PCA.bed_coverage.tsv"] )
        return(file_list)
    else:
        return([])


def run_deseq2():
    if sampleSheet:
        sample_name = os.path.splitext(os.path.basename(sampleSheet))[0]
        if isMultipleComparison:
            file_list = expand("DESeq2_{}".format(sample_name) + ".{compGroup}/DESeq2.session_info.txt",compGroup=cf.returnComparisonGroups(sampleSheet))
        else:
            file_list = ["DESeq2_{}/DESeq2.session_info.txt".format(sample_name)]
        return(file_list)
    return([])


### execute before  starts #####################################################
################################################################################
onstart:
    if "verbose" in config and config["verbose"] and not fromBAM:
        print()
        print("--- Workflow parameters --------------------------------------------------------")
        print("mode:", mode)
        print("samples:", samples)
        print("paired:", pairedEnd)
        print("read extension:", reads)
        print("fastq dir:", fastq_dir)
        print("Sample sheet:", sampleSheet)
        print("-" * 80, "\n")

        print("--- Environment ----------------------------------------------------------------")
        print("$TMPDIR: ", os.getenv('TMPDIR', ""))
        print("$HOSTNAME: ", os.getenv('HOSTNAME', ""))
        print("-" * 80, "\n")

    elif "verbose" in config and config["verbose"] and fromBAM:
        print("--- Workflow parameters --------------------------------------------------------")
        print("samples:" + str(samples))
        print("input dir:" + indir)
        print("paired:", pairedEnd)
        print("-" * 80)  # , "\n"

        print("--- Environment ----------------------------------------------------------------")
        print("$TMPDIR: " + os.getenv('TMPDIR', ""))
        print("$HOSTNAME: " + os.getenv('HOSTNAME', ""))
        print("-" * 80,)

    if toolsVersion:
        usedEnvs = [CONDA_SHARED_ENV, CONDA_NONCODING_RNASEQ_ENV]
        cf.writeTools(usedEnvs, outdir, "ncRNAseq", maindir)
    if sampleSheet:
        cf.copySampleSheet(sampleSheet, outdir)

### main rule ##################################################################
################################################################################

if not fromBAM:
    
    rule all:
        input:
            expand("TEcount/{sample}.cntTable", sample = samples),
            run_FastQC(fastqc),
            run_Trimming(trim, fastqc),
            run_deepTools_qc(),
            run_deseq2(),
            "multiQC/multiqc_report.html"

else:
    rule all:
        input:
            expand("TEcount/{sample}.cntTable", sample = samples),
            run_deepTools_qc(),
            run_deseq2(),
            "multiQC/multiqc_report.html"

### execute after  finished ####################################################
################################################################################
onsuccess:
    if "verbose" in config and config["verbose"]:
        print("\n--- ncRNAseq workflow finished successfully! ------------------------------------\n")

onerror:
    print("\n !!! ERROR in ncRNAseq workflow! !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n")
