import numpy as np
import toytree
import toyplot

"""
Explanation: Cookbook gallery
This chapter is simply for displaying beautiful and creative plots made using toytree and toyplot. If you have one of your own please reach out (or make a github pull request) to contribute it. You can use simulated data (see examples below) or show examples with real data. For simulated data please limit data generation to the use of numpy and pandas. If real data, please make sure that the trees you use are available in an archived location (reliable URL) so that the plot can be easily re-created.
End of explanation
"""


# generate a random tree and data
ntips = 20
rseed = 123456
np.random.seed(rseed)
rtre = toytree.rtree.unittree(ntips=ntips, seed=rseed)
randomdata = np.random.uniform(20, 200, ntips)

# set up a toyplot Canvas with 2 axes: (x1, x2, y1, y2)
canvas = toyplot.Canvas(width=375, height=350)
ax0 = canvas.cartesian(bounds=(50, 200, 50, 300), padding=15, ymin=0, ymax=20)
ax1 = canvas.cartesian(bounds=(225, 325, 50, 300), padding=15, ymin=0, ymax=20)

# add tree to first axes
rtre.draw(axes=ax0);
ax0.show = False

# plot the barplot on the second axes
# (y-axis is range 0-ntips);
# (x-axis is bar values transformed to be 0-1)
# baseline is the space between tipnames and bars
ax1.bars(
    np.arange(ntips),
    randomdata,
    along='y',
);

# style axes
ax1.show = True
ax1.y.show = False
ax1.x.ticks.show = True

"""
Explanation: 1. ToyTree + barplot
Aligning a tree with data is sometimes easier on one axis versus two. See the (#1) and (#2) for comparison. Here when plotting on one axis the tree coordinates which map to treeheight and the number of tips can be difficult to align with data (e.g., a barplot) since the data values may be much greater than the treeheight. This can be fixed by tranforming the data and the axis labels. The example on a two axes is a bit easier in this case.
End of explanation
"""


# generate a random tree and data
ntips = 20
rseed = 123456
rtre = toytree.rtree.unittree(ntips=ntips, seed=rseed)

# names of different lengths
names = ["".join(np.random.choice(list("abcd"), i + 1)) for i in range(ntips)]

# make a canvas and coords for two plots
canvas = toyplot.Canvas(width=600, height=350)
ax0 = canvas.cartesian(grid=(1, 2, 0), yshow=False)
ax1 = canvas.cartesian(grid=(1, 2, 1), yshow=False)

# plot the tree with its default spacing for tree and names
rtre.draw(tip_labels=names, axes=ax0);

# plot the tree on the second axis
rtre.draw(tip_labels=names, axes=ax1, shrink=10);

"""
Explanation: 2. Spacing tree vs. tip names
The ratio of tree to tipnames on a plot is automatically adjusted to try to fit the tip names depending on their font and size, but only to an extent before the tipnames are eventually cutoff. If you want to manually adjust this ratio by squeezing the tree to take up less space this can be done by using the shrink parameter, as demonstrated below.
In the plot below I show the x-axis tick marks to highlight where the data are located on the x, and where the domain is by default and when extended. You can see that in both cases the treeheight is between 0 and -1 on the x-axis. But in the latter we extend the max domain from 1 to 3 which better accomodates the really long tipnames. Of course you can also increase the width of the entire canvas as well to increase spacing.
End of explanation
"""


# generate a random tree
ntips = 20
rseed = 123
np.random.seed(rseed)
rtre = toytree.rtree.coaltree(ntips=ntips, seed=rseed)

# assign new feature 'ancstate' as the random integer  
rtre = rtre.set_node_values(
    feature="ancstate",
    values={i: np.random.randint(5, 15) for i in rtre.idx_dict},
)

# get values in node plot order
sizes = rtre.get_node_values("ancstate", True, True)

# use a boolean of whether 'ancstate' is >10 to set color
colors = [toytree.colors[0] if (i and int(i)>10) else toytree.colors[1] for i in sizes]

# draw tree with styles
rtre.draw(
    node_labels=False,
    node_sizes=sizes,
    node_colors=colors,
    node_style={"stroke": "black"}
);

"""
Explanation: 3. Node size/color from features
By default TreeNodes have a number of features associated with them (support, height, dist, idx, name) and these are often useful for styling nodes. You can also add custom features to nodes (see TreeNodes chapter). Here I set the size and color of nodes based on features of nodes in a random tree (the random node names).
End of explanation
"""


rtre.get_edge_values('idx')

