from io import StringIO
from itertools import permutations, combinations
from Bio import Phylo
from Bio.Phylo.BaseTree import Tree, Clade
# import re
import copy
# import numpy
import sys
import operator
import math
from collections import OrderedDict
from yxmath import base_translate
from yxmath.set import jaccord_index
import matplotlib.path as mpath
import matplotlib.patches as mpatches
import networkx as nx


__author__ = 'Yuxing Xu'

"""
from Bio import Phylo
from toolbiox.lib.common.evolution.tree_operate import add_clade_name

tree_file = "/Users/yuxingxu/Downloads/SpeciesTreeAlignment.fa.raxml.bestTree"
tree = Phylo.read(tree_file, 'newick')
tree = add_clade_name(tree)
"""


# see a tree

def lookup_by_names(tree, order='preorder'):
    """
    order ({'preorder', 'postorder', 'level'}) - Tree traversal order: 'preorder' (default) is depth-first search, 'postorder' is DFS with child nodes preceding parents, and 'level' is breadth-first search.

    order: depth-first search (DFS), preorder from root to first leaf one clade by clade, and second leaf... ; postorder, same search but from leaf to root
    level: breadth-frist search (BFS), from root to leaf one level by level

    see: https://biopython.org/DIST/docs/api/Bio.Phylo.BaseTree.TreeMixin-class.html
    """
    names = OrderedDict()
    for clade in tree.find_clades(order=order):
        if clade.name:
            if clade.name in names:
                raise ValueError("Duplicate key: %s" % clade.name)
            names[clade.name] = clade
    return names


def get_newick_string(tree):
    return tree.format('newick')

# clades relationship


def get_ancestors(child_clade, tree):
    """
    return the list path of parent-child links from the tree root to the clade of choice (root will return null list)
    """
    node_path = tree.get_path(child_clade)
    if len(node_path) == 0:
        return []
    else:
        return [tree.root] + node_path[:-1]


def get_parent(child_clade, tree):
    """
    get most close parent node for child_clade in a tree (root will return null list)
    """
    if child_clade == tree.root:
        return []
    node_path = tree.get_path(child_clade)
    if len(node_path) == 1:
        return tree.root
    else:
        return node_path[-2]


def get_sister(clade, tree):
    """
    get sister clade list
    """

    parent_node = get_parent(clade, tree)
    return [i for i in parent_node.clades if i != clade]


def get_sons(parent_clade):
    """
    get most close son nodes for clade in a tree
    """
    return parent_clade.clades


def get_offspring(parent_clade):
    """
    get all offspring clade for a parent clade
    """
    offspring = []
    for i in parent_clade.find_clades(order='level'):
        if i == parent_clade:
            continue
        offspring.append(i)
    return offspring


def get_leaves(parent_clade, return_name=False):
    if return_name:
        return [leaf.name for leaf in parent_clade.get_terminals()]
    else:
        return [leaf for leaf in parent_clade.get_terminals()]


def get_intranode(parent_clade, return_name=False):
    if return_name:
        return [clade.name for clade in parent_clade.get_nonterminals()]
    else:
        return [clade for clade in parent_clade.get_nonterminals()]


def get_closest_leaf(clade, topology_only=False):
    """Returns node's closest descendant leaf and the distance to
    it.

    :argument False topology_only: If set to True, distance
      between nodes will be referred to the number of nodes
      between them. In other words, topological distance will be
      used instead of branch length distances.

    :return: A tuple containing the closest leaf referred to the
      current node and the distance to it.

    """
    min_dist = None
    min_node = None
    if clade.is_terminal():
        return clade, 0.0
    else:
        for ch in get_sons(clade):
            node, d = get_closest_leaf(ch, topology_only=topology_only)
            if topology_only:
                d += 1.0
            else:
                d += ch.branch_length
            if min_dist is None or d < min_dist:
                min_dist = d
                min_node = node
        return min_node, min_dist


def all_clades_parent(tree):
    parents = {}
    for clade in tree.find_clades(order='level'):
        for child in clade:
            parents[child] = clade
    return parents


def is_offspring_of(A, B, tree):
    """
    A is offspring of B
    """
    return B in get_ancestors(A, tree)


def is_ancestors_of(A, B, tree):
    """
    A is ancestors of B
    """
    return A in get_ancestors(B, tree)


def is_son_of(A, B):
    """
    A is son of B
    """
    return A in get_sons(B)


def is_parent_of(A, B, tree):
    """
    A is parent of B
    """
    return A in get_parent(B, tree)


def is_pass_node(clade, tree):
    if clade.is_terminal():
        return False
    elif clade == tree.root:
        return False
    else:
        clade_son = get_sons(clade)
        if len(clade_son) == 1:
            return True
        else:
            return False


def monophyly(leaf_set, tree, support_threshold=0.8, jaccard_threshold=0.8, exclude_leaf_set=None):
    """
    give a list of gene (gene_set) to check if this gene in a single clade in the tree
    :param tree:
    :param gene_set:
    :param support_threshold:
    :param jaccard_threshold:
    :return:
    """
    tree_old = copy.deepcopy(tree)
    pass_monophyly = []
    for leaf in tree_old.get_terminals():
        if not leaf.name in leaf_set:
            # unroot tree
            tree = copy.deepcopy(tree_old)
            tree = add_clade_name(tree)
            node_dict = lookup_by_names(tree)

            # reroot tree
            tree_new, node_dict_new = reroot_by_outgroup_clade(
                tree, node_dict, leaf.name)
            if not support_threshold is None:
                for clade_name in get_low_support_clade(tree_new, support_threshold):
                    tree_new.collapse(node_dict_new[clade_name])
            node_dict_new = lookup_by_names(tree_new)

            # check monophyly

            for clade_name in node_dict_new:
                clade_leaf_list = [
                    clade_leaf.name for clade_leaf in node_dict_new[clade_name].get_terminals()]
                if exclude_leaf_set:
                    clade_leaf_list_without_exclude = set(
                        clade_leaf_list) - set(exclude_leaf_set)
                else:
                    clade_leaf_list_without_exclude = set(clade_leaf_list)

                if jaccord_index(clade_leaf_list_without_exclude, leaf_set) >= jaccard_threshold:
                    pass_monophyly.append(clade_leaf_list)

    if len(pass_monophyly) > 0:
        min_monophyly_list = sorted(
            pass_monophyly, key=lambda x: jaccord_index(x, leaf_set), reverse=True)[0]
        return True, list(min_monophyly_list)
    else:
        return False, []


def get_MRCA(A, B, tree):
    """
    get most recent common ancestor MRCA
    """
    A_ancest = get_ancestors(A, tree) + [A]
    B_ancest = get_ancestors(B, tree) + [B]

    if len(set(A_ancest) & set(B_ancest)) == 0:
        return tree.root
    else:
        i = 0
        MRCA = None
        while i < len(A_ancest) and i < len(B_ancest) and A_ancest[i] == B_ancest[i]:
            MRCA = A_ancest[i]
            i = i + 1
        return MRCA


def get_MRCA_from_list(clade_list, tree):
    """
    get most recent common ancestor MRCA
    """
    ancest_list = []
    for i in clade_list:
        ancest_list.extend(get_ancestors(i, tree) + [i])

    A_ancest = get_ancestors(clade_list[0], tree) + [clade_list[0]]

    MRCA = A_ancest[0]
    for clade in A_ancest:
        if ancest_list.count(clade) == len(clade_list):
            MRCA = clade
        else:
            break

    return MRCA


def get_distance_to_root(A, tree):
    """
    get distance (sum of branch length) from given clade to root
    """
    ancest_list = get_ancestors(A, tree) + [A]
    distance = 0.0
    for i in ancest_list:
        distance = distance + i.branch_length
    return distance


