import os
import snakePipes.common_functions as cf

### 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"]))
# return the pipeline version in the log
cf.get_version()

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

# This will eventually be written as a yaml file
organismDictionary = {'genome_size': 0,
                      'genome_fasta': '',
                      'genome_index': '',
                      'genome_2bit': '',
                      'genome_dict': '',
                      'bowtie2_index': '',
                      'hisat2_index': '',
                      'bwa_index': '',
                      'bwa_mem2_index': '',
                      'bwameth_index': '',
                      'bwameth2_index': '',
                      'known_splicesites': '',
                      'star_index': '',
                      'salmon_index': '',
                      'salmon_velocity_index': '',
                      't2g_velocity': '',
                      'genes_bed': '',
                      'genes_gtf': '',
                      'genes_t2g': '',
                      'spikein_genes_gtf': '',
                      'extended_coding_regions_gtf': '',
                      'blacklist_bed': '',
                      'spikein_blacklist_bed': '',
                      'ignoreForNormalization': '',
                      'rmsk_file': ''}

targets = [os.path.join(outdir, x) for x in ["genome_fasta/genome.fa",
                                             "genome_fasta/genome.fa.fai",
                                             "genome_fasta/genome.dict",
                                             "genome_fasta/genome.2bit"]]

organismDictionary['genome_fasta'] = os.path.join(outdir, "genome_fasta/genome.fa")
organismDictionary['genome_index'] = os.path.join(outdir, "genome_fasta/genome.fa.fai")
organismDictionary['genome_2bit'] = os.path.join(outdir, "genome_fasta/genome.2bit")
organismDictionary['genome_dict'] = os.path.join(outdir, "genome_fasta/genome.dict")

if gtfURL:
    targets.extend([os.path.join(outdir, x) for x in ["annotation/genes.gtf",
                                                      "annotation/genes.bed",
                                                      "annotation/genes.t2g",
                                                      "annotation/genes.slop.gtf"]])
    organismDictionary['genes_gtf'] = os.path.join(outdir, "annotation/genes.gtf")
    organismDictionary['genes_bed'] = os.path.join(outdir, "annotation/genes.bed")
    organismDictionary['genes_t2g'] = os.path.join(outdir, "annotation/genes.t2g")
    organismDictionary['extended_coding_regions_gtf'] = os.path.join(outdir, "annotation/genes.slop.gtf")

if spikeinGtfURL:
    targets.extend([os.path.join(outdir, x) for x in ["annotation/spikein_genes.gtf"]])
    organismDictionary['spikein_genes_gtf'] = os.path.join(outdir, "annotation/spikein_genes.gtf")

if rmskURL:
    targets.extend([os.path.join(outdir, x) for x in ["annotation/rmsk.txt"]])
    organismDictionary['rmsk_file'] = os.path.join(outdir, "annotation/rmsk.txt")

if 'all' in tools or 'bowtie2' in tools:
    targets.append(os.path.join(outdir, "BowtieIndex/genome.rev.2.bt2"))
    organismDictionary['bowtie2_index'] = os.path.join(outdir, "BowtieIndex/genome")

if 'all' in tools or 'hisat2' in tools:
    targets.append(os.path.join(outdir, "HISAT2Index/genome.6.ht2"))
    organismDictionary['hisat2_index'] = os.path.join(outdir, "HISAT2Index/genome")
    if gtfURL:
        targets.append(os.path.join(outdir, "HISAT2Index/splice_sites.txt"))
        organismDictionary['known_splicesites'] = os.path.join(outdir, "HISAT2Index/splice_sites.txt")

if 'all' in tools or 'star' in tools:
    targets.append(os.path.join(outdir, "STARIndex/SAindex"))
    organismDictionary['star_index'] = os.path.join(outdir, "STARIndex/")

if 'all' in tools or 'salmon' in tools:
    targets.append(os.path.join(outdir, "SalmonIndex/seq.bin"))
    organismDictionary['salmon_index'] = os.path.join(outdir, "SalmonIndex/")
    targets.append(os.path.join(outdir, "SalmonIndex_RNAVelocity/seq.bin"))
    organismDictionary['salmon_velocity_index'] = os.path.join(outdir, "SalmonIndex_RNAVelocity/")
    targets.append(os.path.join(outdir, "annotation/cDNA_introns.joint.t2g"))
    organismDictionary['t2g_velocity'] = os.path.join(outdir, "annotation/cDNA_introns.joint.t2g")

