#!/usr/bin/env python3

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

# Run treedist as a command
# Version  Feb. 25, 2020
# Synopsis: treedist.py 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 -----
h_MSGFILE.write('                   TREEDIST')
h_MSGFILE.write('')

# 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()


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

#----------------- generate keyboard input to send to program -----
comfile_h = open('TreedistComfile', 'w')
# Choose distance measure
if DISTTYPE == "s":                # Symmetric
    comfile_h.write('d\n')
# else:                            # Branch Score Distance (default - DISTTYPE = "b")
# DO NOTHING!

# Should tree be rooted?
if ROOTED == "y":
    comfile_h.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)
    comfile_h.write('2\n')
    comfile_h.write('a\n')
    if OUTSTYLE == "s":
        comfile_h.write(OUTSTYLE + '\n')
    else: # or OUTSTYLE == "v"
        comfile_h.write('v\n')
elif WHICHDIST == "p":                # Dist. between all possible pairs in one tree file
    comfile_h.write('2\n')
    comfile_h.write('p\n')
    if OUTSTYLE in ("f", "s"):
        comfile_h.write(OUTSTYLE + '\n')
    else: # or OUTSTYLE == "v"
        comfile_h.write('v\n')

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

#-------- 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
os.nice(10)
comfile_h = open('TreedistComfile', 'r')
p_TREEDIST = subprocess.Popen(['treedist'], stdin=comfile_h)
p_TREEDIST.wait()
comfile_h.close()
os.nice(0)

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