def get_distance(A, B, tree):
    """
    get distance (sum of branch length) for two clade
    """

    MRCA = get_MRCA(A, B, tree)

    MRCA_to_root = get_distance_to_root(MRCA, tree)
    A_to_root = get_distance_to_root(A, tree)
    B_to_root = get_distance_to_root(B, tree)

    if A_to_root < MRCA_to_root or B_to_root < MRCA_to_root:
        raise ValueError("evolution distance small than ancestor, maybe BUG")

    return (A_to_root - MRCA_to_root) + (B_to_root - MRCA_to_root)


def get_path_to_root(A, tree):
    """
    get path number from A to root
    """
    if A == tree.root:
        return []
    else:
        ancest_list = get_ancestors(A, tree) + [A]
        return ancest_list


def get_path_step_to_root(A, tree):
    """
    get path number from A to root
    """
    if A == tree.root:
        return 0
    else:
        ancest_list = get_ancestors(A, tree) + [A]
        return len(ancest_list)


def get_path_step(A, B, tree):
    """
    get the path number between clades
    """
    MRCA = get_MRCA(A, B, tree)

    MRCA_to_root = get_path_step_to_root(MRCA, tree)
    A_to_root = get_path_step_to_root(A, tree)
    B_to_root = get_path_step_to_root(B, tree)

    if A_to_root < MRCA_to_root or B_to_root < MRCA_to_root:
        raise ValueError("evolution distance small than ancestor, maybe BUG")

    return (A_to_root - MRCA_to_root) + (B_to_root - MRCA_to_root)


# change treeadd_clade_name
def num_to_chr(num, digit):
    """
    when digit is 5:
    let 0 -> AAAAA
        1 -> AAAAB
    """

    base_list = base_translate(num, 26)

    output_list = []
    for i in [0]*(digit-len(base_list)) + base_list:
        output_list.append(chr(65+i))

    return ''.join(output_list)


def add_clade_name(tree, chr_name=False):
    used_name = []
    for clade in tree.find_clades():
        if clade.name:
            used_name.append(clade.name)

    num = 0
    for clade in tree.find_clades():
        if clade.name is None:
            if chr_name:
                new_id = "Node_%s" % num_to_chr(num, 3)
            else:
                new_id = "N%d" % num
            while new_id in used_name:
                num = num + 1
                if chr_name:
                    new_id = "Node_%s" % num_to_chr(num, 3)
                else:
                    new_id = "N%d" % num
            clade.name = new_id
            used_name.append(new_id)
    return tree


def clade_rename(tree, rename_dict):

    for clade in tree.find_clades():
        if clade.name in rename_dict:
            clade.name = rename_dict[clade.name]

    return tree


def equalCladeTopo(clade1, clade2):
    if clade1.is_terminal() and clade2.is_terminal():
        if clade1.name == clade2.name:
            return True
        else:
            return False
    elif (not clade1.is_terminal()) and (not clade2.is_terminal()):
        if len(clade1.clades) != len(clade2.clades):
            return False
        all_true_flag = True
        for sub_clade1 in clade1.clades:
            true_flag = False
            for sub_clade2 in clade2.clades:
                if equalCladeTopo(sub_clade1, sub_clade2):
                    true_flag = True
            if not true_flag:
                all_true_flag = False
        return all_true_flag
    else:
        return False


def get_low_support_clade(tree, threshold=0.8):
    bad_support = []
    for clade in tree.get_nonterminals():
        if clade.confidence:
            if clade.confidence < threshold:
                bad_support.append(clade.name)
    return bad_support


def collapse_clade_list(tree, name_dict, want_to_collapse_list):
    for j in want_to_collapse_list:
        tree.collapse(name_dict[j])
    return tree


def collapse_low_support(tree, threshold=0.8):
    tree = add_clade_name(tree)
    tree_name_dir = lookup_by_names(tree)
    low_support_clade_list = get_low_support_clade(tree, threshold)
    for j in low_support_clade_list:
        try:
            tree.collapse(tree_name_dir[j])
        except:
            pass
    return tree


def ignore_branch_length(tree, to_None=False):
    """
    remove length to 1.0
    """
    tree = add_clade_name(tree)
    tree_name_dir = lookup_by_names(tree)
    for clade_name in tree_name_dir:
        if to_None:
            tree_name_dir[clade_name].branch_length = None
        else:
            tree_name_dir[clade_name].branch_length = 1.0
    return tree


def remove_given_node(clade, remove_name_list):
    if clade.name in remove_name_list:
        return None
    else:
        if clade.is_terminal():
            return clade
        else:
            son_clade_list = []
            for i in clade.clades:
                filtered_clade = remove_given_node(i, remove_name_list)
                if filtered_clade:
                    son_clade_list.append(filtered_clade)

            if len(son_clade_list) == 0:
                return None
            else:
                clade.clades = son_clade_list
                return clade


def remove_given_node_from_tree(tree, remove_name_list):
    tree_clade = copy.deepcopy(tree.clade)
    clade = remove_given_node(tree_clade, remove_name_list)

    new_tree = Tree(clade, tree.rooted)
    new_tree = remove_pass_node(new_tree)

    return new_tree


def remove_all_node_name(clade):
    clade_now = copy.deepcopy(clade)
    if clade_now.is_terminal():
        return clade_now
    else:
        clade_now.name = None
        new_sub_clade = []
        for i in clade_now.clades:
            new_sub_clade.append(remove_all_node_name(i))
        clade_now.clades = new_sub_clade
        return clade_now


def remove_tree_node_name(tree):
    tree_clade = copy.deepcopy(tree.clade)
    clade = remove_all_node_name(tree_clade)
    new_tree = Tree(clade, tree.rooted)

    return new_tree


def remove_pass_node(tree):
    iter_flag = True
    new_tree = copy.deepcopy(tree)

    while iter_flag:
        tree_tmp = copy.deepcopy(new_tree)
        tree_tmp = add_clade_name(tree_tmp)
        tree_tmp_name_dir = lookup_by_names(tree_tmp, order='level')

        # new_tree_root = tree_name_dir[list(tree_name_dir.keys())[0]]

        # build give son to find parent
        son_to_parent_hash = {}
        for clade_name in tree_tmp_name_dir:
            for son_clade in tree_tmp_name_dir[clade_name].clades:
                if son_clade.name in son_to_parent_hash:
                    raise ValueError("same son for different parent")
                son_to_parent_hash[son_clade.name] = clade_name

        # fund pass node
        pass_node_list = [clade_name for clade_name in tree_tmp_name_dir if
                          is_pass_node(tree_tmp_name_dir[clade_name], tree_tmp)]

        if len(pass_node_list) != 0:
            # for pass node give son to parent
            for pass_node_name in pass_node_list:
                parent_name = son_to_parent_hash[pass_node_name]
                if len(tree_tmp_name_dir[pass_node_name].clades) != 1:
                    raise ValueError("find a error pass node")
                son_name = tree_tmp_name_dir[pass_node_name].clades[0].name
                son_clade = tree_tmp_name_dir[son_name]
                if hasattr(tree_tmp_name_dir[pass_node_name], 'branch_length') and not (tree_tmp_name_dir[pass_node_name].branch_length is None):
                    son_clade.branch_length += tree_tmp_name_dir[pass_node_name].branch_length

                brother_list = tree_tmp_name_dir[parent_name].clades
                new_brother_list = []
                for i in brother_list:
                    if i.name == pass_node_name:
                        new_brother_list.append(son_clade)
                    else:
                        new_brother_list.append(tree_tmp_name_dir[i.name])

                tree_tmp_name_dir[parent_name].clades = new_brother_list

            # build new tree
            root_id = list(tree_tmp_name_dir.keys())[0]
            new_tree = Tree(
                root=tree_tmp_name_dir[root_id], rooted=tree_tmp.rooted)

        else:
            iter_flag = False

    return new_tree


def make_unique_name(tree):
    tree_new = copy.deepcopy(tree)
    names = {}
    for clade in tree_new.find_clades():
        if clade.name:
            if clade.name in names:
                names[clade.name].append(clade.name)
                clade.name = "%s.%d" % (clade.name, len(names[clade.name]))
            names[clade.name] = [clade.name]
    return tree_new


