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 ########################################
################################################################################

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

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

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

# BAM filtering
include: os.path.join(maindir, "shared", "rules", "bam_filtering.snakefile")

#umi_tools extract and dedup
include: os.path.join(maindir, "shared", "rules", "umi_tools.snakefile")
#Sambamba Markdup
include: os.path.join(maindir, "shared", "rules", "sambamba.snakefile")

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

# deepTools QC
include: os.path.join(maindir, "shared", "rules", "deepTools_qc.snakefile")

# Qualimap BAM QC
include: os.path.join(maindir, "shared", "rules", "Qualimap_bamqc.snakefile")

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

## Allele-specific JOBs
if "allelic-mapping" in mode:
    genome_alias = os.path.splitext(os.path.basename(genome))[0]
    # Updated global vars if mode = "allelic-mapping"
    if allele_mode == 'create_and_map':
        bowtie2_index_allelic = 'snp_genome/bowtie2_Nmasked/Genome.1.bt2'
        if len(strains) == 1:
            allele_hybrid = 'single'
            SNPFile = "snp_genome/all_SNPs_" + strains[0] + "_" + genome_alias + ".txt.gz"
        elif len(strains) == 2:
            allele_hybrid = 'dual'
            SNPFile = "snp_genome/all_" + strains[1] + "_SNPs_" + strains[0] + "_reference.based_on_" + genome_alias + ".txt"

        include: os.path.join(maindir, "shared", "rules", "masked_genomeIndex.snakefile")
    elif allele_mode == 'map_only':
        bowtie2_index_allelic = NMaskedIndex
        SNPFile = SNPfile
    ## mapping rules
    include: os.path.join(maindir, "shared", "rules", "Bowtie2_allelic.snakefile")
    ## SNPsplit
    include: os.path.join(maindir, "shared", "rules", "SNPsplit.snakefile")
    # deepTools QC
    include: os.path.join(maindir, "shared", "rules", "deepTools_qc_allelic.snakefile")
else:
    if aligner == "Bowtie2":
        # Bowtie2 mapping, duplicate marking, BAM filtering and indexing
        include: os.path.join(maindir, "shared", "rules", "Bowtie2.snakefile")
    elif aligner == "bwa":
        include: os.path.join(maindir, "shared", "rules", "bwa.snakefile")
    elif aligner == "bwa-mem2":
        include: os.path.join(maindir, "shared", "rules", "bwa-mem2.snakefile")

### conditional/optional rules #################################################
################################################################################

def run_FastQC(fastqc):
    if fastqc:
        return( expand("FastQC/{sample}{read}_fastqc.html", sample = samples, read = reads[:idxRange]) )
    else:
        return([])

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

def run_computeGCBias(GCBias):
    file_list = []
    if GCBias:
        if not downsample:
            file_list = expand("deepTools_qc/computeGCBias/{sample}.filtered.GCBias.png", sample = samples)
    return(file_list)

def run_deepTools_qc():
    file_list = ["deepTools_qc/bamPEFragmentSize/fragmentSize.metric.tsv"]
    if len(samples) <= 20:
        file_list.append( ["deepTools_qc/plotCoverage/read_coverage.tsv"] )
    if len(samples)>1 and len(samples)<=20:
        file_list.append( [
            "deepTools_qc/plotCorrelation/correlation.pearson.read_coverage.tsv",
            "deepTools_qc/plotCorrelation/correlation.spearman.read_coverage.tsv",
            "deepTools_qc/plotPCA/PCA.read_coverage.tsv" ])
        if 'allelic-mapping' in mode:
            file_list.append( [
                "deepTools_qc/plotCorrelation/correlation.pearson.read_coverage_allelic.tsv",
                "deepTools_qc/plotCorrelation/correlation.spearman.read_coverage_allelic.tsv",
                "deepTools_qc/plotPCA/PCA.read_coverage_allelic.tsv" ] )
    file_list.append(expand("deepTools_qc/estimateReadFiltering/{sample}_filtering_estimation.txt",sample = samples))
    return (file_list)

def run_Qualimap():
    file_list = []
    if qualimap:
        file_list += expand("Qualimap_qc/{sample}.filtered.bamqc_report.html", sample = samples)
        file_list += expand("Qualimap_qc/{sample}.filtered.bamqc_results.txt", sample = samples)
    return (file_list)
# allele specific
def make_nmasked_genome():
    if allele_mode == 'create_and_map':
        genome1 = "snp_genome/" + strains[0] + '_SNP_filtering_report.txt'
        file_list = [
                genome1,
                SNPFile,
                bowtie2_index_allelic
                ]
        return(file_list)
    else:
        return([])

def run_allelesp_mapping():
    if "allelic-mapping" in mode:
        allele_suffix = ['allele_flagged', 'genome1', 'genome2', 'unassigned']
        file_list = [
        expand("allelic_bams/{sample}.{suffix}.sorted.bam", sample = samples,
                        suffix = allele_suffix),
        expand("allelic_bams/{sample}.{suffix}.sorted.bam.bai", sample = samples,
                        suffix = allele_suffix),
        expand("bamCoverage/allele_specific/{sample}.{suffix}.seq_depth_norm.bw", sample = samples,
                        suffix = ['genome1', 'genome2'])
        ]
        return(file_list)
    else:
        return([])

### execute before workflow starts #############################################
################################################################################
onstart:
    if "verbose" in config and config["verbose"]:
        print("--- Workflow parameters --------------------------------------------------------")
        print("samples:", samples)
        print("paired:", pairedEnd)
        print("read extension:", reads)
        print("fastq dir:", fastq_dir)
        print("maximum insert size (Bowtie2 -X):", insertSizeMax)
        print("-" * 80, "\n")

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

    if toolsVersion:
        usedEnvs = [CONDA_SHARED_ENV, CONDA_DNA_MAPPING_ENV]
        cf.writeTools(usedEnvs, outdir, "DNAmapping", maindir)


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

rule all:
    input:
        run_FastQC(fastqc),
        run_Trimming(trim, fastqc),
        expand(aligner + "/{sample}.markdup.bam", sample = samples),
        expand("Sambamba/{sample}.markdup.txt", sample = samples),
        expand("filtered_bam/{sample}.filtered.bam", sample = samples),
        expand("bamCoverage/{sample}.seq_depth_norm.bw", sample = samples),
        expand("bamCoverage/{sample}.filtered.seq_depth_norm.bw", sample = samples),
        run_computeGCBias(GCBias),
        run_deepTools_qc(),
        run_Qualimap(),
        make_nmasked_genome(),
        run_allelesp_mapping(),
        "multiQC/multiqc_report.html"

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

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