if 'all' in tools or 'bwa' in tools:
    targets.append(os.path.join(outdir, "BWAIndex/genome.fa.sa"))
    organismDictionary['bwa_index'] = os.path.join(outdir, "BWAIndex/genome.fa")

if 'all' in tools or 'bwa-mem2' in tools:
    targets.append(os.path.join(outdir, "BWA-MEM2Index/genome.fa.bwt.2bit.64"))
    organismDictionary['bwa_mem2_index'] = os.path.join(outdir, "BWA-MEM2Index/genome.fa")

if 'all' in tools or 'bwameth' in tools:
    targets.append(os.path.join(outdir, "BWAmethIndex/genome.fa.bwameth.c2t.sa"))
    organismDictionary['bwameth_index'] = os.path.join(outdir, "BWAmethIndex/genome.fa")

if 'all' in tools or 'bwameth2' in tools:
    targets.append(os.path.join(outdir, "BWAmeth2Index/genome.fa.bwameth.c2t.bwt.2bit.64"))
    organismDictionary['bwameth2_index'] = os.path.join(outdir, "BWAmeth2Index/genome.fa")

if blacklist:
    targets.append(os.path.join(outdir, "annotation/blacklist.bed"))
    organismDictionary['blacklist_bed'] = os.path.join(outdir, "annotation/blacklist.bed")

if spikeinBlacklist:
    targets.append(os.path.join(outdir, "annotation/spikein_blacklist.bed"))
    organismDictionary['spikein_blacklist_bed'] = os.path.join(outdir, "annotation/spikein_blacklist.bed")

ignore = []
if ignoreForNormalization:
    f = open(ignoreForNormalization)
    ignore = [line.strip() for line in f]
    f.close()
    organismDictionary['ignoreForNormalization'] = " ".join(ignore)

if effectiveGenomeSize and effectiveGenomeSize > 0:
    os.makedirs(os.path.join(outdir, "genome_fasta"), exist_ok=True)
    f = open(os.path.join(outdir, "genome_fasta", "effectiveSize"), "w")
    f.write("{}\n".format(effectiveGenomeSize))
    f.close()
else:
    targets.append(os.path.join(outdir, "genome_fasta", "effectiveSize"))

globals().update(organismDictionary)
include: os.path.join(maindir, "shared", "rules", "createIndices.snakefile")

### execute before workflow starts #############################################
################################################################################
onstart:
    if "verbose" in config and config["verbose"]:
        print("--- Workflow parameters --------------------------------------------------------")
        print("output directory: ", outdir)
        print("genome name: ", genome)
        print("genome URL/path: ", genomeURL)
        print("GTF URL/path: ", gtfURL)
        print("tools: ", " ".join(tools))
        print("effectiveGenomeSize: ", effectiveGenomeSize)
        print("blacklist file: ", blacklist)
        print("ignore for normalization: ", " ".join(ignore))
        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_CREATE_INDEX_ENV]
        cf.writeTools(usedEnvs, outdir, "createIndices", maindir)

### main rule ##################################################################
################################################################################
rule all:
    input: targets

### execute after workflow finished ############################################
################################################################################
onsuccess:
    f = open(os.path.join(outdir, "genome_fasta", "effectiveSize"))
    organismDictionary['genome_size'] = int(float(f.read().strip()))
    f.close()
    orgDir_yamlPath = os.path.join(maindir, organismsDir, f"{genome}.yaml")
    outDir_yamlPath = os.path.join(outdir, f"{genome}.yaml")
    try:
        cf.write_configfile(orgDir_yamlPath, organismDictionary,trafo=None)
        print(f"{genome} yaml written to {orgDir_yamlPath}")
    except (PermissionError, FileNotFoundError) as e:
        print(f"Error writing in the organism dir {e}")
    cf.write_configfile(outDir_yamlPath, organismDictionary,trafo=None)
    print(f"{genome} yaml written to {outDir_yamlPath}")
    if "verbose" in config and config["verbose"]:
        print("\n--- index creation workflow finished successfully! --------------------------------\n")

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