def mid_unrooted_tree(tree):
    """
    root a unrooted tree at midpoint
    """
    tree.root_at_midpoint()
    tree_new = add_clade_name(tree)
    node_dict = lookup_by_names(tree_new)
    return tree_new, node_dict


def reroot_by_outgroup_clade(tree, node_dict, outgroups_clade_name, keep_old=False):
    outgroup_clade = node_dict[outgroups_clade_name]
    outgroup_clade_old = copy.deepcopy(outgroup_clade)

    if keep_old:
        tree_old = copy.deepcopy(tree)

    tree.root_with_outgroup(outgroup_clade)
    tree_new = add_clade_name(tree)

    sis = [i for i in tree_new.clade.clades if i not in outgroup_clade_old.clades][0]

    new_outgroup_clade = copy.deepcopy(outgroup_clade_old)

    if sis.confidence is None:
        sis.confidence = 100
    new_outgroup_clade.confidence = sis.confidence

    root_clade = Clade(clades=[sis, new_outgroup_clade], name='root')
    root_clade.confidence = sis.confidence

    tree_new = Tree(root=root_clade, rooted=True)
    node_dict_new = lookup_by_names(tree_new)

    if keep_old:
        tree_old = add_clade_name(tree_old)
        node_dict_old = lookup_by_names(tree_old)
        return tree_new, node_dict_new, tree_old, node_dict_old
    else:
        return tree_new, node_dict_new


def reroot_by_outgroup_seqid(tree, outgroups_seq_list):
    tree_old = copy.deepcopy(tree)
    for leaf in tree_old.get_terminals():
        if not leaf.name in outgroups_seq_list:
            # unroot tree
            tree = copy.deepcopy(tree_old)
            tree = add_clade_name(tree)
            node_dict = lookup_by_names(tree)

            tree_new, node_dict_new = reroot_by_outgroup_clade(
                tree, node_dict, leaf.name)
            for clade_name in node_dict_new:
                clade_leaf_list = [
                    clade_leaf.name for clade_leaf in node_dict_new[clade_name].get_terminals()]
                if set(clade_leaf_list) == set(outgroups_seq_list):
                    tree_out, node_dict_out = reroot_by_outgroup_clade(
                        tree_new, node_dict_new, clade_name)
                    return tree_out, node_dict_out
    return None


# root a tree rewrite from orthofinder
# need more test

class RootMap(object):
    def __init__(self, setA, setB, GeneToSpecies):
        self.setA = setA
        self.setB = setB
        self.GeneToSpecies = GeneToSpecies

    def GeneMap(self, gene_name):
        sp = self.GeneToSpecies(gene_name)
        if sp in self.setA:
            return True
        elif sp in self.setB:
            return False
        else:
            print(gene_name)
            print(sp)
            raise Exception


def StoreSpeciesSets(t, GeneMap, tag="sp_"):
    tag_up = tag + "up"
    tag_down = tag + "down"
    for clade in t.find_clades(order='postorder'):
        if clade.is_terminal():
            setattr(clade, tag_down, {GeneMap(clade.name)})
        elif clade == t.root:
            continue
        else:
            setattr(clade, tag_down, set.union(
                *[ch.__getattribute__(tag_down) for ch in get_sons(clade)]))

    for clade in t.find_clades(order='preorder'):
        if clade == t.root:
            setattr(clade, tag_up, set())
        else:
            parent = get_parent(clade, t)
            if parent == t.root:
                others = [ch for ch in get_sons(parent) if ch != clade]
                setattr(clade, tag_up, set.union(
                    *[other.__getattribute__(tag_down) for other in others]))
            else:
                others = [ch for ch in get_sons(parent) if ch != clade]
                sp_downs = set.union(
                    *[other.__getattribute__(tag_down) for other in others])
                setattr(clade, tag_up, parent.__getattribute__(
                    tag_up).union(sp_downs))
    # setattr(t, tag_down, set.union(*[ch.__getattribute__(tag_down) for ch in get_sons(t)]))


def OutgroupIngroupSeparationScore(sp_up, sp_down, sett1, sett2, N_recip, n1, n2):
    f_dup = len(sp_up.intersection(sett1)) * len(sp_up.intersection(sett2)) * len(sp_down.intersection(sett1)) * len(
        sp_down.intersection(sett2)) * N_recip
    f_a = len(sp_up.intersection(sett1)) * (n2 - len(sp_up.intersection(sett2))) * (
        n1 - len(sp_down.intersection(sett1))) * len(sp_down.intersection(sett2)) * N_recip
    f_b = (n1 - len(sp_up.intersection(sett1))) * len(sp_up.intersection(sett2)) * len(sp_down.intersection(sett1)) * (
        n2 - len(sp_down.intersection(sett2))) * N_recip
    choice = (f_dup, f_a, f_b)
    #    print(choice)
    return max(choice)


def map_to_method(gene_to_species_map):
    def f(gene_id):
        return gene_to_species_map[gene_id]

    return f


def get_roots_by_species(gene_tree, species_rooted_tree, gene_to_species_map):
    """
    rewrite from orthofinder trees2ologs_of.py GetRoots
    """
    # reads tree
    gene_tree = add_clade_name(gene_tree)
    gene_tree_node_dict = lookup_by_names(gene_tree)
    species_rooted_tree = add_clade_name(species_rooted_tree)
    species_rooted_tree_node_dict = lookup_by_names(species_rooted_tree)
    GeneToSpecies = map_to_method(gene_to_species_map)

    speciesObserved = set([GeneToSpecies(leaf_name)
                           for leaf_name in get_leaves(gene_tree, True)])

    if len(speciesObserved) == 1:
        # arbitrary root if all genes are from the same species
        return [next(n for n in gene_tree.clade)]

    # use species tree to find correct outgroup according to what species are present in the gene tree
    n = species_rooted_tree.clade
    children = get_sons(n)
    leaves = [set(get_leaves(ch, True)) for ch in children]
    # make sure two group have species in gene tree
    have = [len(l.intersection(speciesObserved)) != 0 for l in leaves]
    while sum(have) < 2:
        n = children[have.index(True)]
        children = get_sons(n)
        leaves = [set(get_leaves(ch, True)) for ch in children]
        have = [len(l.intersection(speciesObserved)) != 0 for l in leaves]

    # Get splits to look for
    roots_list = []
    scores_list = []  # the fraction completeness of the two clades
    #    roots_set = set()
    for i in range(0, len(leaves)):
        t1 = leaves[i]
        t2 = set.union(*[l for j, l in enumerate(leaves) if j != i])
        # G - set of species in gene tree
        # First relevant split in species tree is (A,B), such that A \cap G \neq \emptyset and A \cap G \neq \emptyset
        # label all nodes in gene tree according the whether subsets of A, B or both lie below node
        StoreSpeciesSets(gene_tree, GeneToSpecies)  # sets of species
        root_mapper = RootMap(t1, t2, GeneToSpecies)
        sett1 = set(t1)
        sett2 = set(t2)
        nt1 = float(len(t1))
        nt2 = float(len(t2))
        N_recip = 1. / (nt1 * nt1 * nt2 * nt2)
        GeneMap = root_mapper.GeneMap
        # ingroup/outgroup identification
        StoreSpeciesSets(gene_tree, GeneMap, "inout_")
        # find all possible locations in the gene tree at which the root should be

        T = {True, }
        F = {False, }
        TF = set([True, False])
        for m in gene_tree.find_clades(order='postorder'):
            if m.is_terminal():
                if len(m.inout_up) == 1 and m.inout_up != m.inout_down:
                    # this is the unique root
                    return [m]
                    # print(m)
            else:
                if len(m.inout_up) == 1 and len(m.inout_down) == 1 and m.inout_up != m.inout_down:
                    # this is the unique root
                    return [m]
                    # print(m)
                nodes = get_sons(m) if m == gene_tree.root else [
                    m] + get_sons(m)
                clades = [ch.inout_down for ch in nodes] if m == gene_tree.root else (
                    [m.inout_up] + [ch.inout_down for ch in get_sons(m)])
                # do we have the situation A | B or (A,B),S?
                if len(nodes) == 3:
                    if all([len(c) == 1 for c in clades]) and T in clades and F in clades:
                        # unique root
                        if clades.count(T) == 1:
                            return [nodes[clades.index(T)]]
                            # print([nodes[clades.index(T)]])
                        else:
                            return [nodes[clades.index(F)]]
                            # print([nodes[clades.index(F)]])
                    elif T in clades and F in clades:
                        # AB-(A,B) or B-(AB,A)
                        ab = [c == TF for c in clades]
                        i = ab.index(True)
                        roots_list.append(nodes[i])
                        sp_down = nodes[i].sp_down
                        sp_up = nodes[i].sp_up
                        #                        print(m)
                        scores_list.append(
                            OutgroupIngroupSeparationScore(sp_up, sp_down, sett1, sett2, N_recip, nt1, nt2))
                    elif clades.count(TF) >= 2:
                        # (A,A,A)-excluded, (A,A,AB)-ignore as want A to be bigest without including B, (A,AB,AB), (AB,AB,AB)
                        i = 0
                        roots_list.append(nodes[i])
                        sp_down = nodes[i].sp_down
                        sp_up = nodes[i].sp_up
                        #                        print(m)
                        scores_list.append(
                            OutgroupIngroupSeparationScore(sp_up, sp_down, sett1, sett2, N_recip, nt1, nt2))
                elif T in clades and F in clades:
                    roots_list.append(m)
                    scores_list.append(0)  # last choice
    # If we haven't found a unique root then use the scores for completeness of ingroup/outgroup to root
    if len(roots_list) == 0:
        return []  # This shouldn't occur

    return [sorted(zip(scores_list, roots_list), key=lambda x: x[0], reverse=True)[0][1]]


