#! /usr/bin/env python
# -*- coding: utf-8 -*-
'''
Infer trees from germinal center data
'''
import os
import atexit
import shutil
import re


def which(executable):
    for path in os.environ["PATH"].split(os.pathsep):
        if os.path.exists(os.path.join(path, executable)):
            return os.path.realpath(os.path.join(path, executable))
    return None

# the following must be exported by parent SConstruct/SConscript
Import('env gctree igphyml dnaml quick frame fasta fasta2 outdir naiveID converter CommandRunner')


def delete_on_failure():
    from SCons.Script import GetBuildFailures
    for bf in GetBuildFailures():
        try:
            if False:  # Turn on if failed folder should be deleted
                # shutil.rmtree(os.path.dirname(bf.node.abspath))
                print('There was an error for target:', bf.node.abspath)
                print('To resolve this the base folder was deleted. '
                      'Restarting will possibly solve the issue in case '
                      'this was just a bad simulation. Deleted folder:', os.path.dirname(bf.node.abspath))
        except OSError:
            pass  # Allow this
# atexit.register(delete_on_failure)

return_list = []


if gctree or igphyml or dnaml:
    # fasta_str = str(fasta)
    if fasta2 is not None: #isinstance(fasta, list):
	basename = 'combined'
    elif isinstance(fasta, str):
        basename = '.'.join(os.path.basename(fasta).split('.')[:-1])
    else:
        basename = 'gctree.simulation'
    # parse fasta file to phylip, interpreting integer names as frequencies
    frame_arg = ' --frame {} '.format(frame) if frame is not None else ''
    converter_arg = ' --converter {} '.format(converter) if converter is not None else ''
    phylip = env.Command([os.path.join(outdir, os.path.splitext(os.path.basename(fasta))[0] + '.phylip') if fasta2 is None and isinstance(fasta, str) else os.path.join(outdir, 'gctree.combined.phylip') if fasta2 is not None else os.path.join(outdir, 'gctree.simulation.phylip'),
                          os.path.join(outdir, basename)+'.counts',
                          os.path.join(outdir, basename)+'.idmap'],
                          fasta,
                          'stdbuf -oL python bin/fasta2phylip.py $SOURCES --countfile ${TARGETS[1]} --idmapfile ${TARGETS[2]} '+ frame_arg+converter_arg+
                          '--naive '+naiveID+' > ${TARGETS[0]}')
# Update the naive ID:
naiveID = naiveID.lower()

if dnaml:
    dnaml_config = env.Command(os.path.join(outdir, 'dnaml/config.cfg'),
                                 phylip[0],
                                 'stdbuf -oL python bin/mkconfig.py $SOURCE dnaml --naive ' + naiveID + ' > $TARGET')

    # run dnaml (from phylip package) to generate maximum likelihood tree
    dnaml = CommandRunner(map(lambda x: os.path.join(outdir, x), ['dnaml/outtree', 'dnaml/outfile', 'dnaml/log.log']),
                       dnaml_config,
                       'cd ' + outdir + '/dnaml && rm -f outfile outtree && dnaml < ${SOURCE.file} > ${TARGETS[2].file}')
    # Manually depend on phylip so that we rerun dnaml if the input sequences change (without this, dnaml will
    # only get rerun if one of the targets are removed or if the input dnaml_config file is changed).
    env.Depends(dnaml, phylip)

    dnaml_tree = env.Command([os.path.join(outdir, 'dnaml.inferred_tree.p'),
                                 os.path.join(outdir, 'dnaml.inferred_tree.log')],
                                [dnaml[1], phylip[1]],
                                'TMPDIR=/tmp xvfb-run -a python bin/phylip_parse.py ${SOURCES} --outputfile ${TARGETS[0]} --naive '+naiveID+' > ${TARGETS[1]}')
    return_list.append(dnaml_tree)

if gctree:
    # make config file for dnapars
    quick_arg = '--quick' if quick else ''
    dnapars_config = env.Command(os.path.join(outdir, 'dnapars/config.cfg'),
                                 phylip[0],
                                 'stdbuf -oL python bin/mkconfig.py $SOURCE dnapars --naive ' + naiveID + ' ' + quick_arg + ' > $TARGET')

    # run dnapars (from phylip package) to generate parsimony trees
    dnapars = CommandRunner(map(lambda x: os.path.join(outdir, x), ['dnapars/outtree', 'dnapars/outfile', 'dnapars/log.log']),
                       dnapars_config,
                       'cd ' + outdir + '/dnapars && rm -f outfile outtree && dnapars < ${SOURCE.file} > ${TARGETS[2].file}')
    # Manually depend on phylip so that we rerun dnapars if the input sequences change (without this, dnapars will
    # only get rerun if one of the targets are removed or if the input dnapars_config file is changed).
    env.Depends(dnapars, phylip)

    # ML tree from parsimony trees
    # NOTE: xvfb-run is needed because of issue https://github.com/etetoolkit/ete/issues/101
    # NOTE: TMPDIR is needed due to a an xvfb conflict that Ben (scicomp ninja) discovered
    gctree_outbase = os.path.join(outdir, 'gctree')
    frame_arg = ' --frame {} '.format(frame) if frame is not None else ''
    gctree_infer = CommandRunner([gctree_outbase+'.inference.parsimony_forest.p',
                                gctree_outbase+'.inference.log'],
                                [dnapars[1], phylip[1]],
                                'TMPDIR=/tmp xvfb-run -a stdbuf -oL python bin/gctree.py infer ${SOURCES[0]} ${SOURCES[1]} --naive '+naiveID+
                                ' --outbase '+gctree_outbase+
                                frame_arg+
                                ' > ${TARGETS[1]}')
    return_list.append(gctree_infer)


