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

rc = Client(profile='sge')
dview = rc[:]
lview = rc.load_balanced_view()
len(rc)

with dview.sync_imports():
    import os
    import rpy2
    import rpy2.robjects
    import random
    import string
    import tempfile
    import dendropy
    import socket
    import scipy
    from multiprocessing import Pool
    
def setup_cluster_engines():
    os.environ['R_HOME'] = '/home/cfriedline/lib64/R'
    r = rpy2.robjects.r
    ape = r('library(ape)')
    return socket.gethostname(), os.getpid(), os.environ['R_HOME'], rpy2.__version__
dview['setup_cluster_engines'] = setup_cluster_engines
dview.apply(setup_cluster_engines).get()

r = rpy2.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 = rpy2.robjects.r
    rpy2.robjects.globalenv['numtips'] = num_tips
    rpy2.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()
    rpy2.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
    """
    choices = "%s%s" % (string.letters,string.digits)
    s = ''.join(scipy.random.choice(list(choices),10))
    if s[0] not in string.letters:
        return _get_random_string(length)
    return s

tree = create_tree(100, "T")
tree.as_newick_string()[0:80]

def reseed(args):
    scipy.random.seed()
    return os.getpid()

def pool_reseed(pool, jobs):
    res = pool.map(reseed, range(jobs))
    if len(set(res)) != jobs:
        return pool_reseed(pool, jobs)
    return True, res

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_ip = []
    jobs = []
    res = []
    pool = Pool(num_trees)
    print pool_reseed(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))
        jobs_ip.append(lview.apply_async(create_tree_mp, num_taxa))
    pool.close()
    pool.join()
    res.append(jobs)
    res.append([x.get() for x in jobs_mp])
    res.append([x.get() for x in jobs_ip])
    return res

dview['create_tree'] = create_tree
dview['_get_random_string'] = _get_random_string
dview['ape_to_dendropy'] = ape_to_dendropy
dview['create_tree_mp'] = create_tree_mp
dview['get_taxa_trees'] = get_taxa_trees

# jobs = 10
# pool = Pool(jobs)
# pool_reseed(pool, jobs)
# if pool_reseed:
#     for i in xrange(10):
#         print pool.apply_async(_get_random_string, (20,)).get()
# pool.close()
# pool.join()

trees = get_taxa_trees(5, 10)

for t in trees[0]:
    print t.as_newick_string()[0:80]

for t in trees[1]:
    print t.as_newick_string()[0:80]

for t in trees[2]:
    print t.as_newick_string()[0:80]