def get_root_by_species(gene_tree, species_rooted_tree, gene_to_species_map):
    roots = get_roots_by_species(
        gene_tree, species_rooted_tree, gene_to_species_map)
    if len(roots) > 0:
        root_dists = [get_closest_leaf(r)[1] for r in roots]
        i, _ = max(enumerate(root_dists), key=operator.itemgetter(1))
        return roots[i]
    else:
        return None  # single species tree


def get_rooted_tree_by_species_tree(unrooted_gene_tree, species_tree, gene_to_species_map_dict):
    unrooted_gene_tree = add_clade_name(unrooted_gene_tree)
    unrooted_gene_tree_node_dict = lookup_by_names(unrooted_gene_tree)

    best_root_clade = get_root_by_species(
        unrooted_gene_tree, species_tree, gene_to_species_map_dict)

    gene_tree_rooted, gene_tree_rooted_node_dict, gene_tree, gene_tree_node_dict = reroot_by_outgroup_clade(unrooted_gene_tree,
                                                                                                            unrooted_gene_tree_node_dict,
                                                                                                            best_root_clade.name,
                                                                                                            True)

    return gene_tree_rooted


def get_root_by_species_for_file(unrooted_tree_file, species_tree_file, gene_to_species_map_dict, output_tree_file):
    species_tree = Phylo.read(species_tree_file, 'newick')
    unrooted_gene_tree = Phylo.read(unrooted_tree_file, 'newick')
    unrooted_gene_tree = add_clade_name(unrooted_gene_tree)
    unrooted_gene_tree_node_dict = lookup_by_names(unrooted_gene_tree)

    best_root_clade = get_root_by_species(
        unrooted_gene_tree, species_tree, gene_to_species_map_dict)

    gene_tree_rooted, gene_tree_rooted_node_dict, gene_tree, gene_tree_node_dict = reroot_by_outgroup_clade(unrooted_gene_tree,
                                                                                                            unrooted_gene_tree_node_dict,
                                                                                                            best_root_clade.name,
                                                                                                            True)

    for i in gene_tree_rooted_node_dict:
        c = gene_tree_rooted_node_dict[i]
        if not c.is_terminal():
            c.name = None

    Phylo.write(gene_tree_rooted, output_tree_file, format='newick')

# for nhx format


def map_node_species_info(rooted_gene_tree, species_tree, GeneToSpecies_dict):
    species_tree_node_dict = lookup_by_names(species_tree)

    for clade in rooted_gene_tree.find_clades(order='preorder'):
        if clade.is_terminal():
            clade.taxon = GeneToSpecies_dict[clade.name]
        else:
            sp_list = list(set([GeneToSpecies_dict[i.name]
                                for i in clade.get_terminals()]))
            sp_clade_list = [species_tree_node_dict[j] for j in sp_list]
            mrca_name = get_MRCA_from_list(sp_clade_list, species_tree).name
            clade.taxon = mrca_name

            num = 0
            sub_clades = {}
            for i in clade.clades:
                sub_clades[num] = i
                num += 1

            duplication = False
            speciation = True
            for a, b in combinations(list(sub_clades.keys()), 2):
                s1 = sub_clades[a]
                s2 = sub_clades[b]

                s1_sp_list = list(set([GeneToSpecies_dict[i.name]
                                       for i in s1.get_terminals()]))
                s2_sp_list = list(set([GeneToSpecies_dict[i.name]
                                       for i in s2.get_terminals()]))

                s1_mrca = get_MRCA_from_list(
                    [species_tree_node_dict[i] for i in s1_sp_list], species_tree)
                s2_mrca = get_MRCA_from_list(
                    [species_tree_node_dict[i] for i in s2_sp_list], species_tree)

                if is_ancestors_of(s1_mrca, s2_mrca, species_tree) or is_ancestors_of(s2_mrca, s1_mrca, species_tree) or s1_mrca.name == s2_mrca.name:
                    duplication = True
                    speciation = False
                    break

            clade.duplication = duplication
            clade.speciation = speciation

    return rooted_gene_tree


def get_nhx_comment(gene_tree):
    for clade in gene_tree.find_clades(order='preorder'):
        if clade.is_terminal():
            clade.comment = "&&NHX:S=%s" % clade.taxon
        else:
            if clade.confidence:
                if isinstance(clade.confidence, float):
                    confidence = min(100, int(clade.confidence * 100))
                else:
                    confidence = clade.confidence
            else:
                confidence = 0
            clade.confidence = None

            clade.comment = "&&NHX:S=%s:D=%s:B=%d" % (
                clade.taxon, 'Y' if clade.duplication else 'N', confidence)

    return gene_tree


def get_nhx_tree_with_taxon_info(gene_tree, species_tree_rooted, GeneToSpecies_dict, output_file):
    gene_tree = add_clade_name(gene_tree)
    species_tree_rooted = add_clade_name(species_tree_rooted)

    gene_tree = map_node_species_info(
        gene_tree, species_tree_rooted, GeneToSpecies_dict)
    nhx_tree = get_nhx_comment(gene_tree)

    with open(output_file, 'w') as f:
        Phylo.write(nhx_tree, f, 'newick')


def nhx_comment_parse(comment_string):
    output_dict = {
        "speciation": 0,
        "duplication": 0,
        "confidence": 0,
        "scientific name": None,
    }

    if comment_string == "&&NHX:None":
        return output_dict
    else:
        info_dict = {}
        for info in comment_string.split(":"):
            info = info.split("=")
            if len(info) > 1:
                info_dict[info[0]] = info[1]

        if 'D' in info_dict and info_dict['D'] == 'Y':
            output_dict["duplication"] = 1
            output_dict["speciation"] = 0
        else:
            output_dict["duplication"] = 0
            output_dict["speciation"] = 1

        if 'B' in info_dict:
            output_dict["confidence"] = int(info_dict['B'])

        if 'S' in info_dict:
            output_dict["scientific name"] = info_dict['S']

        return output_dict


def add_nhx_for_node(clade):
    nhx_dict = nhx_comment_parse(clade.comment)
    clade.nhx_dict = nhx_dict

    if "scientific name" in nhx_dict:
        # clade.species = nhx_dict["scientific name"]
        clade.taxon = nhx_dict["scientific name"]
    if "speciation" in nhx_dict:
        clade.speciation = True if nhx_dict["speciation"] else False
    if "duplication" in nhx_dict:
        clade.duplication = True if nhx_dict["duplication"] else False
    if "confidence" in nhx_dict:
        clade.confidence = nhx_dict["confidence"]

    return clade


