'''Generates a tree from a prior distribution (currently uniform over all topologies
with an exponential branch length).'''
import sys
import random
from PIPRes.tree import Tree
from PIPRes.nexus.public_blocks import *
from PIPRes.util.io import cipresGetLogger
_LOG = cipresGetLogger('example.treeFromPrior')
_version = '0.02'
_options = [    ['-v', 'prints version'],
                ['-h', 'prints help'],
                ['-n', 'output in a NEXUS-ish format'],
                ['-l=<int>', '# of leaves (if no inputfile is specified) - default is 4'],
                ['-t=<int>', '# of trees to draw - default is 1'],
                ['-e=<float>', 'mean of the egde length prior (exponential distribution used for each edge).'],
                ['-s=<long>', 'seed for the random number generator'],
                ['<infile>', 'path to a NEXUS file with the taxa to use in a data block.'],
           ]


def getVersion():
    return 'treeFromPrior.py version %s (MTH)\n' % _version
def printVersion(outS):
    outS.write(getVersion())
def printHelp(outS):
    printVersion(outS)
    outS.write('Generates trees from a prior distribution (currently uniform over all topologies with an exponential branch length)\nOptions:\n    ')
    outS.write('\n    '.join(['%10s  %s' % (i[0],i[1]) for i in _options]))
    outS.write('\n\n')

def purePythonGetDataBlock(inFilename):
    matrices = []
    try:    
        NexusParsing.pushCurrentFile(inFilename)
        inF = open(inFilename, 'rU')
        NexusParsing.statusMessage('Reading %s ...' % inFilename)
        for b in  NexusBlockStream(inF, ALL_PUBLIC_BLOCKS, True):
            if isinstance(b, NexusCharactersBlock):
                matrices.append(b)
                break
        NexusParsing.popCurrentFile()
        inF.close()
    except NexusError, x:
        import os
        print >>sys.stderr,  x
        NexusParsing.statusStream = None
    return matrices

if __name__ == '__main__':
    filename = None
    version, help, nexus, nLeaves, nTrees, meanBrLen, seed = False, False, False, 4, 1, 1.0, None
    for arg in sys.argv[1:]:
        if arg[0] == '-':
            if arg == '-h':
                help = True
            elif arg == '-v':
                version = True
            elif arg == '-n':
                nexus = True
            elif arg.startswith('-l='):
                nLeaves = int(arg[3:])
            elif arg.startswith('-t='):
                nTrees = int(arg[3:])
            elif arg.startswith('-s='):
                seed = long(arg[3:])
            elif arg.startswith('-e='):
                meanBrLen = float(arg[3:])
            else:
                printHelp(sys.stderr)
                sys.exit('Unrecognized option: %s' % arg)
        else:
            if filename is not None:
                sys.exit('Only one NEXUS file can be specified (2 found %s and %s)' % (filename, arg))
            filename = arg
    if help:
        printHelp(sys.stdout)
        sys.exit(0)
    if version:
        printVersion(sys.stdout)
        sys.exit(0)

    if filename is None:
        tNames = map(lambda x: 't%d' % (x + 1), range(nLeaves))
    else:
        mats = purePythonGetDataBlock(filename)
        if not mats:
            sys.exit('No Data matrices found')
        mat = mats[0] # what to do if we see multiple?
        tNames = mat.getTaxLabels()
        if not (tNames):
            sys.exit('Error.  Could not get a taxon list from the data matrix')
        nLeaves = len(tNames)

    rng = random.Random()
    if seed is not None:
        rng.seed(seed)

    expLambda = 1.0/meanBrLen

    if nexus:
        print '#NEXUS\n'
        commentStart =  '[! random trees from %s\nInvocation:\n\tpython treeFromPrior.py -n -e=%f -t=%d' % (getVersion(), meanBrLen, nTrees)
        seedStatement = seed is not None and ('-s=%ld' % seed) or ''
        print commentStart, seedStatement,
        if filename is None:
            print '-l=%d\n]\n' % nLeaves
        else:
            print os.path.abspath(filename), '\n]'
        print 'BEGIN taxa;\n\tdimensions ntax = %d;\n\ttaxlabels %s;\nEND;' % (nLeaves, ' '.join(tNames))
        print "BEGIN trees;\n\ttranslate %s;\n" % (', '.join(['%d %s' % (i[0]+1, i[1]) for i in enumerate(tNames)]))

    for i in range(nTrees):
        t = Tree.randTopology(nLeaves, rng=rng)
        for edge in t.iterEdges():
            edge.end.edgeLength = rng.expovariate(expLambda)
        t.hasEdgeLengths = True
        s = str(t)
        if nexus:
            print 'tree Tree%d = [&U] %s;' % (i+1, s)
        else:
            print s
    if nexus:
        print 'END;'