import sys
import re
import shutil

if "workdir" in config:
    workdir: config["workdir"]

sys.setrecursionlimit(1000000)

SCRIPT="python3 " + srcdir("seqseqpan.py")
JAR="java -jar " + srcdir("java/XMFAConsensus/XMFAconsensus_majorityvote.jar")
GENOMEDESC_SCRIPT= "python3 " + srcdir("genomedescription.py")

with open(config["genomefile"], "r") as f:
    GENOMES = f.readlines()

GENOME_DICT = {str(idx): GENOMES[idx].strip() for idx in range(len(GENOMES))}
GENOME_NAMES = [str(idx) for idx in range(len(GENOMES))]
MERGE = config.get("merge", "")
PREV_PAN = config.get("pangenome", "")

if MERGE.lower() == "false":
    MERGE = False
else:
    MERGE = True

if PREV_PAN != "":
     PREV_PAN = re.sub("\.xmfa$", "", PREV_PAN)

PMAUVE = config.get("pmauve", "progressiveMauve")
BLAT = config.get("blat", "blat")

ruleorder: use_prev_pan > resolve


# determine which file is translated to output xmfa
def output_xmfa(wildcards):
    xmfa_f = reconstruct_or_not({"genome": str(len(GENOME_NAMES)-1)})
    return xmfa_f

def is_first_align(genome):
    if "_single" in genome:
        return False
    else:
       return(int(genome) <=1)

# reconstructing not possible/necessary for alignment of first two genomes
def reconstruct_or_not(wildcards):
    return ("add"+wildcards["genome"]+"_joined.xmfa" if (is_first_align(wildcards["genome"]) and PREV_PAN == "") else "add"+wildcards["genome"]+"_reconstruct.xmfa")

def not_resolved(wildcards):
    return (is_first_align(wildcards["genome"]) and PREV_PAN == "" )

# resolve not possible/necessary for alignment of first two genomes
def resolved_or_not(wildcards):
    return ("add"+wildcards["genome"]+"_align.xmfa" if not_resolved(wildcards) else "add"+wildcards["genome"]+"_resolve.xmfa")

# merge after alignment
def merged_or_not(wildcards):
    return("add"+wildcards["genome"]+"_merge.xmfa" if MERGE else resolved_or_not(wildcards))

def merged_or_not_param(wildcards):
    return ("-m" if MERGE else "")

def prev_pan_or_not_param(wildcards):
    return ("-a "+ PREV_PAN +"_genomedescription.txt" if PREV_PAN != "" else "")

def mauve_input(wildcards):
    return (GENOME_DICT[wildcards["genome"]])

def consensus_file(wildcards):
    new_nr = prev_genome(wildcards)
    if int(new_nr) > 0 or (int(new_nr) <= 0 and PREV_PAN != ""):
       consgen = ["add"+new_nr+"_consensus.fasta.blockseparated.fasta", "add"+new_nr+"_consensus.fasta", "add"+new_nr+"_consensus.fasta.blockseparated.idx"]
    elif PREV_PAN == "":
       consgen = [GENOME_DICT[new_nr]]
    else:
        print("Error in defining the second file for alignment.")
    return consgen

def prev_genome(wildcards):
    return(str(int(wildcards['genome']) - 1))

def prev_joined_filename(wildcards):
    return "".join(["add", prev_genome(wildcards), "_joined.xmfa"])


#  "xmfa_prev" always needed when using consensus! --> xmfa file used to build consensus


rule all:
    input:  config["outfilename"]+"_genomedescription.txt",
            config["outfilename"]+"_consensus.fasta.blockseparated.fasta",
            config["outfilename"]+"_consensus.fasta.blockseparated.idx",
            config["outfilename"]+"_consensus.fasta.idx",
            config["outfilename"]+"_consensus.fasta",
            config["outfilename"]+".xmfa",
            config["outfilename"]+".maf"
    
