#!/usr/bin/env python
# coding: utf-8

# ### BPP parallelization without blocking

# In[1]:


import ipcoal
import toytree
import ipyrad.analysis as ipa
import ipyparallel as ipp


# ### Start an ipcluster instance
# 
# Here I assume that you already started an ipcluster instance in a terminal using the command below, or by starting engines in the Ipython Clusters tab in Jupyter. Remember that when you pull in new updates and restart your kernel you also need to restart your cluster instance. 
# 
# 
# ```bash
# ipcluster start --n=4
# ```

# In[2]:


# connect to a running client
ipyclient = ipp.Client()

# show number of engines
ipyclient.ids


# ### Simulate loci under a known scenario

# In[3]:


# make a random tree
tree = toytree.rtree.unittree(ntips=5, treeheight=5e5, seed=1243)
tree.draw(ts='p');

# simulate loci and write to HDF5
model = ipcoal.Model(tree, Ne=1e5, nsamples=4)
model.sim_loci(100, 500)
model.write_loci_to_hdf5(name="test", outdir="/tmp", diploid=True)


# ### Setup BPP

# In[4]:


# create an IMAP 
IMAP = {
    'r' + str(i): [j for j in model.alpha_ordered_names if int(j[1]) == i][:2] 
    for i in range(5)
}
IMAP


# In[5]:


# init bpp tool.
bpp1 = ipa.bpp(
    data="/tmp/test.seqs.hdf5",
    name="test1", 
    workdir="/tmp",
    guidetree=tree,
    imap=IMAP,
    maxloci=100,
    burnin=1000,
    nsample=5000,
)
bpp1.kwargs


# ### Submit BPP jobs to run on cluster (using `._run()`)

# In[6]:


# submit 2 jobs to ipyclient
bpp1._run(nreps=2, ipyclient=ipyclient, force=True, block=False, dry_run=False)


# ### Submit more jobs on the same ipyclient
# Here I use the `.copy()` function for convenience, but you could just create a new BPP object and call the `._run()` command with the same ipyclient object.

# In[7]:


# submit X other jobs to ipyclient (e.g., using diff job name)
bpp2 = bpp1.copy("test2")
bpp2._run(nreps=4, ipyclient=ipyclient, force=True, block=False, dry_run=False)


# ### The asynchronous job objects

# In[8]:


# see the jobs that are submitted
bpp1.asyncs


# In[9]:


bpp2.asyncs


# ### Block until jobs finish (or don't)

# In[11]:


# see outstanding jobs (optional, this does NOT BLOCK)
ipyclient.outstanding


# In[12]:


# BLOCK until all jobs on ipyclient are finished (returns True when done)
ipyclient.wait()


# ### Summarize results (WHEN FINISHED)

# In[13]:


res, mcmc = bpp1.summarize_results("00")
res


# In[14]:


res, mcmc = bpp2.summarize_results("00")
res