#!/usr/bin/env python3
import argparse
from collections import defaultdict
from ete3 import Tree
import pandas as pd


def name_internal_nodes(tree: Tree) -> Tree:
    node_i = 0
    # Name the internal nodes with a number and the root with "Root"
    for n in tree.traverse():
        if not n.is_leaf() and n.name.strip() == "":
            n.name = f"node_{node_i}"
            node_i += 1
        if n.is_root():
            n.name = "Root"
            continue
        assert n.dist > 0.0
    return tree


def main(input_tree):

    # Convert newick to nexus including root
    tree = name_internal_nodes(Tree(input_tree, format=1))
    # Write as a .tsv file, one line per node
    dico_output = defaultdict(list)
    feature_list = sorted(set([f for node in tree.traverse() for f in node.features]) - {"support", "name"})
    for node in tree.traverse():
        dico_output["node_name"].append(node.name)
        for feature in feature_list:
            dico_output[feature].append(getattr(node, feature))

    df = pd.DataFrame(dico_output)
    df.to_csv(input_tree.replace(".nhx", ".tsv"), sep="\t", index=False)
    tree.write(format=1, outfile=input_tree.replace(".nhx", ".tree"))

if __name__ == '__main__':
    # Parse arguments
    descr = "A script to convert a New Hampshire X (.nhx) formatted tree (containing annotations) to a newick tree (topology) and a tsv file (annotations).\n"
    descr += "The script assumes that the input file is ending with '.nhx' and will replace it with '.tree' and '.tsv' in the output files.\n"
    descr += "The '.tree' file will contain the topology of the tree with branch lengths and internal node names.\n"
    descr += "The '.tsv' file will contain the annotations of the tree with one line per node.\n"
    descr += "This script is meant to process '.nhx' files generated by the option '--newick' of the program 'readnodeomega', 'readnodemutsel' and 'readnodetraits'.\n"
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, description=descr)
    parser.add_argument('-i', '--input', required=True, type=str, dest="input")
    args = parser.parse_args()
    main(args.input)