if igphyml:
    # basename = 'igphyml'
    # Need the dedup phy as fa for ASR:
    dedup_fasta = env.Command([os.path.join(outdir, basename+'.dedup.fasta'),
                               os.path.join(outdir, basename+'.dedup.log')],
                              phylip[0],
                              'seqmagick convert $SOURCE ${TARGETS[0]} > ${TARGETS[1]}')

    # Run IgPhyML with the GY94 model to get the initial topology:
    igphyml_gy94_topology = CommandRunner([os.path.join(outdir, basename+'.phylip_igphyml_tree.txt_gy94'),
                                           os.path.join(outdir, basename+'.phylip_igphyml_tree.txt_gy94.log')],
                                           phylip[0],
                                           'igphyml -i $SOURCE -m GY -w M0 -t e --run_id gy94 > ${TARGETS[1]}')

    # Run IgPhyML with the HLP16 model using the GY94 tree topology as starting point:
    igphyml_opti = 'tlr'  # <--- can be changed to 'lr' or 'r'
    ### Full motif model: --motifs WRC_2:0,GYW_0:1,WA_1:2,TW_0:3,SYC_2:4,GRS_0:5 --hotness e,e,e,e,e,e
    ### Less parameter rich motif model: --motifs WRC_2:0,GYW_0:0,WA_1:1,TW_0:2,SYC_2:3,GRS_0:3 --hotness e,e,e,e
    igphyml_hlp16 = CommandRunner([os.path.join(outdir, basename+'.phylip_igphyml_tree.txt_hlp16'),
                              os.path.join(outdir, basename+'.phylip_igphyml_stats.txt_hlp16'),
                              os.path.join(outdir, basename+'.phylip_igphyml_tree.txt_hlp16.log')],
                             [phylip[0], igphyml_gy94_topology[0]],
                             'igphyml --motifs WRC_2:0,GYW_0:1,WA_1:2,TW_0:3,SYC_2:4,GRS_0:5 --hotness e,e,e,e,e,e -i ${SOURCES[0]} -u ${SOURCES[1]} -m HLP17 --root '+naiveID+' -o '+igphyml_opti+' --run_id hlp16 > ${TARGETS[2]}')

    # On the HLP16 output tree make the naive sequence an outgroup:
    naive_outgroup = env.Command([os.path.join(outdir, basename+'.phylip_igphyml_tree.txt_hlp16.outgroup'),
                                  os.path.join(outdir, basename+'.phylip_igphyml_tree.txt_hlp16.outgroup.log')],
                                 igphyml_hlp16[0],
                                 'python igphyml_files/igphyml_tools.py reroot --tree $SOURCE --reroot_tree ${TARGETS[0]} --pattern '+naiveID+' --outgroup 1 > ${TARGETS[1]}')

    igphyml_dir = which('igphyml')
    assert(igphyml_dir is not None)
    igphyml_dir = re.sub(r'/src/\w+$', '', igphyml_dir)
    # Run IgPhyML ASR script:
    run_ASR = CommandRunner([os.path.join(outdir, basename+'.MLcodons.fa'),
                        os.path.join(outdir, basename+'.igphyml_hlp16.MLcodons.log')],
                       [igphyml_hlp16[1], naive_outgroup[0], dedup_fasta[0]],
                       'stdbuf -oL perl igphyml_files/ancReconstructHLP16.pl PLACEHOLDER_CONFIG_FILE -rooted 1 -length D -stats ${SOURCES[0]} -tree ${SOURCES[1]} -seqfile ${SOURCES[2]} -outdir '+outdir+' -stem '+basename+' -rootid '+naiveID+' -igphyml '+igphyml_dir+'> ${TARGETS[1]}')

    # Convert the ASR output to an collapsed forest with an ete3 tree and pickle it:
    igphyml_infer = env.Command([os.path.join(outdir, 'igphyml.inferred_tree.p'),
                                 os.path.join(outdir, 'igphyml.inferred_tree.log')],
                                [naive_outgroup[0], phylip[1], run_ASR],
                                'TMPDIR=/tmp xvfb-run -a python igphyml_files/igphyml_tools.py ASR_parser --tree ${SOURCES[0]} --counts ${SOURCES[1]} --asr_seq ${SOURCES[2]} --outbase '+os.path.join(outdir, 'igphyml.inferred_tree')+ ' --naive '+naiveID+' > ${TARGETS[1]}')
    return_list.append(igphyml_infer)

Return('return_list')
