import os
os.environ['R_HOME'] = '/home/cfriedline/lib64/R'
import rpy2.robjects as robjects
import random
import string
import tempfile
import dendropy
from multiprocessing import Pool

r = robjects.r
ape = r('library(ape)')

def create_tree(num_tips, type):
    """
    creates the taxa tree in R
    @param num_tips: number of taxa to create
    @param type: type for naming (e.g., 'taxa')
    @return: a dendropy Tree
    @rtype: dendropy.Tree
    """
    r = robjects.r
    robjects.globalenv['numtips'] = num_tips
    robjects.globalenv['treetype'] = type
    name = _get_random_string(20)
    if type == "T":
        r("%s = rtree(numtips, rooted=T, tip.label=paste(treetype, seq(1:(numtips)), sep=''))" % name)
    else:
        r("%s = rtree(numtips, rooted=F, tip.label=paste(treetype, seq(1:(numtips)), sep=''))" % name)
    tree = r[name]
    return ape_to_dendropy(tree)

def ape_to_dendropy(phylo):
    """
    converts an ape tree to dendropy tree
    @param phylo: ape instance from rpy2
    @return: a dendropy tree
    @rtype: dendropy.Tree
    """
    f = tempfile.NamedTemporaryFile()
    robjects.r['write.nexus'](phylo, file=f.name)
    tree = dendropy.Tree.get_from_path(f.name, "nexus")
    f.close()
    return tree

def _get_random_string(length):
    """
    gets a random string of letters/numbers, ensuring that it does not start with a
    number
    @param length: length of the string
    @return: the random string
    @rtype: string
    """
    s = ''.join(random.choice(string.letters + string.digits) for i in xrange(length))
    if not s[0] in string.letters:
        return _get_random_string(length)
    return s

def create_tree_mp(num_taxa):
    t = create_tree(num_taxa, "T")
    return t

def get_taxa_trees(num_trees, num_taxa):
    jobs_mp = []
    jobs = []
    res = []
    pool = Pool(num_trees)
    for i in xrange(num_trees):
        jobs_mp.append(pool.apply_async(create_tree_mp, (num_taxa,)))
        jobs.append(create_tree_mp(num_taxa))
    pool.close()
    pool.join()
    res.append(jobs)
    res.append([x.get() for x in jobs_mp])
    return res

trees = get_taxa_trees(5, 10)