import sys
import os,sys                                                              
import glob
import itertools
import subprocess
from math import log, ceil
from os import listdir
from os.path import isfile, join
from snakemake import logger as snake_logger
from shutil import which

CFG=workflow.overwrite_configfile

REF = config["reference"]["genome"]
# TODO: this is hack that should be solved more neatly
if "out_dir" in config:
    OUT_DIR = config["out_dir"]
else:
    OUT_DIR = config["prepare"]["files"]["output_dir"]
if "threads" in config:
    THREADS = int(config["threads"])
else:
    THREADS = 1

if "mikado" in config:
    MIKADO_MODES=config["mikado"]["modes"]
else:
    MIKADO_MODES = ["split", "nosplit", "permissive", "lenient", "stringent"]

# Directory shortcuts
OUT_DIR_FULL = os.path.abspath(OUT_DIR)
MIKADO_DIR = OUT_DIR + "/5-mikado"
MIKADO_DIR_FULL = os.path.abspath(MIKADO_DIR)
BLAST_DIR = MIKADO_DIR + "/blast"
BLAST_DIR_FULL = os.path.abspath(BLAST_DIR)
TDC_DIR = MIKADO_DIR + "/transdecoder"
TDC_DIR_FULL = os.path.abspath(TDC_DIR)
PROD_DIR = MIKADO_DIR + "/prodigal"
PROD_DIR_FULL = os.path.abspath(PROD_DIR)

CWD = os.getcwd()

BLASTX_TARGET = config["blastx"]["prot_db"]
BLASTX_MAX_TARGET_SEQS = config["blastx"]["max_target_seqs"]
BLASTX_EVALUE = config["blastx"]["evalue"]
BLASTX_CHUNKS = int(config["blastx"]["chunks"])
ASM_COLLECT = which("asm_collect.py")

CHUNK_ARRAY = []
if len(BLASTX_TARGET) > 0:
	for a in range(1,BLASTX_CHUNKS+1):
		val=str(a).zfill(3)
		CHUNK_ARRAY.append(val)


def loadPre(config, program):
	cc = config.get("load", dict()).get(program, "")
	if not cc:
		return ""
	else:
		return "set +u && {} &&".format(cc)


#########################
# Rules

rule all:
	input:
		MIKADO_DIR+"/pick/comparison.stats"
	output: touch(os.path.join(MIKADO_DIR, "all.done"))

rule clean:
	shell: "rm -rf {OUT_DIR}"

rule mikado_prepare:
	input: 
		ref=REF
	output:
		gtf=MIKADO_DIR+"/mikado_prepared.gtf",
		fa=MIKADO_DIR+"/mikado_prepared.fasta"
	params:
		load=loadPre(config, "mikado"),
		cfg=CFG
	threads: THREADS
	message: "Preparing transcripts using mikado"
	shell: "{params.load} mikado prepare --start-method=spawn --procs={threads} --fasta={input.ref} --json-conf={params.cfg} -od {MIKADO_DIR} 2>&1"

rule create_blast_database:
    input: fa=BLASTX_TARGET
    output: BLAST_DIR+"/index/blastdb-proteins.fa"
    message: "Creating the BLASTX database"
    params:
         fastas=" ".join(BLASTX_TARGET)
    shell: """sanitize_blast_db.py --out {output} {params.fastas}"""

rule split_fa:
	input: tr=rules.mikado_prepare.output.fa
	output: BLAST_DIR+"/fastas/split.done"
	params: 
		outdir=BLAST_DIR+"/fastas/chunk",
		chunks=config["blastx"]["chunks"],
		load=loadPre(config, "mikado")
	threads: 1
	message: "Splitting fasta: {input.tr}"
	shell: "{params.load} split_fasta.py -m {params.chunks} {input.tr} {params.outdir} && touch {output}"

if config["mikado"]["use_diamond"] is False:
    # Default, use BLAST
    rule make_blast:
        input: fa=BLAST_DIR+"/index/blastdb-proteins.fa"
        output: BLAST_DIR+"/index/blastdb-proteins.pog"
        params:
            db=BLAST_DIR+"/index/blastdb-proteins",
            load=loadPre(config, "blast")
        log: BLAST_DIR+"/blast.index.log"
        message: "Making BLAST protein database for: {input.fa}"
        shell: "{params.load} makeblastdb -in {input.fa} -out {params.db} -dbtype prot -parse_seqids > {log} 2>&1"

    rule blastx:
        input:
            db=rules.make_blast.output,
            split=rules.split_fa.output
        output: BLAST_DIR+"/xmls/chunk-{chunk_id}-proteins.xml.gz"
        params:
            tr=BLAST_DIR+"/fastas/chunk_{chunk_id}.fasta",
            db=BLAST_DIR + "/index/blastdb-proteins",
            load=loadPre(config, "blast"),
            uncompressed=BLAST_DIR+"/xmls/chunk-{chunk_id}-proteins.xml",
        log: BLAST_DIR + "/logs/chunk-{chunk_id}.blastx.log"
        threads: THREADS
        message: "Running BLASTX for mikado transcripts against: {params.tr}"
        shell: "{params.load} if [ -s {params.tr} ]; then blastx -num_threads {threads} -outfmt 5 -query {params.tr} -db {params.db} -evalue {BLASTX_EVALUE} -max_target_seqs {BLASTX_MAX_TARGET_SEQS} > {params.uncompressed} 2> {log}; else touch {params.uncompressed}; fi && gzip {params.uncompressed}"