def load_nhx_tree(nhx_file):
    try:
        nhx_tree = Phylo.read(nhx_file, 'newick')
        nhx_tree = add_clade_name(nhx_tree)
        nhx_tree_node_dict = lookup_by_names(nhx_tree)

        # for i in gene_tree_node_dict:
        #     if not gene_tree_node_dict[i].is_terminal():
        #         gene_tree_node_dict[i] = add_nhx_for_node(gene_tree_node_dict[i])

        for i in nhx_tree_node_dict:
            nhx_tree_node_dict[i] = add_nhx_for_node(nhx_tree_node_dict[i])

        return nhx_tree, nhx_tree_node_dict
    except:
        raise ValueError(nhx_file)


def get_top_taxon_clade(taxon_labeled_tree, species_tree, taxonomy, node_type=None, confidence_threshold=None):
    sp_tree_dict = lookup_by_names(species_tree)
    sp_list = [i.name for i in sp_tree_dict[taxonomy].get_terminals()]

    candi_clade_list = []
    for clade in taxon_labeled_tree.find_clades(order='preorder'):
        good_type = True

        if clade.taxon == taxonomy:
            good_type = False

            if node_type == 'speciation' and hasattr(clade, 'speciation') and clade.speciation:
                good_type = True
            elif node_type == 'duplication' and hasattr(clade, 'duplication') and clade.duplication:
                good_type = True
            elif node_type is None:
                good_type = True

            if confidence_threshold and clade.confidence < confidence_threshold:
                good_type = False

        if good_type and len(set([l.taxon for l in clade.get_terminals()]) - set(sp_list)) == 0:
            candi_clade_list.append(clade)

    top_clade_list = []
    for i in candi_clade_list:
        top_flag = True
        for j in candi_clade_list:
            if is_offspring_of(i, j, taxon_labeled_tree):
                top_flag = False
        if top_flag:
            top_clade_list.append(i)

    return top_clade_list


# plot tree

# draw function fork from biopython

