#!/usr/bin/env python
import sys
import traceback
import time
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
    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 CipresIDL_api1.NexusException, x:
        if x.fileName:
            p =  os.path.abspath(x.fileName)
        else:
            p = "the input file"
        sys.exit("\nNexusError at line %d of %s:\n%s\n" % (x.lineNumber, p, x.errorMsg))
    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)
    tree = None
    try:
        tree = nex_blob.treeList[0]
    except:
        if require_tree:
            sys.exit('The file must contain a tree.\n'\
                     'The file "%s" does not appear to.' % fn)
    return tax, mat, tree

def run_file(fn, tiObjRef, event_consumer=None, use_start_tree=True):
    """Returns a (taxa, matrix, tree) tuple with the tree coming from running 
    a CIPRES tree-improver analysis on the taxa, matrix and first tree from the
    file `fn` using `app_name` as the tree improver.
    """
    if isinstance(event_consumer, EventConsumer):
        event_consumer = event_consumer.toIDL()
    tax, mat, tree = parse_file(fn, require_tree=use_start_tree)
    tiObjRef.setMatrix(mat)
    
    startTime = time.time()
    
    if use_start_tree:
        startTree = CipresTree(tree)
        tiObjRef.setTree(startTree)
        startTime = time.time()
        tList = tiObjRef.improveTreeReturnAll(event_consumer)
    else:
        tList = tiObjRef.inferTreeReturnAll(event_consumer)
    
    endTime = time.time()
    duration = endTime - startTime
    sys.stderr.write("Search took %f seconds\n" % duration)
    
    if not tList:
        sys.exit("No trees returned")
    ctList = []
    for i in tList:
        ctList.append(CipresTree(i))
    return tax, mat, ctList

if __name__ == '__main__':
    parser = OptionParser()
    parser.add_option("-a", "--app", dest="app_name", default="org.cipres.garli",
                      help="The tree improver service application name")
    parser.add_option("-c", "--command", dest="command", default="",
                      help="Execute command for service")
    parser.add_option("-x", "--xml", dest="xml", default="",
                      help="path to a file containing service commands in Cipres CommandObject XML",
                      metavar="FILE")
    parser.add_option("-s", "--status", dest="status", default="",
                      help="destination for status output from the tree service",
                      metavar="FILE")
    parser.add_option("-n", "--infer", dest="ignore_start_tree", 
                      default=False, action="store_true",
                      help="Do not use (or require) a  tree in the file as a starting tree (use the TreeInference interface instead of TreeImprove)",)
    (options, args) = parser.parse_args()
    if len(args) < 1:
        sys.exit("Expecting a data file as a command line argument")
        
    registry = getCipresRegistry()

    if options.status:
        status_out = open(options.status, "w")
    else:
        status_out = None
    use_start_tree = not options.ignore_start_tree
    for fn in args:
        commands = []
        tiObjRef = None
        #event_consumer = createEventConsumer(EventConsumer)
        event_consumer = ToNexusEventConsumer(buffered=False)
        assert event_consumer is not None
        event_consumer.setStream(status_out)
        try:
            if options.xml:
                cmdObj = readCmdObjFromXMLFile(options.xml)
                c = PIPRes.corba.api1.CipresIDL_api1.TreeImprove
                tiObjRef = getTopLevelObjectReference(cmdObj)
                if len(cmdObj.children) > 0:
                    r, s = tiObjRef.executeRecursive(cmdObj)
                    if not r:
                        m = "%s: Error executing command:\n%s\n" % (sys.argv[0], s)
                        sys.exit(m)
                    if s:
                        m = "%s: Command completed: %s\n" % (sys.argv[0], s)
                    else:
                        m = "%s: Succeeded\n" % sys.argv[0]
                else:
                    commands = cmdObj.commands
            else:
                app = options.app_name
                if not app:
                    m = "Service must be supplied using an --app or --xml argument"
                    sys.exit(m)
                cmd = options.command
                if use_start_tree:
                    interface='TreeImprove'
                    cls = CipresIDL_api1.TreeImprove
                else:
                    interface='TreeInfer'
                    cls = CipresIDL_api1.TreeInfer
                regEnt  = composeRegistryQuery(interface=interface, 
                                               applicationName=app)
                try:
                    tiObjRef = registry.getServiceOrThrow(cls, 
                                                      registryEntry=regEnt)
                except CipresIDL_api1.NoRegistryMatchException:
                    q = composeRegistryQuery(interface=interface)
                    rel = registry.find(q)
                    apps = [i.applicationName for i in rel]
                    if apps:
                        tail = "\nValid choices for this interface are:\n\t%s" % "\n\t".join(apps)
                    else:
                        tail = ""
                    sys.exit("Your CIPRES registry is not aware of an implementation of %s that is called %s%s" % (interface, app, tail))
                    
                                                      
                if options.command:
                    commands = [options.command]
            for cmd in commands:
                m = "%s: Executing command:\n%s\n" % (sys.argv[0], cmd)
                sys.stderr.write(m)
                r, s = tiObjRef.execute(cmd)
                if not r:
                    m = "%s: Error executing command:\n%s\n" % (sys.argv[0], s)
                    sys.exit(m)
                if s:
                    m = "%s: Command completed: %s\n" % (sys.argv[0], s)
                else:
                    m = "%s: Succeeded\n" % sys.argv[0]
                sys.stderr.write(m)
            tax, mat, tre = run_file(fn, 
                                     tiObjRef, 
                                     event_consumer, 
                                     use_start_tree=use_start_tree
                                     )
            writeTreesBlock(sys.stdout, tre, tax)
        finally:
            if not (tiObjRef is None):
                tiObjRef.remove()
            event_consumer.remove()
            event_consumer = None