else:
    # EXPERIMENTAL, DIAMOND itself at the moment does *not* support the necessary features!
    rule diamond_index:
        input:
            fa=BLAST_DIR+"/index/blastdb-proteins.fa"
        output: BLAST_DIR+"/index/blastdb-proteins.dmnd"
        params:
            db= BLAST_DIR+"/index/blastdb-proteins",
            load=loadPre(config, "diamond")
        log: BLAST_DIR+"/diamond.index.log"
        message: "Making DIAMOND protein database for: {input.fa}"
        threads: THREADS
        shell: "{params.load} diamond makedb --threads THREADS --in {input.fa} --db {params.db} 2> {log} > {log}"

    rule diamond:
        input:
            db=rules.diamond_index.output,
            split=rules.split_fa.output
        output: BLAST_DIR+"/xmls/chunk-{chunk_id}-proteins.xml.gz"
        params:
            load=loadPre(config, "diamond"),
            tr=BLAST_DIR+"/fastas/chunk_{chunk_id}.fasta"
        threads: THREADS
        log: BLAST_DIR + "/logs/chunk-{chunk_id}.blastx.log"
        shell: "{params.load} if [ -s {params.tr} ]; then diamond blastx --threads {threads} --outfmt xml --compress 1 --out {output} --max-target-seqs {BLASTX_MAX_TARGET_SEQS} --evalue {BLASTX_EVALUE} --db {input.db} --salltitles --query {params.tr} --sensitive > {log} 2> {log}; else touch {output}; fi"

rule blast_all:
	input: expand(BLAST_DIR + "/xmls/chunk-{chunk_id}-proteins.xml.gz", chunk_id=CHUNK_ARRAY)
	output: BLAST_DIR + "/blastx.all.done"
	shell: "touch {output}"

if config.get("mikado", dict()).get("use_prodigal", False) is False:
    rule transdecoder_lo:
        input: rules.mikado_prepare.output.fa
        output: TDC_DIR+"/transcripts.fasta.transdecoder_dir/longest_orfs.gff3"
        params:
            outdir=TDC_DIR_FULL,
            tr="transcripts.fasta",
            tr_in=MIKADO_DIR_FULL+"/mikado_prepared.fasta",
            load=loadPre(config, "transdecoder"),
            minprot=config["transdecoder"]["min_protein_len"]
        log: TDC_DIR_FULL+"/transdecoder.longorf.log",
            # ss="-S" if MIKADO_STRAND else ""
        threads: 1
        message: "Running transdecoder longorf on Mikado prepared transcripts: {input}"
        run:
            if config.get("transdecoder", dict()).get("execute", True) is True:
                shell("{params.load} cd {params.outdir} && ln -sf {params.tr_in} {params.tr} && TransDecoder.LongOrfs -m {params.minprot} -t {params.tr} > {log} 2>&1")
            else:
                shell("mkdir -p $(dirname {output}) && touch {output}")

    rule transdecoder_pred:
        input:
            mikado=rules.mikado_prepare.output.fa,
            trans=rules.transdecoder_lo.output
        output: TDC_DIR+"/transcripts.fasta.transdecoder.bed"
        params: outdir=TDC_DIR_FULL,
            tr_in=MIKADO_DIR_FULL+"/mikado_prepared.fasta",
            lolog=TDC_DIR_FULL+"/transdecoder.longorf.log",
            plog=TDC_DIR_FULL+"/transdecoder.predict.log",
            tr="transcripts.fasta",
            load=loadPre(config, "transdecoder")
            # ss="-S" if MIKADO_STRAND else ""
        log: TDC_DIR_FULL+"/transdecoder.predict.log"
        threads: THREADS
        message: "Running transdecoder predict on Mikado prepared transcripts: {input}"
        run:
            if config.get("transdecoder", dict()).get("execute", True) is True:
                shell("{params.load} cd {params.outdir} && TransDecoder.Predict -t {params.tr} --cpu {threads} > {log} 2>&1")
            else:
                shell("mkdir -p $(dirname {output}) && touch {output}")
    orf_out = rules.transdecoder_pred.output