rule fn_translate:
    input:  cb = "add"+str(len(GENOME_NAMES)-1)+"_consensus.fasta.blockseparated.fasta",
            cbi = "add"+str(len(GENOME_NAMES)-1)+"_consensus.fasta.blockseparated.idx",
            ci = "add"+str(len(GENOME_NAMES)-1)+"_consensus.fasta.idx",
            c = "add"+str(len(GENOME_NAMES)-1)+"_consensus.fasta",
            xmfa = output_xmfa
            
    output: cb = config["outfilename"]+"_consensus.fasta.blockseparated.fasta",
            cbi = config["outfilename"]+"_consensus.fasta.blockseparated.idx",
            ci = config["outfilename"]+"_consensus.fasta.idx",
            c = config["outfilename"]+"_consensus.fasta",
            xmfa = config["outfilename"]+".xmfa"
            
    shell:
        '''
        name={input.c}
        out={output.c}
        name="${{name//_consensus.fasta}}"
        out="${{out//_consensus.fasta}}"
        cp {input.cb} {output.cb}
        cp {input.cbi} {output.cbi}
        cp {input.ci} {output.ci}
        cp {input.c} {output.c}
        cp {input.xmfa} {output.xmfa}
        
        sed -i "1 s|$name|$out|g" {output.c}
        sed -i "1 s|$name|$out|g" {output.cb}
        sed -i "1 s|$name|$out|" {output.ci}
        sed -i "2 s|$name|$out|" {output.ci}
        sed -i "1 s|$name|$out|" {output.cbi}
        sed -i "2 s|$name|$out|" {output.cbi}
        '''

        
rule consensus:
    input: reconstruct_or_not
    
    output: temp("add{genome,-*\d+(_[12]_single)*}_consensus.fasta.blockseparated.fasta"),
            temp("add{genome,-*\d+(_[12]_single)*}_consensus.fasta.blockseparated.idx"),
            temp("add{genome,-*\d+(_[12]_single)*}_consensus.fasta.idx"),
            temp("add{genome,-*\d+(_[12]_single)*}_consensus.fasta")
    shell:
        #"{SCRIPT} consensus -x {input} -p . -n add{wildcards.genome}"
        "{JAR} --xmfa {input} --outfasta {output[3]} --outidx {output[2]} --outblocksepfasta {output[0]} --outblocksepidx {output[1]}"


rule consensus_singleton:
    input: "add{genome}_single.xmfa"

    output: temp("add{genome,\d+_[12]}_single_consensus.fasta.blockseparated.fasta"),
            temp("add{genome,\d+_[12]}_single_consensus.fasta.blockseparated.idx"),
            temp("add{genome,\d+_[12]}_single_consensus.fasta.idx"),
            temp("add{genome,\d+_[12]}_single_consensus.fasta")
    shell:
        #"{SCRIPT} consensus -x {input} -p . -n add{wildcards.genome}"
        "{JAR} --xmfa {input} --outfasta {output[3]} --outidx {output[2]} --outblocksepfasta {output[0]} --outblocksepidx {output[1]}"


rule resolve:
    input:  xmfa= "add{genome}_align.xmfa",
            consensus=consensus_file
    output: temp("add{genome,\d+}_resolve.xmfa")
    shell:
        "{SCRIPT} resolve -x {input.xmfa} -p . -n add{wildcards.genome} -c {input.consensus[1]}"
        
rule reconstruct:
    input: xmfa_prev= lambda wildcards: reconstruct_or_not({"genome":prev_genome(wildcards)}),
           xmfa= "add{genome}_joined.xmfa",
           consensus=consensus_file
    output: temp("add{genome,\d+}_reconstruct.xmfa")
    shell:
        "{SCRIPT} reconstruct -x {input.xmfa} -p . -n add{wildcards.genome} -c {input.consensus[1]}"

