import toytree
import toyplot
import numpy as np

## create a random tree with N tips
tre = toytree.rtree.coaltree(24)

## assign random edge lengths and supports to each node
for node in tre.treenode.traverse():
    node.dist = np.random.exponential(1)
    node.support = int(np.random.uniform(50, 100))

"""
Explanation: Random trees
There are some simple tools in toytree, based on functions in ete3, that can be used to generate random trees which can be useful for exploring data and plotting methods.
End of explanation
"""


## make a colormap
colormap = toyplot.color.brewer.map(name="Spectral", domain_min=0, domain_max=100)

## get node support values in plot order 
vals = np.array(tre.get_node_values("support", show_root=1, show_tips=1))
vals

## get colors mapped to values
colors = toyplot.color.broadcast((vals, colormap), shape=vals.shape)
colors = [toyplot.color.to_css(i) for i in colors]

"""
Explanation: Color-mapping
End of explanation
"""


tre.draw(
    height=400, width=300,
    node_labels=tre.get_node_values(None, show_root=True, show_tips=False),
    node_colors=colors,
    node_sizes=10,
    node_style={"stroke": "black"},
    tip_labels_align=True,
    );

"""
Explanation: Show nodes but not values (text)
Use None as the argument to get_node_values(). This allows you to still toggle whether you want the root and tip nodes to show with the following arguments to get_node_values. This argument would commonly be used if you wanted to set different node sizes or colors to represent values instead of having text show at nodes.
End of explanation
"""


tre.draw(
    height=400, width=350,
    node_labels=tre.get_node_values("support", False, False),
    node_colors=colors,
    node_sizes=15,
    node_style={"stroke": "black"},
    use_edge_lengths=False,
    );

"""
Explanation: Show nodes and node values
Similar to above except now enter the name of a real node "feature" into the get_node_values functions. This will return the values of the feature in node plotting order. The features "support", "dist", "name", "height", and "idx" are always available. Additional features can be added to nodes.
End of explanation
"""


tre.draw(
    height=400, width=350,
    use_edge_lengths=False,
    node_labels=tre.get_node_values("support", False, False),
    node_colors=colors,
    node_sizes=0,
    node_style={"stroke": "black"},
    node_labels_style={
        "baseline-shift": "8px", 
        "-toyplot-anchor-shift": "-15px",
        "font-size": "10px"},
    );

"""
Explanation: Show node values but not the nodes themselves
End of explanation
"""


