#!/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 get_dico_df(input_trees: str, include_internal_nodes: bool) -> dict[str, pd.DataFrame]:
    # Read the trees
    trees = []
    for line in open(input_trees):
        # Skip comments and empty lines
        if line.startswith("#") or line.strip() == "":
            continue
        # Read the tree and name the internal nodes
        trees.append(name_internal_nodes(Tree(line.strip(), format=1)))

    # Create a set of all node names and features
    node_names = set()
    features = set()
    for tree in trees:
        for node in tree.traverse():
            node_names.add(node.name)
            features.update(node.features)
    features = sorted(features - {"support", "name", "dist"})

    # For each feature, create a Dataframe.
    # Each Dataframe with has many rows as trees (points of MCMC), and as many columns as nodes
    dico_df = {}
    for feature in features:
        ddf = defaultdict(list)
        # For each tree, get the values for each node
        for tree in trees:
            for node in tree.traverse():
                # Skip the nodes that do not have the feature (e.g. the root for some features)
                if feature not in node.features:
                    continue
                # Skip if the node is not a leaf and we don't want to include internal nodes
                if node.is_leaf() or include_internal_nodes:
                    ddf[node.name].append(float(getattr(node, feature)))
        # Each Dataframe with has many rows as trees (points of MCMC), and as many columns as nodes
        dico_df[feature] = pd.DataFrame(ddf)
    return dico_df


def compute_stats(df: pd.DataFrame):
    # For each column (node), get the mean and std
    for col in df.columns:
        mean = df[col].mean()
        std = df[col].std()
        print(f"{col}: {mean:.4f} ± {std:.4f}")
    # Get a new column with the mean of across the all columns
    df["global"] = df.mean(axis=1)
    # Compute the 5 and 95 percentiles of the global column
    q5, q95 = df["global"].quantile([0.05, 0.95])
    return q5, q95


def main(input_trees, output_file):
    # get the dictionary of dataframes, including the internal nodes if needed otherwise only the leaves
    dico_df = get_dico_df(input_trees, include_internal_nodes=True)

    # Show an example by computing the 5 and 95 percentiles of a new column which is the average across all nodes, for each feature
    dico_output = defaultdict(list)
    # Add a column with the percentiles (5% and 95%)
    dico_output["Percentile"] = ["5%", "95%"]
    for feature, df in dico_df.items():

        print(f"Computing stats for {feature}:")
        # This is in the compute_stats function that we compute the percentiles
        q5, q95 = compute_stats(df)
        # Add the percentiles to the output dictionary
        dico_output[feature].extend([q5, q95])

    # Create a DataFrame from the dictionary
    output_df = pd.DataFrame(dico_output)

    # Write the DataFrame to a .tsv file
    if output_file is None:
        output_file = input_trees + ".tsv"
    output_df.to_csv(output_file, sep="\t", index=False)


if __name__ == '__main__':
    descr = "A script to parse a list New Hampshire X (.nhx) formatted tree (containing annotations) and compute statistics.\n"
    descr += "This script is meant to process '.trees' files generated by the option '--newick_trees' of the program 'readnodeomega', 'readnodemutsel' and 'readnodetraits'.\n"
    descr += "This script is meant to be used a base and to be modified to compute the statistics you want.\n"
    descr += "As an example, for each feature, it computes the mean and std at each node and the global (average across nodes) mean and std\n"
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, description=descr)
    parser.add_argument('-i', '--input', required=True, type=str, dest="input")
    parser.add_argument('-o', '--output', required=False, type=str, dest="output")
    args = parser.parse_args()
    main(args.input, args.output)
