import toytree

"""
Explanation: Plotting IQtree concordance metrics (odd newick format)
End of explanation
"""


NEWICK = """
(MjavVW4_1:0.0029143962,(((((MjavVW4_2:0.0041175791,MareHarA_2
:0.0046787476)100/46.2/56.5:0.0027390681,(MfloSJF1_2:0.0036565039,
MincW1_2:0.0041099329)100/42.6/60.7:0.0024466057)100/46.9/70.6:
0.0024712752,((MentL30_2:0.0131877117,MentL30_1:0.0131141768)
100/61.2/64.9:0.0037704010,(Mhaplanaria_1:0.0193708158,
Mhaplanaria_2:0.0105167729)64/16.7/33.4:0.0010075264)100/28.6/
51.2:0.0005818881)100/34.5/65.9:0.0034515356,MfloSJF1_1:0.0077
086495)100/34.6/43:0.0034574478,MincW1_1:0.0042945317)100/46.1
/60.5:0.0022392730,MareHarA_1:0.0031851283);
""".replace("\n", "")

# parse tree and extract values on nodes (names)
tree = toytree.tree(NEWICK, tree_format=1)

"""
Explanation: The newick format below stores the concordance metrics as node annotations (names) and also includes branch lengths in the format (node)name:dist for internal nodes, and name:dist for tip nodes. The concordance metrics are stored as a string x/y/z where the three values are of interest.
End of explanation
"""


# extract names from internal nodes
tree.get_node_values("name")

# Example: you could extract each element manually
# split values on / to extract nth element 
val0 = [
    i.split("/")[0] if "/" in i else "" for i in 
    tree.get_node_values("name")
]

# show the parsed values
val0

"""
Explanation: Accessing features
The annotations in the newick string will be stored as features in the toytree object. If the newick string is in NHX format, like mrbayes produces, then the features will be stored with unique names, whereas if it is a standard newick string then the data are parsed to standard keys 'dist' and 'name'.
End of explanation
"""


# extract the 0th element
ctree = tree.set_node_values(
    feature="concordance",
    values={
        i: j.name.split("/")[0]
        for (i, j) in tree.idx_dict.items()
        if "/" in j.name
    }
)

ctree.draw(node_labels="concordance", node_sizes=18);

# extract the 1th element
ctree = ctree.set_node_values(
    feature="concordance-1",
    values={
        i: j.name.split("/")[1]
        for (i, j) in tree.idx_dict.items()
        if "/" in j.name
    }
)

ctree.draw(node_labels="concordance-1", node_sizes=21, height=350, width=400, ts='s');

# extract the 2th element
ctree = ctree.set_node_values(
    feature="concordance-2",
    values={
        i: j.name.split("/")[2]
        for (i, j) in tree.idx_dict.items()
        if "/" in j.name
    }
)

ctree.draw(node_labels="concordance-2", node_sizes=21, height=350, width=400, ts='s');

"""
Explanation: Set values as features on a tree
The .set_node_values function allows you to set node values on the tree object itself with new feature names. Here I set a new concordance feature (name it whatever you want) and set the first value from the concordance string to each node. Below I repeat for the other two values.
End of explanation
"""


