#!/usr/bin/env python

import math
import os
import os.path
import phylip
import shutil
import subprocess
import sys
import time

#Version  8/28/2007
# Run distance programs as a command
#Synopsis: distance.csh infile utree ufn replicates jumble jseed numjum method \
#           negbranch outgroup subrep global termout printdata outfile treefile

#  It is important that when running with the USERTREE option, 
#  the usertree should come from the same program that is currently
#  being run. eg. don't use a kitsch-generated tree as input to fitch.

""" ensure that there are enough command line arguments to parse """
if len(sys.argv) < 17:
    print("Usage: distance.py  INFILE  UTREE  UFN  REPLICATES  JUMBLE  JSEED  NUMJUM")
    print("         METHOD  NEGBRANCH  OUTGROUP  SUBREP  GLOBAL  TERMOUT  PRINTDATA")
    print("         OUTFILE  TREEFILE")
    exit();

#Convert arguments to variables
INFILE        = sys.argv[1]
UTREE         = sys.argv[2]
UFN           = sys.argv[3]
REPLICATES    = sys.argv[4]
JUMBLE        = sys.argv[5]
JSEED         = sys.argv[6]
NUMJUM        = sys.argv[7]
METHOD        = sys.argv[8]
NEGBRANCH     = sys.argv[9]
OUTGROUP      = sys.argv[10]
SUBREP        = sys.argv[11]
GLOBAL        = sys.argv[12]
TERMOUT       = sys.argv[13]
PRINTDATA     = sys.argv[14]
OUTFILE       = sys.argv[15]
TREEFILE      = sys.argv[16]


# Remember where we started
STARTDIR = os.getcwd()
print 'Starting Directory: ' + STARTDIR

time_elapsed = time.time()

os.nice(10)
termout_h = open(TERMOUT, 'w')
msgfile_h = open('MSGFILE', 'a')
p_PROGRAM = subprocess.Popen([PROGRAM], stdin=subprocess.PIPE, stdout=termout_h)

#----------------- generate keyboard input to send to program -----

# Make a temporary directory in which to run the program
TEMPDIR = 'DISTANCE.' + os.getpid()
os.mkdir(TEMPDIR)
shutil.copy(INFILE, os.path.join(TEMPDIR, 'infile'))

if UTREE == 'y':
    # Make sure that treefile begins with number of trees on first
    # line of file. If first line in file has parentheses, the
    # number must be added.
    infile_h = open(os.path.join(TEMPDIR, 'infile'), 'a')
    ufn_h = open (UFN, 'r')
    ufn_lines = ufn_h.readlines()
    ufn_h.close()

    sc_count = 0
    if ufn_lines[0].contains('{'):
        for line in ufn_lines:
            if line.contains(';'):
                sc_count = sc_count + 1
        if math.log(sc_count) < 7:
            infile.write('       '.substring(math.log(sc_count)))
        infile.write(sc_count)

    for line in ufn_lines:
        infile.write(line)
    infile_h.close()

os.chdir(TEMPDIR)

# Choose program and set program-specific parameters
if METHOD == "f" or METHOD == "k":
    if UTREE == 'y':
        p_PROGRAM.stdin.write('u\n')
    if METHOD == 'f':
        PROGRAM = 'fitch'
        # Global rearrangements
        if GLOBAL == 'y':
            p_PROGRAM.stdin.write('g\n')

        # Outgroup
        OUTGROUP = phylip.do_outgroup(OUTGROUP, p_PROGRAM.stdin)
    else:
        PROGRAM = 'kitsch'

    # Negative branch lengths allowed?
    if NEGBRANCH == 'y':
        p_PROGRAM.stdin.write('-\n')
else:
    PROGRAM = 'neighbor'
    # For Neighbor Joining, choose NJ or UPGMA
    if METHOD == "u":
        p_PROGRAM.stdin.write('n\n')
    if UTREE == 'y':
        msgfile_h.write('>>> USERTREE option NOT AVAILABLE in NEIGHBOR\n')

# Subreplicates?
if SUBREP == 's':
    p_PROGRAM.stdin.write('s\n')

# if bootstrapping, REPLICATES will be > 1
if REPLICATES > 1:
    p_PROGRAM.stdin.write('m\n')
    p_PROGRAM.stdin.write(REPLICATES + '\n')

# Jumble
if JUMBLE == 'y':
    jumble(p_PROGRAM.stdin, msgfile_h, True)
	
# Should sequence data be printed?
if PRINTDATA == 'y':
    p_PROGRAM.stdin.write('1\n')

# When resampling or jumbling, turn off printing trees to outfile
if REPLICATES > 1 or JUMBLE == 'y':
    p_PROGRAM.stdin.write('3\n')

#accept current settings and do the analysis
p_PROGRAM.stdin.write('y\n')

#-------- Run distance program -----------
p_PROGRAM.wait()
time_elapsed = time.time() - time_elapsed;
termout_h.close()
msgfile_h.close()

outfile_h = open('outfile', 'a')
outfile_h.write('Execution times on `hostname`: ' + time_elapsed + '\n')
outfile_h.close()

#----------- Return results to calling directory----------------
# When using resampling or jumbling, filter the treefile through
# consense to generate an unrooted consensus tree.
if REPLICATES > 1 or NUMJUM > 1:
    outfile_h = open('outfile', 'a')
    print >> outfile_h, " "
    print >> outfile_h, "-------------------------------------------"
    print >> outfile_h, "* * * * * * * * * *  C O N S E S U S   T R E E  * * * * * * * *"
    outfile_h.close()

    # FITCH requires a rooted tree. KITSCH requires an unrooted tree.
    if PROGRAM == 'fitch':
        subprocess.call(['consense.csh', 'treefile', OUTGROUP, 'y', 'outfile.consense', os.path.join(STARTDIR, TREEFILE)])
    else:
        subprocess.call(['consense.csh', 'treefile', OUTGROUP, 'n', 'outfile.consense', os.path.join(STARTDIR, TREEFILE)])

    outfile_consense_h = open('outfile.consense', 'a')
    print >> outfile_consense_h, "\n"
    print >> outfile_consense_h, ">>>> TO GENERATE BRANCH LENGTHS"
    if PROGRAM == 'fitch':
        print >> outfile_consense_h, ">>>> USE TREE FILE AS INPUT FOR fitch"
    else:
        print >> outfile_consense_h, ">>>> USE TREE FILE AS INPUT FOR kitsch"
    print >> outfile_consense_h, ">>>> USING THE USERTREE OPTION"
    outfile_consense_h.close()

    phylip.write_tree(os.path.join(STARTDIR, OUTFILE), True, False)
else:
    shutil.move('treefile', os.path.join(STARTDIR, TREEFILE))
    phylip.write_tree(os.path.join(STARTDIR, OUTFILE), False)

os.chdir('..')

shutil.rmtree(TEMPDIR)