#!/usr/bin/env python

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

#Version  5/15/2007
# Run distance programs as a command
#Synopsis: genericdist.csh matfile tconmeth power subrep global negbranch outgroup jumble jseed numjum\
#             termout printdata outfile treefile

#Convert arguments to variables
MATFILE       = sys.argv[1]
NUMMAT        = sys.argv[2]
TCONMETH      = sys.argv[3]
POWER         = sys.argv[4]
SUBREP        = sys.argv[5]
GLOBAL        = sys.argv[6]
NEGBRANCH     = sys.argv[7]
OUTGROUP      = sys.argv[8]
JUMBLE        = sys.argv[9]
JSEED         = sys.argv[10]
NUMJUM        = sys.argv[11]
TERMOUT       = sys.argv[12]
PRINTDATA     = sys.argv[13]
OUTFILE       = sys.argv[14]
TREEFILE      = sys.argv[15]

# Remember where we started
STARTDIR = os.getcwd()

# Make a temporary directory in which to run the program
TEMPDIR = 'GENERICDIST.' + os.getpid()
os.mkdir(TEMPDIR)
shutil.copyfile(MATFILE, os.path.join(TEMPDIR, 'infile'))
os.chdir(TEMPDIR)

#----------------- generate keyboard input to send to GENERICDIST -----

h_MSGFILE = open("MSGFILE", "w")
print >> h_MSGFILE, "Generic Distance Matrix Phylogeny Methods"
print >> h_MSGFILE, "---------------------  CONSTRUCTING TREE(S) ---------------------"
print >> h_MSGFILE, ""

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

# Choose method for constructing distance matrix
# w = weighbor, weighted Neighbor-Joining
# F = FITCH, Fitch-Margoliash method
# f = FITCH, Minimum evolution
# K = KITSCH, Fitch-Margoliash method
# k = KITSCH, Minimum evolution method
# N = Neighbor-joining
# U = Neighbor-joining
# default = FITCH, Fitch-Margoliash method
if TCONMETH in ("w"):
    PROGRAM = "weighbor"
    print >> h_MSGFILE, 'Please cite:'
    print >> h_MSGFILE, 'WEIGHBOR - Weighted Neighbor Joining.'
    print >> h_MSGFILE, 'William J. Bruno, Nicholas D. Socci, and Aaron L. Halpern'
    print >> h_MSGFILE, 'Weighted Neighbor Joining: A Likelihood-Based Approach to'
    print >> h_MSGFILE, 'Distance-Based Phylogeny Reconstruction,'
    print >> h_MSGFILE, 'Mol. Biol. Evol. 17 (1): 189-197 (2000).'
    print >> h_MSGFILE, " "


    #Read length of sequence alignment from INFILE
    h_infile = open("infile.temp", "r")
    SEQLENGTH = h_infile.readline().split()[1]
    h_infile.close()

    print "SEQLENGTH = " + SEQLENGTH

    time_elapsed = time.time()
    os.nice(8)
    subprocess.call(["weighbor", "-L", SEQLENGTH, "-b", "4", "-i", "infile", "-o", "outtree", "-v"], stdout=open(TERMOUT, "a"))
    shutil.move("weighbor.out", "outfile")
    time_elapsed = time.time() - time_elapsed

    h_smalloutfile = open("outfile", "a")
    print >> h_smalloutfile, "Execution times on " + socket.gethostname() + ": " + time_elapsed
    h_smalloutfile.close()
else:
    if TCONMETH in ("K", "k"):
        PROGRAM = "kitsch"
    elif TCONMETH in ("N", "U"):
        PROGRAM = "neighbor"
    else:
        PROGRAM = "fitch"

    #----------------- Run FITCH, KITCH or NEIGHBOR   -----
    termout_h = open(TERMOUT, 'a')
    os.nice(8)
    time_elapsed = time.time()
    p_TREE = subprocess.Popen([PROGRAM], stdin=subprocess.PIPE, stdout=termout_h)

    if TCONMETH in ("k", "f"):
        print p_TREE.stdin, "d"
    elif TCONMETH == "U":
        print p_TREE.stdin, "n"

    if METHOD in ("b", "d"):
        print p_TREE.stdin, "m"
        print p_TREE.stdin, "w"
        print p_TREE.stdin, REPLICATES
    elif METHOD in ("ps", "po", "pw"):
        print p_TREE.stdin, "m"
        print p_TREE.stdin, "d"
        print p_TREE.stdin, REPLICATES
    #case "n": DO NOTHING


    # Jumble - When multiple datasets are analyzed, DNAPARS automatically
    # jumbles, and prompts for a random number seed for jumbling. Othersise,
    # jumbling must be explicitly set.
    jumble_select (METHOD, JUMBLE, p_TREE.stdin, h_MSGFILE)

    # Subreplicates
    if SUBREP == 'n':
        print p_TREE.stdin, 's'

    # Global rearrangements
    if GLOBAL == 'y':
        print p_TREE.stdin, 'g'

    # Negative branch lengths
    if NEGBRANCH == 'y':
        print p_TREE.stdin, '-'

    # Outgroup
    OUTGROUP = phylip.do_outgroup(OUTGROUP, p_TREE.stdin)


    # Should sequence data be printed?
    if PRINTDATA == "y":
        print p_TREE.stdin, "1"

    # When resampling, turn off printing trees to outfile
    # METHOD = ""
    #if METHOD in ('b', 'd', 'ps', 'po', 'pw'):
    #    print p_TREE.stdin, "3"

    #accept current settings and do the analysis
    print p_TREE.stdin, "y"

    #----------------- Run FITCH, KITCH or NEIGHBOR   -----
    p_TREE.stdin.close()
    p_TREE.wait()
    termout_h.close()
    msgfile_h.close()
    time_elapsed = time_elapsed - time.time()
    os.nice(0)

#----------- Return results to calling directory----------------

shutil.move('outtree', os.path.join(STARTDIR, TREEFILE))
phylip.write_tree(os.path.join(STARTDIR, OUTFILE), False)

os.chdir(STARTDIR)
shutil.rmtree(TEMPDIR)

print PROGRAM + " completed."