#!/usr/bin/env python

import re
import os
import os.path
import phylip
import shutil
import subprocess
import sys

# Run treedist as a command
# Version  2/27/2007
# Synopsis: treedist.csh intree disttype rooted whichdist outstyle outfile
if len(sys.argv) < 7:
    print("Usage: treedist.py  INFILE  DISTTYPE  ROOTED  OUTFILE  WHICHDIST OUTSTYLE")
    exit();

#Convert arguments to variables
INFILE      = sys.argv[1]
DISTTYPE    = sys.argv[2]
ROOTED      = sys.argv[3]
WHICHDIST   = sys.argv[4]
OUTSTYLE    = sys.argv[5]
OUTFILE     = sys.argv[6]

# Make a temporary directory in which to run the program
STARTDIR = os.getcwd()
TEMPDIR = 'TREEDIST.' + str(os.getpid())
os.mkdir(TEMPDIR)
shutil.copyfile(INFILE, os.path.join(TEMPDIR, 'intree'))
os.chdir(TEMPDIR)

h_MSGFILE = open("MSGFILE", "w")

#----------------- Generate a header to append to the beginning of the output file -----
print('                   TREEDIST', file=h_MSGFILE)
print('', file=h_MSGFILE)

# Check for branch lengths in first few lines of input file.
# If no branch lengths exist, we must use Symmetric 
BRLEN = 0
line_count = 0
h_intree = open("intree", "r")
for line in h_intree:
    if line_count == 10:
        break
    if re.search(".*\:[0-9]*\.[0-9]*", line):
        BRLEN = BRLEN + 1
h_intree.close()

#----------------- generate keyboard input to send to program -----
os.nice(10)
p_TREEDIST = subprocess.Popen(['treedist'], stdin=subprocess.PIPE)


if DISTTYPE == "b" and BRLEN == 0:
    DISTTYPE = "s"
    print('>>>>> Warning: No branch lengths.', file=h_MSGFILE)
    print('>>>>> Using Symmetric Difference.', file=h_MSGFILE)
    print('', file=h_MSGFILE)
h_MSGFILE.close()

# Choose distance measure
if DISTTYPE == "s":                # Symmetric
    p_TREEDIST.stdin.write('d\n')
# else:                            # Branch Score Distance (default - DISTTYPE = "b")
# DO NOTHING!

# Should tree be rooted?
if ROOTED == "y":
    p_TREEDIST.stdin.write('r\n')

# Set the style for the report written to outfile
# C and L not currently implemented.
if WHICHDIST == 'a':                  # Dist. between adjacent pairs in one tree file (default)
    p_TREEDIST.stdin.write('2\n')
    p_TREEDIST.stdin.write('a\n')
    if OUTSTYLE == "s":
        p_TREEDIST.stdin.write(OUTSTYLE + '\n')	
    else: # or OUTSTYLE == "v"
        p_TREEDIST.stdin.write('v\n')
elif WHICHDIST == "p":                # Dist. between all possible pairs in one tree file
    p_TREEDIST.stdin.write('2\n')
    p_TREEDIST.stdin.write('p\n')
    if OUTSTYLE in ("f", "s"):
        p_TREEDIST.stdin.write(OUTSTYLE + '\n')
    else: # or OUTSTYLE == "v"
        p_TREEDIST.stdin.write('v\n')

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

#-------- Run treedist, sending terminal output to /dev/null -----------
# high nice level is set just in case we do get an infinite loop
# Remember, treedist reads a file called intree, and writes to outfile
p_TREEDIST.wait()
os.nice(0)

phylip.merge_msg(os.path.join(STARTDIR, OUTFILE), False)
os.chdir(STARTDIR)
shutil.rmtree(TEMPDIR)