rtre.get_node_values('idx', 1, 1)

# generate a random tree
ntips = 7
rseed = 12345
rtre = toytree.rtree.coaltree(ntips=ntips, seed=rseed)

# edge mapping 1: enter a dictionary mapping clade members to colors
ecolors = rtre.get_edge_values_mapped({
    ("r0", "r1", "r2"): toytree.colors[0],  # <- using tips to define a clade
    ("r3"): toytree.colors[1],              # <- using tips to define a clade 
    11: toytree.colors[2],                  # <- using node idx to define a clade
})

# edge mapping 2: map specific edges (here 3,6,10,11,12) to edge width value
elabels = rtre.get_edge_values('idx')
ewidths = [5.0 if i in (3, 6, 10, 11, 12) else 2.5 for i in elabels]

# draw tree with edge colors, edge_widths, and node idx labels shifted to edges
c, a, m = rtre.draw(
    width=200,
    edge_colors=ecolors, 
    edge_widths=ewidths,
    node_labels=rtre.get_node_values('idx', True, False),
    node_labels_style={
        "-toyplot-anchor-shift": "-10px",
        "baseline-shift": "5px",
        "font-size": "10px",
    },
);

"""
Explanation: 4. Variable edge colors and widths
The function .get_edge_values_from_dict() is the most convenient way to apply a style value to parts of the tree. It returns a list with the value (e.g., color or width) mapped to the correct index of the list to apply to the correct edge when entered as an argument to draw. This is convenient for applying styles to clades. If alternatively you want to apply style to individual edges it is best to use .get_edge_values() and use the 'idx' argument to return the index order in which edges are plotted. You can then create a list of edge values based on this order. Both examples are shown below, as well as a way of shifting node labels so they are arranged over edges. The 'idx' label of nodes is used to refer to edges subtending nodes.
End of explanation
"""


# generate a random tree
rtre = toytree.rtree.unittree(20, seed=12345)

# make the canvas and axes
canvas = toyplot.Canvas(width=250, height=400)
axes = canvas.cartesian()
axes.show = True

# draw a rectangle (x1, x2, y1, y2)
axes.rectangle(
    -0.75, 0.35, -0.5, 4.5, 
    opacity=0.25,
    color=toytree.colors[0],
)

# draw a rectangle (x1, x2, y1, y2)
axes.rectangle(
    -0.75, 0.35, 4.5, 8.5, 
    opacity=0.25,
    color=toytree.colors[1],
)

# add tree to the axes 
rtre.draw(axes=axes);

"""
Explanation: 5. Colored rectangles to highlight clades
The easiest way to add colored shapes to a plot is with the Toyplot .rectangle or .fill() functions of cartesian axes objects. For this you simply need to know the coordinates of the area that you wish to fill (See Coordinates for details on this). The example below draws two rectangles in the coordinate space and then adds a tree on top of these. You could make more complex polygon shapes using the fill function (see Toyplot docs). Remember you can use axes.show=True to see the axes coordinates if you need a reminder of how to set the x and y coordinates of the rectangles.
End of explanation
"""


# we'll use scipy.stats to get prob. density func. of normal dist
import scipy.stats as sc

# generate a random tree with N tips
ntips = 40
tre = toytree.rtree.baltree(ntips).mod.node_slider(seed=123)

# generate a distribution between -10 and 10 for each tip in the tree
points = np.linspace(-10, 10, 50)
dists = {}
for tip in tre.get_tip_labels():
    dists[tip] = sc.norm.pdf(points, loc=np.random.randint(-5, 5, 1), scale=2)

# set up canvas for two panel plot
canvas = toyplot.Canvas(width=300, height=400)

# add tree to canvas
ax0 = canvas.cartesian(bounds=(50, 180, 50, 350), ymin=0, ymax=ntips, padding=15)
tre.draw(axes=ax0, tip_labels=False)
ax0.show = False

# add histograms to canvas
ax1 = canvas.cartesian(bounds=(200, 275, 50, 350), ymin=0, ymax=ntips, padding=15)

# iterate from top to bottom (ntips to 0)
for tip in range(tre.ntips)[::-1]:
    
    # select a color for hist
    color = toytree.colors[int((tip) / 10)]
    
    # get tip name and get hist from dict
    tipname = tre.get_tip_labels()[tip]
    probs = dists[tipname]
    
    # fill histogram with slightly overlapping histograms
    ax1.fill(
        points, probs / probs.max() * 1.25,
        baseline=[tip] * len(points), 
        style={"fill": color, "stroke": "white", "stroke-width": 0.5},
        title=tipname,
    )
    
    # add horizontal line at base
    ax1.hlines(tip, opacity=0.5, color="grey", style={"stroke-width": 0.5})

