#!/usr/bin/env python3

import os
import os.path
import shutil
import subprocess
import sys

# Version Feb. 25, 2020
#Synopsis: consense.py intree method fraction outgroup rooted outfile outtree

# Run consense as a command
# At present, this script has no way to check the number of
# "species" in the trees in intree. If OUTGROUP is greater
# than the number of sequences, consense will go into an
# infinite loop prompting for OUTGROUP. dnapars.csh does
# check the number of "species". Similar checks should exist
# in any procedure that calls this script.

""" ensure that there are enough command line arguments to parse """
import os
if len(sys.argv) < 8:
    print("Usage: consense.py  INFILE  METHOD  FRACTION  OUTGROUP  ROOTED  OUTFILE  TREEFILE")
    exit();

#Convert arguments to variables
INFILE      = sys.argv[1]
METHOD      = sys.argv[2]
FRACTION    = sys.argv[3]
OUTGROUP    = sys.argv[4]
ROOTED      = sys.argv[5]
OUTFILE     = sys.argv[6]
TREEFILE    = sys.argv[7]

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

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

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


#----------------- generate keyboard input to send to program -----
# Generate keyboard input for protdist and save it to ConsenseComfile.
# While it is possible to pipe the output directly to consense, we have
# discovered that on systems with a high load average, the program doesn't
# always complete. Besides, it is WAY easier to debug when we save keyboard
# output to a file that can be examined.
comfile_h = open('ConsenseComfile', 'w')

# Choose resampling method
if METHOD == "s":                # strict majority rule
    comfile_h.write('c\n')
elif METHOD == "m":                # Majority rule
    comfile_h.write('c\n')
    comfile_h.write('c\n')
elif METHOD == "l":                # Ml
    comfile_h.write('c\n')
    comfile_h.write('c\n')
    comfile_h.write('c\n')
#elif METHOD == "e":                # Majority rule (extended)
    # do nothing
#else:
    # do nothing

# Outgroup
if int(OUTGROUP) > 1:
    comfile_h.write('o\n')
    comfile_h.write(OUTGROUP + '\n')

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

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

# Enter fraction for M1 method. This can either
# be a real number, in which case it is directly
# used, or an integer between 50 and 100, which
# is converted to a real number.
if METHOD == "l":
    # consense.item generates an integer between 50 and 100,
    # which we can convert into a fraction
    FINT = int(FRACTION)
    if (50 < FINT < 100):
        FRACTION = str(float(FINT)/100.0)
    else:
        FRACTION = '0.5'
    comfile_h.write(FRACTION + '\n')
    print('FRACTION set to ' + FRACTION)
comfile_h.close()

# Run consense
comfile_h = open('ConsenseComfile', 'r')
p = subprocess.Popen(['consense'], stdin=comfile_h)
p.wait()
comfile_h.close()

#Copy the output files to the destination directory.
shutil.copyfile('outfile', os.path.join(STARTDIR, OUTFILE))
shutil.copyfile('outtree', os.path.join(STARTDIR, TREEFILE))
os.chdir(STARTDIR)
shutil.rmtree(TEMPDIR)
#os.nice(0)