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

# <h1><span style="color:gray">ipyrad-analysis toolkit:</span> mrbayes</h1>
# 
# In these analyses our interest is primarily in inferring accurate branch lengths under a relaxed molecular clock model. This means that tips are forced to line up at the present (time) but that rates of substitutions are allowed to vary among branches to best explain the variation in the sequence data. 
# 
# There is a huge range of models that can be employed using mrbayes by employing different combinations of parameter settings, model definitions, and prior settings. The ipyrad-analysis tool here is intended to make it easy to run such jobs many times (e.g., distributed in parallel) once you have decided on your settings. In addition, we provide a number of pre-set models (e.g., clock_model=2) that may be useful for simple scenarios. 
# 
# Here we use simulations to demonstrate the accuracy of branch length estimation when sequences come from a single versus multiple distinct genealogies (e.g., gene tree vs species tree), and show an option to fix the topology to speed up analyses when your only interest is to estimate branch lengths. 

# In[2]:


# conda install ipyrad -c conda-forge -c bioconda
# conda install mrbayes -c conda-forge -c bioconda
# conda install ipcoal -c conda-forge


# In[3]:


import toytree
import ipcoal
import ipyrad.analysis as ipa


# ### Simulate a gene tree with 14 tips and MRCA of 1M generations

# In[4]:


TREE = toytree.rtree.bdtree(ntips=8, b=0.8, d=0.2, seed=123)
TREE = TREE.mod.node_scale_root_height(1e6)
TREE.draw(ts='o', layout='d', scalebar=True);


# ### Simulate sequences on single gene tree and write to NEXUS
# When Ne is greater the gene tree is more likely to deviate from the species tree topology and branch lengths. By setting recombination rate to 0 there will only be one true underlying genealogy for the gene tree. We set nsamples=2 because we want to simulate diploid individuals. 

# In[5]:


# init simulator
model = ipcoal.Model(TREE, Ne=2e4, nsamples=2, recomb=0)

# simulate sequence data on coalescent genealogies
model.sim_loci(nloci=1, nsites=20000)

# write results to database file
model.write_concat_to_nexus(name="mbtest-1", outdir='/tmp', diploid=True)


# In[6]:


# the simulated genealogy of haploid alleles
gene = model.df.genealogy[0]

# draw the genealogy
toytree.tree(gene).draw(ts='o', layout='d', scalebar=True);


# ### View an example locus
# This shows the 2 haploid samples simulated for each tip in the species tree.

# In[7]:


model.draw_seqview(idx=0, start=0, end=50);


# ### (1) Infer a tree under a relaxed molecular clock model

# In[12]:


# init the mb object
mb = ipa.mrbayes(
    data="/tmp/mbtest-1.nex",
    name="itest-1",
    workdir="/tmp",
    clock_model=2,
    constraints=TREE,
    ngen=int(1e6),
    nruns=2,
)

# modify a parameter
mb.params.clockratepr = "normal(0.01,0.005)"
mb.params.samplefreq = 5000

# summary of priors/params
print(mb.params)


# In[13]:


# start the run
mb.run(force=True)


# In[15]:


# load the inferred tree
mbtre = toytree.tree("/tmp/itest-1.nex.con.tre", 10)

# scale root node to 1e6
mbtre = mbtre.mod.node_scale_root_height(1e6)

# draw inferred tree 
c, a, m = mbtre.draw(ts='o', layout='d', scalebar=True);

# draw TRUE tree in orange on the same axes
TREE.draw(
    axes=a, 
    ts='o', 
    layout='d', 
    scalebar=True, 
    edge_colors="darkorange",
    node_sizes=0,
    fixed_order=mbtre.get_tip_labels(),
);


# In[16]:


# check convergence statistics
mb.convergence_stats


# ### (2) Concatenated sequences from a species tree
# Here we use concatenated sequence data from 100 loci where each represents one or more distinct genealogies. In addition, Ne is increased to 1e5, allowing for more genealogical variation. *We expect the accuracy of estimated edge lengths will decrease since we are not adequately modeling the genealogical variation when using concatenation.* Here I set the recombination rate *within* loci to be zero. There is free recombination among loci, however, since they are unlinked.

# In[17]:


# init simulator
model = ipcoal.Model(TREE, Ne=1e5, nsamples=2, recomb=0)

# simulate sequence data on coalescent genealogies
model.sim_loci(nloci=100, nsites=200)

# write results to database file
model.write_concat_to_nexus(name="mbtest-2", outdir='/tmp', diploid=True)


# In[25]:


# the simulated genealogies of haploid alleles
genes = model.df.genealogy[:4]

# draw the genealogies of the first four loci
toytree.mtree(genes).draw(ts='o', layout='r', height=250);


# In[26]:


# init the mb object
mb = ipa.mrbayes(
    data="/tmp/mbtest-2.nex",
    workdir="/tmp",
    name="itest-2",
    clock_model=2,
    constraints=TREE,
    ngen=int(1e6),
    nruns=2,
)

# summary of priors/params
print(mb.params)

# start the run
mb.run(force=True)


# In[27]:


# load the inferred tree
mbtre = toytree.tree("/tmp/itest-2.nex.con.tre", 10)

# scale root node from unitless to 1e6
mbtre = mbtre.mod.node_scale_root_height(1e6)

# draw inferred tree 
c, a, m = mbtre.draw(ts='o', layout='d', scalebar=True);

# draw true tree in orange on the same axes
TREE.draw(
    axes=a, 
    ts='o', 
    layout='d', 
    scalebar=True, 
    edge_colors="darkorange",
    node_sizes=0,
    fixed_order=mbtre.get_tip_labels(),
);


# In[28]:


mb.convergence_stats


# ### To see the NEXUS file (data, parameters, priors):

# In[30]:


mb.print_nexus_string()


# ### (3) Tree inference (not fixed topology) and plotting support values
# Here we will try to infer the topology from a concatenated data set (i.e., not set a constraint on the topology). I increased the ngen setting since the MCMC chain takes longer to converge when searching over topology space. Take note that the support values from mrbayes newick files are available in the "prob{percent}" feature, as shown below. 

# In[31]:


# init the mb object
mb = ipa.mrbayes(
    data="/tmp/mbtest-2.nex",
    name="itest-3",
    workdir="/tmp",
    clock_model=2,
    ngen=int(2e6),
    nruns=2,
)

# summary of priors/params
print(mb.params)

# start run
mb.run(force=True)


# The tree topology was correctly inferred 

# In[32]:


# load the inferred tree
mbtre = toytree.tree("/tmp/itest-3.nex.con.tre", 10)

# scale root node from unitless to 1e6
mbtre = mbtre.mod.node_scale_root_height(1e6)

# draw inferred tree 
c, a, m = mbtre.draw(
    layout='d',
    scalebar=True, 
    node_sizes=18, 
    node_labels="prob{percent}",
);


# The branch lengths are not very accurate in this case:

# In[33]:


# load the inferred tree
mbtre = toytree.tree("/tmp/itest-3.nex.con.tre", 10)

# scale root node from unitless to 1e6
mbtre = mbtre.mod.node_scale_root_height(1e6)

# draw inferred tree 
c, a, m = mbtre.draw(ts='o', layout='d', scalebar=True);

# draw true tree in orange on the same axes
TREE.draw(
    axes=a, 
    ts='o', 
    layout='d', 
    scalebar=True, 
    edge_colors="darkorange",
    node_sizes=0,
    fixed_order=mbtre.get_tip_labels(),
);