def draw_ascii(raw_tree, file=None, column_width=80, clade_name=False, clade_confidence=False, rename_dict=None):
    """Draw an ascii-art phylogram of the given tree.

    The printed result looks like::

                                        _________ Orange
                         ______________|
                        |              |______________ Tangerine
          ______________|
         |              |          _________________________ Grapefruit
        _|              |_________|
         |                        |______________ Pummelo
         |
         |__________________________________ Apple


    :Parameters:
        file : file-like object
            File handle opened for writing the output drawing. (Default:
            standard output)
        column_width : int
            Total number of text columns used by the drawing.

    """
    if file is None:
        file = sys.stdout

    tree = copy.deepcopy(raw_tree)

    if rename_dict:
        for clade in tree.find_clades(order='postorder'):
            if clade.name and clade.name in rename_dict:
                clade.name = rename_dict[clade.name]

    taxa = tree.get_terminals()
    # Some constants for the drawing calculations
    max_label_width = max(len(str(taxon)) for taxon in taxa)
    drawing_width = column_width - max_label_width - 1
    drawing_height = 2 * len(taxa) - 1

    def get_col_positions(tree):
        """Create a mapping of each clade to its column position."""
        depths = tree.depths()
        # If there are no branch lengths, assume unit branch lengths
        if not max(depths.values()):
            depths = tree.depths(unit_branch_lengths=True)
        # Potential drawing overflow due to rounding -- 1 char per tree layer
        fudge_margin = int(math.ceil(math.log(len(taxa), 2)))
        cols_per_branch_unit = ((drawing_width - fudge_margin) /
                                float(max(depths.values())))
        return dict((clade, int(blen * cols_per_branch_unit + 1.0))
                    for clade, blen in depths.items())

    def get_row_positions(tree):
        positions = dict((taxon, 2 * idx) for idx, taxon in enumerate(taxa))

        def calc_row(clade):
            for subclade in clade:
                if subclade not in positions:
                    calc_row(subclade)
            positions[clade] = ((positions[clade.clades[0]] +
                                 positions[clade.clades[-1]]) // 2)

        calc_row(tree.root)
        return positions

    col_positions = get_col_positions(tree)
    row_positions = get_row_positions(tree)
    char_matrix = [[' ' for x in range(drawing_width)]
                   for y in range(drawing_height)]

    def draw_clade(clade, startcol):
        thiscol = col_positions[clade]
        thisrow = row_positions[clade]
        # Draw a horizontal line
        for col in range(startcol, thiscol):
            char_matrix[thisrow][col] = '_'
        if clade.clades:
            # Draw a vertical line
            toprow = row_positions[clade.clades[0]]
            botrow = row_positions[clade.clades[-1]]
            for row in range(toprow + 1, botrow + 1):
                char_matrix[row][thiscol] = '|'
            # NB: Short terminal branches need something to stop rstrip()
            if (col_positions[clade.clades[0]] - thiscol) < 2:
                char_matrix[toprow][thiscol] = ','

            # add label for clade draw, by xuyuxing
            clade_string = ""
            if clade_name and clade.name:
                clade_string = clade_string + clade.name + ":"
            if clade_confidence and clade.confidence:
                if len(str(clade.confidence)) <= len("%.2e" % clade.confidence):
                    clade_string = clade_string + str(clade.confidence)
                else:
                    clade_string = clade_string + "%.2e" % clade.confidence
            for i in range(len(clade_string)):
                if len(clade.get_terminals()) % 2:
                    char_matrix[thisrow - 1][thiscol + 1 + i] = clade_string[i]
                else:
                    char_matrix[thisrow][thiscol + 1 + i] = clade_string[i]

            # Draw descendents
            for child in clade:
                draw_clade(child, thiscol + 1)

    draw_clade(tree.root, 0)
    # Print the complete drawing
    for idx, row in enumerate(char_matrix):
        line = ''.join(row).rstrip()
        # Add labels for terminal taxa in the right margin
        if idx % 2 == 0:
            line += ' ' + str(taxa[idx // 2])
        file.write(line + '\n')
    file.write('\n')


def scale_to_give_size(a_point_coord, target_size, raw_size, blank=(0, 0, 0, 0), reverse_x=False, reverse_y=False):
    """
    scale a point in a raw size to a new size

    point = (0,0)
    raw_size = (-2,2,-2,2) # x leaf,x right, y tail,y head

    target_size = (0,20,0,100)
    new point = (10,50)

    """

    b_x_L, b_x_R, b_y_T, b_y_H = blank
    a_x_L, a_x_R, a_y_T, a_y_H = target_size
    w_x_L, w_x_R, w_y_T, w_y_H = a_x_L + b_x_L, a_x_R - \
        b_x_R, a_y_T + b_y_T, a_y_H - b_y_H
    r_x_L, r_x_R, r_y_T, r_y_H = raw_size

    if reverse_x:
        o_x = w_x_R - (a_point_coord[0] - r_x_L) / \
            (r_x_R - r_x_L) * (w_x_R - w_x_L)
    else:
        o_x = (a_point_coord[0] - r_x_L) / \
            (r_x_R - r_x_L) * (w_x_R - w_x_L) + w_x_L

    if reverse_y:
        o_y = w_y_H - (a_point_coord[1] - r_y_T) / \
            (r_y_H - r_y_T) * (w_y_H - w_y_T)
    else:
        o_y = (a_point_coord[1] - r_y_T) / \
            (r_y_H - r_y_T) * (w_y_H - w_y_T) + w_y_T

    return o_x, o_y


def line_segment_plot(ax, start, end, facecolor='w', lw=2):
    verts = [start, end]
    codes = [mpath.Path.MOVETO, mpath.Path.LINETO]

    path = mpath.Path(verts, codes)
    patch = mpatches.PathPatch(path, facecolor=facecolor, lw=lw)
    ax.add_patch(patch)


def normal_tree_rela_coord(tree):
    """
    add an attr named 'rela_coord' to every clade in the tree. The X is the branch value, and the Y is base on leaf number.
    """
    # get X axis aka branch length from root to clade
    tree = add_clade_name(tree)
    tree_name_dir = lookup_by_names(tree, order='preorder')

    for clade_name in tree_name_dir:
        X_clade = get_distance_to_root(tree_name_dir[clade_name], tree)
        tree_name_dir[clade_name].rela_coord = (X_clade, None)

    # get Y axis
    # first for leaf
    order_leaf_list = [
        clade_name for clade_name in tree_name_dir if tree_name_dir[clade_name].is_terminal()]
    for leaf_name in order_leaf_list:
        X_clade, Y_clade = tree_name_dir[leaf_name].rela_coord
        if (X_clade is None) or (not Y_clade is None):
            raise ValueError("somebady make a rela_coord already")
        Y_clade = float(order_leaf_list.index(leaf_name))
        tree_name_dir[leaf_name].rela_coord = (X_clade, Y_clade)

    # for orther clade, we should cal node which close to leaf
    path_for_target_to_most_far_leaf_dict = {}
    for clade_name in tree_name_dir:
        if clade_name in order_leaf_list:
            continue
        else:
            most_far_path = max([get_path_step(leaf_tmp, tree_name_dir[clade_name], tree) for leaf_tmp in
                                 tree_name_dir[clade_name].get_terminals()])
            if most_far_path not in path_for_target_to_most_far_leaf_dict:
                path_for_target_to_most_far_leaf_dict[most_far_path] = []
            path_for_target_to_most_far_leaf_dict[most_far_path].append(
                clade_name)

    for i in sorted(path_for_target_to_most_far_leaf_dict.keys()):
        for clade_name in path_for_target_to_most_far_leaf_dict[i]:
            X_clade, Y_clade = tree_name_dir[clade_name].rela_coord
            son_Y = [son.rela_coord[1] for son in tree_name_dir[clade_name]]
            Y_clade = (max(son_Y) - min(son_Y)) / 2 + min(son_Y)
            tree_name_dir[clade_name].rela_coord = (X_clade, Y_clade)

    return tree


def normal_tree_abs_coord(tree, figure_size):
    tree_name_dict = lookup_by_names(tree)

    tree_clade_rela_coord = {
        clade_name: tree_name_dict[clade_name].rela_coord for clade_name in tree_name_dict}
    rela_x_list = [x for x, y in tree_clade_rela_coord.values()]
    rela_y_list = [y for x, y in tree_clade_rela_coord.values()]
    rela_fig_size = (min(rela_x_list), max(rela_x_list),
                     min(rela_y_list), max(rela_y_list))

    tree_clade_abs_coord = {
        clade_name: scale_to_give_size(tree_clade_rela_coord[clade_name], figure_size, rela_fig_size) for clade_name in
        tree_clade_rela_coord}

    for clade_name in tree_name_dict:
        tree_name_dict[clade_name].abs_coord = tree_clade_abs_coord[clade_name]

    return tree


def draw_tree(ax, tree):
    tree_name_dict = lookup_by_names(tree, 'level')
    for clade_name in tree_name_dict:
        clade_tmp = tree_name_dict[clade_name]
        son_clade = clade_tmp.clades
        son_y = []
        for son_clade_tmp in son_clade:
            line_segment_plot(ax, (son_clade_tmp.abs_coord[0], son_clade_tmp.abs_coord[1]),
                              (clade_tmp.abs_coord[0], son_clade_tmp.abs_coord[1]))
            son_y.append(son_clade_tmp.abs_coord[1])
        if not clade_tmp.is_terminal():
            line_segment_plot(ax, (clade_tmp.abs_coord[0], min(
                son_y)), (clade_tmp.abs_coord[0], max(son_y)))


# #########################3
#
# def get_speci_id(seq_id):
#     return re.findall(r'^(\d+)\_.*', seq_id)[0]
#
#
# def orthofinder_find_root(gene_tree, speci_tree):
#     """
#     https://www.biorxiv.org/content/biorxiv/early/2019/04/24/466201.full.pdf
#     """
#     pass
#
#
# def get_node_from_target_to_leaf(target_node, leaf):
#     if target_node not in leaf.get_ancestors():
#         raise ValueError("target node is not an ancestor of leaf")
#     output_list = []
#     for i in leaf.get_ancestors():
#         if i == target_node:
#             break
#         output_list.append(i)
#     return output_list
#
#
# def get_node_depth(node, ignore_pass_node):
#     """
#     given a node return the depth of this node, leaf is -1, end node is 0
#     """
#
#     # tmp_dir = {}
#     # for leaves in node.get_leaves():
#     #    tmp_dir[leaves.name]=[tmp_anc for tmp_anc in get_node_from_target_to_leaf(node,leaves) if not pass_node(tmp_anc)]
#
#     if node.is_leaf():
#         return -1
#     else:
#         if ignore_pass_node:
#             return max(
#                 [len([tmp_anc for tmp_anc in get_node_from_target_to_leaf(node, leaves) if not pass_node(tmp_anc)]) for
#                  leaves in node.get_leaves()])
#         else:
#             return max(
#                 [len([tmp_anc for tmp_anc in get_node_from_target_to_leaf(node, leaves)]) for
#                  leaves in node.get_leaves()])
#
#
# if __name__ == '__main__':
#     tree_file = '/lustre/home/xuyuxing/tmp/tree14.tre'
#     tree = Phylo.read(tree_file, "newick")
#     tree = add_clade_name(tree)
#     node_dict = lookup_by_names(tree)
#
#     import pickle
#     import sys
#
#     sys.path.extend(['/lustre/home/xuyuxing/python_project/Genome_work_tools'])
#
#     import argparse
#
#     ###### argument parse
#     parser = argparse.ArgumentParser(
#         prog='PhylTools',
#     )
#
#     subparsers = parser.add_subparsers(title='subcommands', dest="subcommand_name")
#
#     # argparse for get_close_node_info
#     parser_a = subparsers.add_parser('get_close_node_info',
#                                      help='build a protein gene tree from raw sequence',
#                                      description='build a protein gene tree from raw sequence')
#
#     parser_a.add_argument('tree_file', type=str, help='Path of raw protein sequences')
#     parser_a.add_argument('target_speci', type=str, help='Path of raw protein sequences')
#     parser_a.add_argument('taxon_info', type=str, help='Path of raw protein sequences')
#     parser_a.add_argument('taxonomy_dir', type=str, help='Path of raw protein sequences')
#     parser_a.add_argument('taxon_lineage_file', type=str, help='Path of raw protein sequences')
#
#     args = parser.parse_args()
#     args_dict = vars(args)
#
#     # taxon_info = "/lustre/home/xuyuxing/Work/Balanophora/orthofinder/pep2/OrthoFinder/Results_Sep05/WorkingDirectory/tmp/taxon_vs_ortho"
#     # tree_file = "/lustre/home/xuyuxing/Work/Balanophora/orthofinder/pep2/OrthoFinder/Results_Sep05/WorkingDirectory/tmp/OG0000443_tree_id.txt"
#     # taxonomy_dir = "/lustre/home/xuyuxing/Database/genome/info/NCBI/taxonomy"
#     # taxon_lineage_file = "/lustre/home/xuyuxing/Work/Balanophora/orthofinder/pep2/OrthoFinder/Results_Sep05/WorkingDirectory/tmp/lineage.pyb"
#     # target_speci = "15"
#
#     if args_dict["subcommand_name"] == "get_close_node_info":
#
#         from toolbiox.lib.common.fileIO import tsv_file_dict_parse
#
#         tax_info_dict = tsv_file_dict_parse(args.taxon_info, fieldnames=["taxon_id", "speci_ID"], key_col="speci_ID",
#                                             delimiter=",")
#         tree = Phylo.read(args.tree_file, "newick")
#         tree, node_dict = mid_unrooted_tree(tree)
#
#         leaf_dict = {}
#         leaf_speci_hash = {}
#         for leaf in tree.get_terminals():
#             speci_id = get_speci_id(leaf.name)
#             leaf_speci_hash[leaf.name] = speci_id
#             if speci_id not in leaf_dict:
#                 leaf_dict[speci_id] = []
#             leaf_dict[speci_id].append(leaf.name)
#
#         target_leaf = []
#         for leaf in tree.get_terminals():
#             if leaf_speci_hash[leaf.name] == args.target_speci:
#                 target_leaf.append(leaf.name)
#
#         # from toolbiox.lib.resource.ncbi_taxonomy import build_taxon_database
#         #
#         # tax_record_dict = build_taxon_database(taxonomy_dir)
#         #
#         # lineage_dict = {}
#         # for speci_id in leaf_dict:
#         #     taxon_record = tax_record_dict[tax_info_dict[speci_id]['taxon_id']]
#         #     lineage = taxon_record.get_lineage(tax_record_dict)
#         #     lineage_dict[speci_id] = lineage
#         #
#         # OUT = open(taxon_lineage_file, 'wb')
#         # pickle.dump(lineage_dict, OUT)
#         # OUT.close()
#
#         TEMP = open(args.taxon_lineage_file, 'rb')
#         lineage_dict = pickle.load(TEMP)
#         TEMP.close()
#
#
#         def get_close_node_without_same_speci(tree, child_node, leaf_speci_hash, target_speci):
#             parent_node = get_parent(child_node, tree)
#             close_node_leaf = [leaf.name for leaf in parent_node.get_terminals() if
#                                leaf_speci_hash[leaf.name] != target_speci]
#             if len(close_node_leaf) == 0 and parent_node != tree:
#                 return get_close_node_without_same_speci(tree, parent_node, leaf_speci_hash, target_speci)
#             else:
#                 return parent_node
#
#
#         def common_lineage(lineage_list):
#             if len(lineage_list) == 1:
#                 return lineage_list[0][-1]
#
#             common_taxon = set(lineage_list[0])
#             for i in lineage_list:
#                 common_taxon = set(i) & common_taxon
#
#             a = []
#             for i in lineage_list[0]:
#                 if not i in common_taxon:
#                     return a[-1]
#                 else:
#                     a.append(i)
#
#
#         for leaf_name in leaf_dict[args.target_speci]:
#             close_node = get_close_node_without_same_speci(tree, node_dict[leaf_name], leaf_speci_hash,
#                                                            args.target_speci)
#             close_taxon = common_lineage(
#                 [lineage_dict[get_speci_id(leaf.name)] for leaf in close_node.get_terminals() if
#                  leaf_speci_hash[leaf.name] != args.target_speci])
#             print(args.tree_file, leaf_name, close_node.name, close_taxon[0], close_taxon[1], close_node.confidence)

# for binary tree


class BinaryTreeNode(Clade):
    def __init__(
        self,
        branch_length=None,
        name=None,
        left=None,
        right=None,
        confidence=None,
        color=None,
        width=None,
    ):
        """Define parameters for the BinaryTreeNode Clade."""
        if left is None and right is None:
            clades = None
        else:
            clades = [left, right]
        super(BinaryTreeNode, self).__init__(
            branch_length=branch_length,
            name=name,
            clades=clades,
            confidence=confidence,
            color=color,
            width=width,
        )

        self.left = left
        self.right = right


def normalCalde2BinaryTreeNode(clade):
    if clade.is_terminal():
        return BinaryTreeNode(clade.branch_length, clade.name, None, None, clade.confidence, clade.color, clade.width)
    elif len(clade.clades) == 2:
        left = normalCalde2BinaryTreeNode(clade.clades[0])
        right = normalCalde2BinaryTreeNode(clade.clades[1])
        return BinaryTreeNode(clade.branch_length, clade.name, left, right, clade.confidence, clade.color, clade.width)
    else:
        raise ValueError("clade is not good to be a Binary clade")


def Tree2BinaryTree(tree):
    tree_now = copy.deepcopy(tree)

    tree_clade = tree_now.clade
    BT_clade = normalCalde2BinaryTreeNode(tree_clade)

    new_tree = Tree(BT_clade, tree_now.rooted)

    return new_tree


def allPossibleFBT(N):
    """
    get all possible full binary tree structure, left or right will be treat as difference

    N: all node number, 2 * leaf -1
    return: a list of BinaryTreeNode
    """

    if N == 1:
        return [BinaryTreeNode()]
    res = []

    for i in range(1, N, 2):
        for left in allPossibleFBT(i):
            for right in allPossibleFBT(N-i-1):
                node = BinaryTreeNode(left=left, right=right)
                res.append(node)

    return res


def equalBT(BT1, BT2):
    if BT1.left is None and BT1.right is None and BT2.left is None and BT2.right is None:
        if BT1.name == BT2.name:
            return True
    elif BT1.left is not None and BT1.right is not None and BT2.left is not None and BT2.right is not None:
        if (equalBT(BT1.left, BT2.left) and equalBT(BT1.right, BT2.right)) or (equalBT(BT1.left, BT2.right) and equalBT(BT1.right, BT2.left)):
            return True

    return False


def equalBTtree(BT1, BT2):
    return equalBT(BT1.clade, BT2.clade)


def get_newick(tree):
    f = StringIO()
    Phylo.write(tree, f, 'newick')
    string_tmp = f.getvalue()
    f.close()
    return string_tmp


def merge_same_binary_tree(BT_list):
    edge_list = []
    for i in BT_list:
        for j in BT_list:
            if equalBTtree(i, j):
                edge_list.append((i, j))

    G = nx.Graph()
    G.add_nodes_from(BT_list)
    G.add_edges_from(edge_list)

    output_list = []
    for sub_graph in (G.subgraph(c) for c in nx.connected_components(G)):
        output_list.append(list(sub_graph.nodes))

    return output_list


def all_possible_tree(leaf_list):
    node_num = len(leaf_list) * 2 - 1
    tree_structure_list = [i[0]
                           for i in merge_same_binary_tree(allPossibleFBT(node_num))]

    renamed_tree_list = []
    for tree_structure in tree_structure_list:

        for leaf_order in permutations(leaf_list):
            tree_now = copy.deepcopy(tree_structure)

            num = 0
            for clade in tree_now.find_clades(order='preorder'):
                if clade.is_terminal():
                    # print(num)
                    clade.name = leaf_order[num]
                    num += 1

            renamed_tree_list.append(tree_now)

    renamed_tree_list = [i[0]
                         for i in merge_same_binary_tree(renamed_tree_list)]

    return renamed_tree_list


def gene_loss_model_test(full_model_gene_tree, gene_species_map_dict):
    """
    if we have a gene tree model, but may gene loss will change the topologist of the tree. so test if how many possible case when gene loss

    from io import StringIO
    full_tree_string = "(((Gel_a1,Gel_a2)B1,(Ash_a1,Ash_a2)C1)S1,((Gel_b1,Gel_b2)B2,(Ash_b1,Ash_b2)C2)S2)A,Aco"

    full_model_gene_tree = Phylo.read(StringIO(full_tree_string), 'newick')

    gene_species_map_dict = {
        'Gel_a1': 'Gel',
        'Gel_a2': 'Gel',
        'Gel_b1': 'Gel',
        'Gel_b2': 'Gel',
        'Ash_a1': 'Ash',
        'Ash_a2': 'Ash',
        'Ash_b1': 'Ash',
        'Ash_b2': 'Ash',
    }

    gene not in gene_species_map_dict will not loss

    """

    may_loss_gene_list = gene_species_map_dict.keys()

    keep_case_dict = {}
    num = 0
    for i in range(1, len(may_loss_gene_list)+1):
        for keep_case in combinations(may_loss_gene_list, i):
            keep_case_dict[num] = {
                'keep_gene': keep_case,
                'lost_gene': list(set(may_loss_gene_list) - set(keep_case)),
            }
            num += 1

    # get gene loss tree
    for id_tmp in keep_case_dict:
        lost_gene_id = keep_case_dict[id_tmp]['lost_gene']
        gene_loss_tree = remove_given_node_from_tree(
            full_model_gene_tree, lost_gene_id)
        gene_loss_tree = remove_tree_node_name(gene_loss_tree)
        gene_loss_tree = clade_rename(gene_loss_tree, gene_species_map_dict)
        gene_loss_tree = Tree2BinaryTree(gene_loss_tree)
        keep_case_dict[id_tmp]['BT_tree'] = gene_loss_tree

    # merge
    BT_list = [keep_case_dict[i]['BT_tree'] for i in keep_case_dict]
    uniq_BT_list = [i[0] for i in merge_same_binary_tree(BT_list)]

    uniq_BT_dict = {}
    for uniq_BT in uniq_BT_list:
        uniq_BT_dict[uniq_BT] = []
        for id_tmp in keep_case_dict:
            gene_loss_tree = keep_case_dict[id_tmp]['BT_tree']
            if equalBTtree(uniq_BT, gene_loss_tree):
                uniq_BT_dict[uniq_BT].append(id_tmp)

    return keep_case_dict, uniq_BT_dict


if __name__ == '__main__':
    gene_tree_file = '/lustre/home/xuyuxing/Work/Csp/orthofinder/protein_seq/Results_Apr10/Orthologues_Apr11/WorkingDirectory/Trees_ids/OG0000532_tree_id.txt'
    species_rooted_tree_file = '/lustre/home/xuyuxing/Work/Csp/orthofinder/protein_seq/Results_Apr10/Orthologues_Apr11/SpeciesTree_rooted.id.txt'

    # reads a tree file
    gene_tree = Phylo.read(gene_tree_file, 'newick')
    gene_tree = add_clade_name(gene_tree)
    gene_tree_node_dict = lookup_by_names(gene_tree)

    species_rooted_tree = Phylo.read(species_rooted_tree_file, 'newick')

    # get all leaf
    leaf_name_list = [leaf.name for leaf in gene_tree.get_terminals()]

    gene_to_species_map = {'1_7601': '1', '8_15565': '8', '1_16489': '1', '8_8464': '8', '9_13535': '9', '6_10585': '6',
                           '2_11400': '2', '5_10489': '5', '4_30192': '4', '5_24848': '5', '6_28832': '6',
                           '2_18137': '2', '4_9871': '4', '5_294': '5', '4_9272': '4', '5_24710': '5', '4_9271': '4',
                           '0_26868': '0', '3_2604': '3', '7_24217': '7', '1_9240': '1', '9_19573': '9', '8_12531': '8',
                           '6_14246': '6', '4_20069': '4', '5_23488': '5', '4_15038': '4', '5_3351': '5', '2_9356': '2',
                           '6_7876': '6', '4_23206': '4', '5_41654': '5', '4_43364': '4', '2_12645': '2',
                           '4_23207': '4', '5_4308': '5', '0_3911': '0', '0_13321': '0', '3_8434': '3', '7_5793': '7'}

    # root gene tree by species info
    best_root_clade = get_root_by_species(
        gene_tree, species_rooted_tree, gene_to_species_map)

    gene_tree_rooted, gene_tree_rooted_node_dict, gene_tree, gene_tree_node_dict = reroot_by_outgroup_clade(gene_tree,
                                                                                                            gene_tree_node_dict,
                                                                                                            best_root_clade.name,
                                                                                                            True)

    # draw a tree
    from io import StringIO

    # normal tree
    tree_string = "(((Orange:0.1,Tangerine:0.2):0.3,(Grapefruit:0.4,Pummelo:0.2):0.1):0.2,Apple:0.6)"
    # tree with comment
    tree_string = "(((Orange:0.1,Tangerine:0.2):0.3[50],(Grapefruit:0.4,Pummelo:0.2):0.1[50]):0.2[60],Apple:0.6)"
    # tree with confidence
    tree_string = "(((Orange:0.1,Tangerine:0.2)50:0.3,(Grapefruit:0.4,Pummelo:0.2)50:0.1)60:0.2,Apple:0.6)"
    tree = Phylo.read(StringIO(tree_string), 'newick')

    draw_ascii(tree)

    # work with nhx tree file
    from toolbiox.lib.common.fileIO import tsv_file_dict_parse

    normal_gene_tree_file = "/lustre/home/xuyuxing/Work/Gel/orcWGD_redo/phylogenomics_old/tree/fasttree/OG0002678/tree.phb"
    species_tree_file = "/lustre/home/xuyuxing/Work/Gel/orcWGD_redo/species.tree.txt"
    rename_map = "/lustre/home/xuyuxing/Work/Gel/orcWGD_redo/phylogenomics_old/tree/fasttree/OG0002678/rename.map"

    gene_tree = Phylo.read(normal_gene_tree_file, 'newick')
    gene_tree = add_clade_name(gene_tree)
    species_tree = Phylo.read(species_tree_file, 'newick')
    species_tree = add_clade_name(species_tree)

    tmp_info = tsv_file_dict_parse(rename_map, fieldnames=[
                                   'new_id', 'old_id', 'speci'], key_col='new_id')
    gene_to_species_map_dict = {i: tmp_info[i]["speci"] for i in tmp_info}

    rooted_gene_tree = get_rooted_tree_by_species_tree(
        gene_tree, species_tree, gene_to_species_map_dict)
    rooted_taxon_gene_tree = map_node_species_info(
        rooted_gene_tree, species_tree, gene_to_species_map_dict)

    nhx_tree = get_nhx_comment(rooted_taxon_gene_tree)

    nhx_file = '/lustre/home/xuyuxing/temp/best.nhx'
    nhx_tree, clade_dict = load_nhx_tree(nhx_file)

    for clade_name in clade_dict:
        clade = clade_dict[clade_name]
        if not clade.is_terminal():
            print(clade.name, clade.nhx_dict['Duplications'])

    # get all possible tree
    leaf_list = ['Gel_a1', 'Gel_a2', 'Ash_a1', 'Ash_a2']
    tree_list = all_possible_tree(leaf_list)

    num = 0
    for i in tree_list:
        print(num)
        num += 1
        Phylo.draw_ascii(i)

    # get all possible tree
    leaf_list = ['Gel', 'Gel', 'Ash', 'Ash']
    tree_list = all_possible_tree(leaf_list)

    num = 0
    for i in tree_list:
        print(num)
        num += 1
        Phylo.draw_ascii(i)

    # gene loss tree
    from io import StringIO
    full_tree_string = "(((Gel_a1,Gel_a2)B1,(Ash_a1,Ash_a2)C1)S1,((Gel_b1,Gel_b2)B2,(Ash_b1,Ash_b2)C2)S2)A,Aco"

    tree = Phylo.read(StringIO(full_tree_string), 'newick')
    draw_ascii(tree, clade_name=True)

    """
                                                     _________________ Gel_a1
                                    ________________|B1:
                                   |                |_________________ Gel_a2
                   ________________|S1:
                  |                |                 _________________ Ash_a1
                  |                |________________|C1:
                  |                                 |_________________ Ash_a2
  ________________|A:
 |                |                                  _________________ Gel_b1
 |                |                 ________________|B2:
 |                |                |                |_________________ Gel_b2
_|                |________________|S2:
 |                                 |                 _________________ Ash_b1
 |                                 |________________|C2:
 |                                                  |_________________ Ash_b2
 |
 |________________ Aco
    
    
    """

    gene_species_map_dict = {
        'Gel_a1': 'Gel',
        'Gel_a2': 'Gel',
        'Gel_b1': 'Gel',
        'Gel_b2': 'Gel',
        'Ash_a1': 'Ash',
        'Ash_a2': 'Ash',
        'Ash_b1': 'Ash',
        'Ash_b2': 'Ash',
    }

    keep_case_dict, uniq_BT_dict = gene_loss_model_test(
        tree, gene_species_map_dict)

    Phylo.draw(tree)

    from Bio.AlignIO import NexusIO

    from Bio import SeqIO

    from Bio.SeqIO import parse

    #
