#!/usr/bin/env python
import sys
import traceback
from optparse import OptionParser
from PIPRes.util.cmd_obj_xml import readCmdObjFromXMLFile, getTopLevelObjectReference
from PIPRes.event_consumer import EventConsumer, ToNexusEventConsumer

try:
    import PIPRes
    from PIPRes.basic import readFilePathAsNexus, writeTreesBlock, writeNexusDataMatrix
    from PIPRes.util.client_import import *
except:
    sys.exit("The importing of CIPRES python failed.\nMake sure that your environment is configured correctly.")

def parse_file(fn, require_tree=False):
    "Returns taxa, matrix, and trees from fn"
    if not os.path.exists(fn):
        sys.exit(fn + " not found.")
    try:
        nex_blob =  readFilePathAsNexus(fn)[0]
    except Exception, e:
        print >>sys.stderr, "Stack trace of error message follows:\n"
        traceback.print_exc(sys.stderr)
        sys.exit("\n\nError during NEXUS parsing")
    try:
        mat = nex_blob.matrixList[0]
        tax = nex_blob.taxa
    except:
        sys.exit('The file must contain a matrix.\n'\
                 'The file "%s" does not appear to.' % fn)
    if require_tree and len(nex_blob.treeList) == 0:
        sys.exit('The file must contain a tree.\n'\
                 'The file "%s" does not appear to.' % fn)
    return tax, mat, [CipresTree(i) for i in nex_blob.treeList]

if __name__ == '__main__':
    parser = OptionParser()
    parser.add_option("-t", "--tree_num", dest="tree_num", default=1, 
        type="int",
        help="The number of the tree to use to order the taxa")
    parser.add_option("-v", "--verbose", dest="verbose", default=False, 
        action="store_true",
        help="verbose output")
    (options, args) = parser.parse_args()
    if len(args) != 1:
        sys.exit("Expecting a data file as a command line argument")

    fn = args[0]
    registry = getCipresRegistry()
    
    tax, mat, tree_list = parse_file(fn, require_tree=True)


    try:
        tree = tree_list[options.tree_num - 1]
    except:
        sys.exit("Tree number %d not found in %s" % (options.tree_num, fn))

    if options.verbose:
        print >>sys.stderr, "#NEXUS"
        writeNexusDataMatrix(sys.stderr, mat, tax)
        writeTreesBlock(sys.stderr, tree_list, tax)


    new_mat = []
    new_tax = []
    for tip in tree.iterTips():
        new_mat.append(mat.m_matrix[tip.index])
        new_tax.append(tax[tip.index])
    mat.m_matrix = new_mat
    
    print "#NEXUS"
    writeNexusDataMatrix(sys.stdout, mat, new_tax)
    writeTreesBlock(sys.stdout, [CipresTree(i) for i in tree_list], tax)