rule reconstruct_singleton:
    input: xmfa="add{genome, \d+}_single_realign.xmfa",
           xmfa_prev_1="add{genome, \d+}_1_single.xmfa",
           xmfa_prev_2="add{genome, \d+}_2_single.xmfa",
           consensus_1="add{genome,\d+}_1_single_consensus.fasta",
           consensus_2="add{genome,\d+}_2_single_consensus.fasta",
           cb_1="add{genome,\d+}_1_single_consensus.fasta.blockseparated.fasta",
           cb_2="add{genome,\d+}_2_single_consensus.fasta.blockseparated.fasta",
           cbi_1="add{genome,\d+}_1_single_consensus.fasta.blockseparated.idx",
           cbi_2="add{genome,\d+}_2_single_consensus.fasta.blockseparated.idx",
           ci_1="add{genome,\d+}_1_single_consensus.fasta.idx",
           ci_2="add{genome,\d+}_2_single_consensus.fasta.idx"

    output: temp("add{genome, \d+}_single_temp_reconstruct.xmfa"),
            temp("add{genome, \d+}_single_reconstruct.xmfa")
    #params: merge=merged_or_not_param
    shell:
        '''
        {SCRIPT} reconstruct -x {input.xmfa} -p . -n add{wildcards.genome}_single_temp -c {input.consensus_1} --quiet
        {SCRIPT} reconstruct -x {output[0]} -p . -n add{wildcards.genome}_single -c {input.consensus_2} --quiet
        '''


rule resolve_singleton:
    input: xmfa="add{genome, \d+}_single_align.xmfa",
           xmfa_prev_1="add{genome, \d+}_1_single.xmfa",
           xmfa_prev_2="add{genome, \d+}_2_single.xmfa",
           consensus_1="add{genome,\d+}_1_single_consensus.fasta",
           consensus_2="add{genome,\d+}_2_single_consensus.fasta",
           cb_1="add{genome,\d+}_1_single_consensus.fasta.blockseparated.fasta",
           cb_2="add{genome,\d+}_2_single_consensus.fasta.blockseparated.fasta",
           cbi_1="add{genome,\d+}_1_single_consensus.fasta.blockseparated.idx",
           cbi_2="add{genome,\d+}_2_single_consensus.fasta.blockseparated.idx",
           ci_1="add{genome,\d+}_1_single_consensus.fasta.idx",
           ci_2="add{genome,\d+}_2_single_consensus.fasta.idx"
    output: temp("add{genome, \d+}_single_temp_resolve.xmfa"),
            temp("add{genome, \d+}_single_resolve.xmfa")
    shell:
        '''
        {SCRIPT} resolve -x {input.xmfa} -p . -n add{wildcards.genome}_single_temp -c {input.consensus_1}
        {SCRIPT} resolve -x {output[0]} -p . -n add{wildcards.genome}_single -c {input.consensus_2}
        '''


rule join:
    input:  xmfa_pair="add{genome}_pairblocks.xmfa",
            xmfa_singlealign="add{genome}_single_reconstruct.xmfa",
            xmfa_orig="add{genome}_realign.xmfa"
    output: temp("add{genome, \d+}_joined.xmfa")
    run:
        # check if singleton reconstruct is empty
        count = 0
        with open(input.xmfa_singlealign) as f:
            for line in f:
                if line[0] == ">":
                    count=1
                    break
        empty = (count == 0)

        # if singleton reconstruct is empty (at least one of genomes does not have singletons)
        # copy file from before split and sort it, as singletons are sorted
        if empty:
            shell(" ".join([SCRIPT, "xmfa -x", input.xmfa_orig, "-o 1 -p . -n add"+wildcards["genome"]+"_joined", "--quiet"]))
        else:
            shell(" ".join([SCRIPT, "join -x", input.xmfa_pair, "-y", input.xmfa_singlealign, "-o 1 -p . -n add"+wildcards["genome"], "--quiet"]))


rule realign:
    input: merged_or_not
    output: temp("add{genome,\d+(_single)*}_realign.xmfa")
    params: lambda wildcards: "" if not_resolved(wildcards) else "--quiet"
    shell:
        "{SCRIPT} realign -x {input} -p . -n add{wildcards.genome} --blat {BLAT} {params}"

        
rule merge:
    input: resolved_or_not
    output: temp("add{genome,\d+(_single)*}_merge.xmfa")
    params: lambda wildcards: "" if not_resolved(wildcards) else "--quiet"
    shell:
        "{SCRIPT} merge -x {input} -p . -n add{wildcards.genome} {params}"
        
        