# hide y axis, show x 
ax1.y.show = False
ax1.x.label.text = "Trait value"
ax1.x.ticks.show = True

"""
Explanation: 6. Plot histograms associated with tip trait values (ridge plot)
You can use the .hist() or .fill() functions of toytree to plot histograms. Here we will generate and plot a distribution of a data in order from top to bottom so that the histograms overlap in a "ridge plot" fashion. An analagous function in ggtree seems to have merited an entire publication: https://academic.oup.com/mbe/article/35/12/3041/5142656.
End of explanation
"""


# load tree with variable name lengths
tree = toytree.tree("https://eaton-lab.org/data/Cyathophora.tre")
tree = tree.root(wildcard="prz")

"""
Explanation: 7. Plot tree with matrix/heatmap
End of explanation
"""


# generate some random data for this columns
spdata = np.random.randint(low=1, high=10, size=(tree.ntips, 5))
spdata

# scale tree to be 2X length of number of matrix cols
ctree = tree.mod.node_scale_root_height(spdata.shape[1] * 2)

# get canvas and axes with tree plot
canvas, axes, mark = ctree.draw(
    width=500,
    height=300,
    tip_labels_align=True,
    tip_labels_style={"-toyplot-anchor-shift": "80px"}
);

# add n columns of data (here random data)
ncols = 5
xoffset = 1
for col in range(5):
    
    # select the column of data
    data = spdata[:, col]
    
    # plot the data column
    axes.scatterplot(
        np.repeat(col, tree.ntips) + xoffset, 
        np.arange(tree.ntips),
        marker='s',
        size=10,
        color=toytree.colors[col],           
        opacity=0.1 + data[::-1] / data.max(),
        title=data,
    );

# stretch domain to fit long tip names
axes.x.domain.max = 20

"""
Explanation: Method 1:
The simplest method is to plot the tree and markers on shared coordinate axes. To make it easy to space items on the x-axis I set the tree to be 2X the width of the data (matrix), which allows me to use units of x=1 to space items on the x-axis. Then I generate a canvas and axes by drawing a tree, as usual, and here I add the data as square scatterplot markers with different opacities to represent the (randomly generated) data. 
The only tricky thing here is that you need to use tip_labels_style to offset the x-location of the tre tip labels, and also to extend the x-axis max domain if the names are long to prevent them from getting cut off.
End of explanation
"""


# a random rectangular matrix
matrix = np.arange(tree.ntips * 5).reshape(tree.ntips, 5)
matrix.shape

# create a canvas
canvas = toyplot.Canvas(width=500, height=350);

# add tree 
axes = canvas.cartesian(bounds=(50, 150, 70, 250))
tree.draw(axes=axes, tip_labels=False, tip_labels_align=True)

# add matrix
table = canvas.table(
    rows=13,
    columns=5, 
    margin=0,
    bounds=(175, 250, 65, 255),
)

colormap = toyplot.color.brewer.map("BlueRed")

# apply a color to each cell in the table
for ridx in range(matrix.shape[0]):
    for cidx in range(matrix.shape[1]):
        cell = table.cells.cell[ridx, cidx]
        cell.style = {
            "fill": colormap.colors(matrix[ridx, cidx], 0, 100), 
        }

# style the gaps between cells
table.body.gaps.columns[:] = 3
table.body.gaps.rows[:] = 3 

# hide axes coordinates
axes.show = False

# create a canvas
canvas = toyplot.Canvas(width=500, height=350);

# add tree 
axes = canvas.cartesian(bounds=(50, 150, 70, 250))
tree.draw(axes=axes, tip_labels=False, tip_labels_align=True)

# add matrix
colormap = toyplot.color.brewer.map("BlueRed")
table = canvas.matrix(
    (matrix, colormap),
    bounds=(120, 300, 25, 295), 
    tshow=True,
    tlabel="Traits",
    lshow=False,
    rshow=True,
    margin=0,
    rlocator=toyplot.locator.Explicit(range(tree.ntips), tree.get_tip_labels()[::-1])
)

# hide axes coordinates
axes.show = False

"""
Explanation: Method 2:
Using both a matrix and cartesian axes in toyplot. The key to aligning the two is that matrices have a margin of 50px by default. There aren't as many options to style matrix cells as there are in the option above. Here I used the right-side matrix labels to add and align tip names.
End of explanation
"""


