import os
import snakePipes.common_functions as cf
import subprocess
import sys

### snakemake_workflows initialization ########################################
maindir = os.path.dirname(os.path.dirname(workflow.basedir))
workflow_rscripts=os.path.join(maindir, "shared", "rscripts")

# 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 ########################################
################################################################################
include: os.path.join(maindir, "shared", "rules", "FASTQ.snakefile")
#STARsolo uses filtered GTF for counting
include: os.path.join(maindir, "shared", "rules", "filterGTF.snakefile")
#umi_tools
if mode=="STARsolo":
    include: os.path.join(maindir, "shared", "rules", "scRNAseq_STARsolo.snakefile")
    include: os.path.join(maindir, "shared", "rules", "sambamba.snakefile")
    include: os.path.join(maindir, "shared", "tools" , "deeptools_cmds.snakefile")
    include: os.path.join(maindir, "shared", "rules", "deepTools_RNA.snakefile")
    trim=False
elif mode=="Alevin":
#    include: os.path.join(maindir, "shared", "rules", "Salmon.snakefile")
    include: os.path.join(maindir, "shared", "rules", "scRNAseq_Alevin.snakefile")
    trim=False
pairedEnd = True
include: os.path.join(maindir, "shared", "rules", "FastQC.snakefile")
pairedEnd = False
include: os.path.join(maindir, "shared", "rules", "multiQC.snakefile")


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


def run_deeptools_qc():
    if mode=="STARsolo":
        file_list = [
        expand("bamCoverage/{sample}.RPKM.bw", sample = samples),
        expand("bamCoverage/{sample}.coverage.bw", sample = samples),
        "deepTools_qc/plotEnrichment/plotEnrichment.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_velocyto(skipVelocyto):
    if not skipVelocyto :
        if mode == "STARsolo":
            return([])
        elif mode == "Alevin":
            file_list = [expand("AlevinForVelocity/{sample}/alevin/quants_mat.gz",sample=samples),
            "SingleCellExperiment/AlevinForVelocity/merged_samples.RDS"]
            return(file_list)

    else:
        return([])

### execute before workflow starts #############################################
################################################################################
onstart:
    if "verbose" in config and config["verbose"]:
        print("--- Workflow parameter ---------------------------------------------------------")
        print("Input directory:", indir)
        print("Input files:", infiles)
        print("Samples:", samples)
        print("Paired-end reads:", pairedEnd)
        print("Read extension:", reads)
        print("Genome:", genome)
        print("Downsample:", downsample)
        print("Mode:", mode)
        print("Input directory for mapping:", fastq_dir)
        print("BigWig bin size:", bwBinSize)

        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_scRNASEQ_ENV]
        cf.writeTools(usedEnvs, outdir, "scRNAseq", maindir)

### main rule ##################################################################
################################################################################
if mode=="STARsolo":
    localrules: gzip_STARsolo_for_seurat
    rule all:
        input:
            expand("STARsolo/{sample}.sorted.bam",sample = samples),
            "STARsolo/Report.tsv",
            run_deeptools_qc(),
            "deepTools_qc/bamPEFragmentSize/fragmentSize.metric.tsv",
            expand("deepTools_qc/estimateReadFiltering/{sample}_filtering_estimation.txt",sample = samples),
            "multiQC/multiqc_report.html",
            "Seurat/STARsolo_filtered/merged_samples.RDS",
            "Seurat/STARsolo_raw/merged_samples.RDS",
            "Seurat/STARsolo_raw_RmEmptyDrops/merged_samples.RDS",
            run_velocyto(skipVelocyto)
elif mode=="Alevin":
    localrules: filter_gtf, gtf_to_files, cut_t2g
    rule all:
        input:
            expand("originalFASTQ/{sample}{read}.fastq.gz", sample = samples, read=reads),
            expand("FastQC/{sample}{read}_fastqc.html", sample = samples, read=reads),
            "Annotation/genes.slim.t2g",
            expand("Alevin/{sample}/alevin/quants_mat.gz",sample = samples),
            expand("multiQC/Alevin_{sample}.html", sample=samples),
            "multiQC/multiqc_report.html",
            run_velocyto(skipVelocyto)


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

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