rule align:
    input:  consensus=consensus_file, 
            new=mauve_input
    output: temp("add{genome, \d+}_align.xmfa"),
            temp("add{genome, \d+}_align.xmfa.bbcols"),
            temp("add{genome, \d+}_align.xmfa.backbone")
    shell:
        '''
        input=`basename {input.consensus[0]}`
        input2=`basename {input.new}`
        
        input=temp_$input
        input2=temp_$input2
        
        cp {input.consensus[0]} $input
        cp {input.new} $input2
        
        {PMAUVE} $input $input2 --output={output[0]} > /dev/null
        
        sed -i "s|$input|{input.consensus[0]}|g" {output[0]}
        sed -i "s|$input2|{input.new}|g" {output[0]}
        
        rm $input.sslist
        rm $input
        rm $input2.sslist
        rm $input2
        '''

rule genome_count_split:
    input: "add{genome}_realign.xmfa"
    output: temp("add{genome, \d+}_pairblocks.xmfa"),
            temp("add{genome,\d+}_1_single.xmfa"),
            temp("add{genome,\d+}_2_single.xmfa")
    shell:
        "{SCRIPT} blockcountsplit -x {input} -p . -n add{wildcards.genome}"


rule align_singleton:
    input: consensus_1="add{genome, \d+}_1_single_consensus.fasta.blockseparated.fasta",
           consensus_2="add{genome, \d+}_2_single_consensus.fasta.blockseparated.fasta",
           xmfa_1="add{genome, \d+}_1_single.xmfa",
           xmfa_2="add{genome, \d+}_2_single.xmfa",
           xmfa_pair="add{genome, \d+}_pairblocks.xmfa",

    output: temp("add{genome, \d+}_single_align.xmfa"),
            temp("add{genome, \d+}_single_align.xmfa.bbcols"),
            temp("add{genome, \d+}_single_align.xmfa.backbone"),
            temp("add{genome, \d+}_1_single_consensus.fasta.blockseparated.fasta.sslist"),
            temp("add{genome, \d+}_2_single_consensus.fasta.blockseparated.fasta.sslist")
    run:

        def empty_xmfa_file(cons_1, cons_2, xmfa_pair_f, output_f):
            header = ""
            with open(xmfa_pair_f, "r") as input:
                for line in input:
                    if line[0] == "#":
                        header += line
                    else:
                        break

            m = re.search("Sequence1File\t(.*)\n", header)
            f1 = m.group(1)
            m = re.search("Sequence2File\t(.*)\n", header)
            f2 = m.group(1)

            header = header.replace(f1, cons_1)
            header = header.replace(f2, cons_2)

            with open(output_f, "w") as out:
                out.write(header)


        # check if fasta files only contain header
        count = 0
        length = 0
        with open(input.consensus_1) as f:
            first_line = f.readline()
            for line in f:
                count = 1
                length = len(line)
                break
        one = (count > 0 and length > 23)

        count = 0
        length = 0
        with open(input.consensus_2) as f:
            first_line = f.readline()
            for line in f:
                count = 1
                length = len(line)
                break

        two = (count > 0 and length > 23)

        # align if both fasta files valid
        if one and two:
            shell(" ".join([PMAUVE, input.consensus_1, input.consensus_2, "--output=" + output[0], "> /dev/null"]))
        else: # else create empty temporary files
            for i in range(1,5):
                open(output[i], "w").close()
            empty_xmfa_file(input.consensus_1, input.consensus_2, input.xmfa_pair, output[0])


rule maf:
    input: xmfa="{genome}.xmfa",
           genomedesc=config["outfilename"]+"_genomedescription.txt"
    output: "{genome}.maf"
    shell:
        "{SCRIPT} maf -x {input.xmfa} -p . -n {wildcards.genome} -g {input.genomedesc}"


rule split:
    input: xmfa="add{genome}_reconstruct.xmfa",
           genomedesc=config["outfilename"]+"_genomedescription.txt"
    output: "add{genome}_split.xmfa"
    shell:
        "{SCRIPT} split -x {input.genome} -p . -n add{wildcards.genome} -g {input.genomedesc}"


rule genomedesc:
    input: config["genomefile"]
    output: config["outfilename"]+"_genomedescription.txt"
    params: add=prev_pan_or_not_param
    shell:
        "{GENOMEDESC_SCRIPT} -i {input} -o {output} {params.add}"


rule use_prev_pan:
    input: PREV_PAN+".xmfa"
    output: temp("add-1_reconstruct.xmfa")
    shell:
        "cp {input} {output} "
