## imports
import numpy as np
import ipyrad as ip
import ipyrad.analysis.baba as baba
import ipyparallel as ipp

## print versions
print "  ipyrad v.{}".format(ip.__version__)
print "  ipyparallel v.{}".format(ipp.__version__)
print "  numpy v.{}".format(np.__version__)

"""
Explanation: (&#x1F4D7;) Cookbook: ipyrad.analysis.baba
This notebook demonstrates how to use the baba module in ipyrad to test for admixture and introgression. The code is written in Python and is best implemented in a Jupyter-notebook like this one. Analyses can be split up among many computing cores. Finally, there are some simple plotting methods to accompany tests which I show below. 
import packages
End of explanation
"""


## start 4 engines by running the commented line below in a separate bash terminal. 
##
## ipcluster start --n=4

## connect to client and print connection summary
ipyclient = ipp.Client()
print ip.cluster_info(client=ipyclient)

"""
Explanation: ipcluster setup to run parallel code
You can find more details about this in our [ipyparallel tutorial].
End of explanation
"""


## create a Tree class object with the default tree
tre = baba.Tree()
tre.show(width=400, height=300);

"""
Explanation: Get default (example) 12 taxon tree
The function  baba.Tree()  takes a  newick  string as an argument, or, if no value is passed, it constructs the default 12 taxon tree shown below. The  Tree  Class object holds a representation of the tree, which can be used for plotting, as well as for describing a demographic model for simulating data, as we'll show below.
End of explanation
"""


## pass a newick string to the Tree class object
tre = baba.Tree(newick="((a,b),c);")
tre.show(width=200, height=200);

"""
Explanation: Or get any arbitrary tree
End of explanation
"""


## store newick tree and a migration event list [source, sink, start, stop, rate]
newick = "(((a,b),c), (d,e));"
events = [['e', 'c', 0, 1, 1e-6], 
          ['3', '1', 1, 2, 1e-6]]

## add admixture events to tree
tre = baba.Tree(newick=newick, admix=events)
tre.show(width=250, height=250, yaxis=True);

"""
Explanation: get a tree with an admixture edge
Shown by an arrow pointing backwards from source to sink (indicating that geneflow should be entered as occurring backwards in time). This is how the program msprime will interpret it.
End of explanation
"""


for idx in xrange(tre.verts.shape[0]-1, -1, -1):
    print idx, tre.verts[idx-1]
    row = tre.verts[idx-1]
    
    ## get time of the next event
    if row[1] == 0:
        pass
    else:
        tau = row[1]
    


## when do coalescent events occur?
who = tre.verts[:, 1] > 0
print 'when', tre.verts[who, 1]
print 'who', np.where(who)[0]

## when do migration events occur
## ...

for time in when:
    ## does migration change during the next epoch


for node in tre.tree.traverse("postorder"):
    ## check for migration
    if node.name in [i[0] for i in events]:
        node
        #ms.MigrationRateChange(time=0, rate=m_C_B, matrix_index=(1, 2)),
        #ms.MigrationRateChange(time=Taus[1], rate=0),

"""
Explanation: simulate data on that tree
End of explanation
"""