else:
    rule prodigal:
        input: rules.mikado_prepare.output.fa
        output: PROD_DIR+"/transcripts.fasta.prodigal.gff3"
        params:
            outdir=PROD_DIR_FULL,
            tr="transcripts.fasta",
            tr_in=MIKADO_DIR_FULL+"/mikado_prepared.fasta",
            tr_out="transcripts.fasta.prodigal.gff3",
            load=loadPre(config, "prodigal"),
            minprot=config["transdecoder"]["min_protein_len"]
        log: PROD_DIR_FULL+"/prodigal.log"
        threads: 1
        message: "Running PRODIGAL on Mikado prepared transcripts: {input}"
        shell: "{params.load} mkdir -p {params.outdir} && cd {params.outdir} && ln -sf {params.tr_in} {params.tr} && prodigal -f gff -g 1 -i {params.tr} -o {params.tr_out} > {log} 2>&1"
    orf_out = rules.prodigal.output

rule genome_index:
	input: os.path.abspath(REF)
	output: MIKADO_DIR+"/"+os.path.basename(REF)+".fai"
	params: load=loadPre(config, "samtools"),
		fa=MIKADO_DIR+"/"+os.path.basename(REF)
	threads: 1
	message: "Using samtools to index genome"
	shell: "ln -sf {input} {params.fa} && touch -h {params.fa} && {params.load} samtools faidx {params.fa}"

rule mikado_serialise:
	input: 
		blast=rules.blast_all.output,
		orfs=orf_out,
		fai=rules.genome_index.output,
		transcripts=rules.mikado_prepare.output.fa
	output: db=MIKADO_DIR+"/mikado.db"
	log: MIKADO_DIR+"/mikado_serialise.err"
	params:
	    cfg=CFG,
		blast="--xml=" + BLAST_DIR+"/xmls" if len(BLASTX_TARGET) > 0 else "",
		load=loadPre(config, "mikado"),
		blast_target="--blast_targets=" + BLAST_DIR+"/index/blastdb-proteins.fa" if len(BLASTX_TARGET) > 0 else ""
	threads: THREADS
	message: "Running Mikado serialise to move numerous data sources into a single database"
	run:
	    if config.get("mikado", dict()).get("use_prodigal") is True:
	        params.orfs = "--orfs="+str(rules.prodigal.output)
	    elif config.get("transdecoder", dict()).get("execute", True) is True:
	        params.orfs = "--orfs="+str(rules.transdecoder_pred.output)
	    else:
	        params.orfs = ""
	    shell("{params.load} mikado serialise {params.blast} {params.blast_target} --start-method=spawn --transcripts={input.transcripts} --genome_fai={input.fai} --json-conf={params.cfg} --force {params.orfs} -od {MIKADO_DIR} --procs={threads} > {log} 2>&1")

rule mikado_pick:
	input:
		gtf=rules.mikado_prepare.output.gtf,
		db=rules.mikado_serialise.output
	output:
		loci=os.path.join(MIKADO_DIR, "pick", "{mode}", "mikado-{mode}.loci.gff3")
	log: os.path.join(MIKADO_DIR, "pick", "{mode}", "mikado-{mode}.pick.err")
	params:
		cfg=CFG,
		load=loadPre(config, "mikado"),
		outdir=os.path.join(MIKADO_DIR, "pick", "{mode}")
	threads: THREADS
	message: "Running mikado picking stage"
	shell: "{params.load} mikado pick --source Mikado_{wildcards.mode} --mode={wildcards.mode} --procs={threads} --start-method=spawn --json-conf={params.cfg} -od {params.outdir} --loci_out mikado-{wildcards.mode}.loci.gff3 -lv INFO {input.gtf} -db {input.db} > {log} 2>&1"

rule mikado_stats:
    input:
        rules.mikado_pick.output.loci
    output:
        stats=os.path.join(MIKADO_DIR,
                           "pick", "{mode}",
                           "mikado-{mode}.loci.stats")
    params: load=loadPre(config, "mikado")
    shell: "{params.load} mikado util stats {input} > {output}"

rule mikado_collect_stats:
	input:
	    mikado=expand(MIKADO_DIR+"/pick/{mode}/mikado-{mode}.loci.stats", mode=MIKADO_MODES)
	output: MIKADO_DIR+"/pick/comparison.stats"
	params:
	    load=loadPre(config, "mikado")
	threads: 1
	message: "Collecting mikado statistics"
	shell: "{params.load} {ASM_COLLECT} {input.mikado} > {output}"

rule complete:
  input: rules.mikado_collect_stats.output
