import logging
import random
from collections import defaultdict, Counter
from functools import reduce
from queue import Queue

import numpy as np
import pandas as pd
from ete3 import Tree, TreeNode

DATE = 'date'

REASONABLE_NUMBER_OF_TIPS = 15

CATEGORIES = 'categories'

NODE_SIZE = 'node_size'
NUM_TIPS_INSIDE = 'max_size'
NODE_NAME = 'node_name'

TIPS_INSIDE = 'in_tips'
TIPS_BELOW = 'all_tips'

EDGE_SIZE = 'edge_size'
EDGE_NAME = 'edge_name'
METACHILD = 'metachild'
FONT_SIZE = 'node_fontsize'


def get_dist_to_root(tip):
    dist_to_root = 0
    n = tip
    while not n.is_root():
        dist_to_root += n.dist
        n = n.up
    return dist_to_root


def date_tips(tree, date_df):
    """
    Adds dates to the tips as 'date' attribute.
    :param tree: ete3.Tree
    :param date_df: a pandas.Series with tip ids as indices and dates as values
    :return: void, modifies the initial tree
    """

    def _get_date(d):
        if pd.notnull(d):
            first_jan_this_year = pd.datetime(year=d.year, month=1, day=1)
            day_of_this_year = d - first_jan_this_year
            first_jan_next_year = pd.datetime(year=d.year + 1, month=1, day=1)
            days_in_this_year = first_jan_next_year - first_jan_this_year
            return d.year + day_of_this_year / days_in_this_year
        else:
            return None

    id2tip = {n.name: n for n in tree}
    date_df = date_df[date_df.index.isin(id2tip) & ~pd.isnull(date_df)]
    dated_fraction = len(date_df) / len(id2tip)
    if dated_fraction < .5:
        raise ValueError('Too few dates are provided (only for {}% of tips)!'.format('%g' % (100 * dated_fraction)))

    for id, value in date_df.iteritems():
        id2tip[id].add_feature(DATE, int(_get_date(value)))

    min_date, max_date = int(_get_date(date_df.min())), int(_get_date(date_df.max()))

    if len(date_df) < len(id2tip):
        unique_dates = list(date_df.unique())
        if len(unique_dates) == 1:
            for id in set(id2tip.keys()) - set(date_df.index):
                id2tip[id].add_feature(DATE, int(_get_date(unique_dates[0])))
        else:
            rates = []
            for _ in range(10):
                date1, date2 = random.sample(unique_dates, 2)
                id1, date1 = next(date_df[date_df == date1].sample(1).iteritems())
                id2, date2 = next(date_df[date_df == date2].sample(1).iteritems())

                dist1 = get_dist_to_root(id2tip[id1])
                dist2 = get_dist_to_root(id2tip[id2])
                rate = (_get_date(date2) - _get_date(date1)) / (dist2 - dist1)
                rates.append(rate)
            rate = np.mean(rates)

            for id in set(id2tip.keys()) - set(date_df.index):
                id2tip[id].add_feature(DATE,
                                       min(min_date, max(max_date,
                                                         int(_get_date(date1) + rate * (
                                                                 get_dist_to_root(id2tip[id]) - dist1)))))

    return min(_.date for _ in id2tip.values()), max(_.date for _ in id2tip.values())


def name_tree(tree):
    """
    Names all the tree nodes that are not named, with unique names.
    :param tree: ete3.Tree, the tree to be named
    :return: void, modifies the original tree
    """
    existing_names = Counter((_.name for _ in tree.traverse() if _.name))
    i = 0
    for node in tree.traverse('postorder'):
        if node.is_root():
            node.name = 'ROOT'
            existing_names['ROOT'] += 1
        if not node.name or existing_names[node.name] > 1:
            name = None
            while not name or name in existing_names:
                name = '{}_{}'.format('tip' if node.is_leaf() else 'node', i)
                i += 1
            node.name = name


def collapse_zero_branches(tree):
    num_collapsed = 0
    for n in list(tree.traverse('postorder')):
        for child in list(n.children):
            if not child.is_leaf() and child.dist <= 0:
                n.remove_child(child)
                for grandchild in child.children:
                    n.add_child(grandchild)
                num_collapsed += 1
            if child.is_leaf() and child.dist < 0:
                child.dist = 0
    logging.info('Collapsed {} zero branches.'.format(num_collapsed))


def remove_certain_leaves(tr, to_remove=lambda node: False):
    """
    Removes all the branches leading to leaves identified positively by to_remove function.
    :param tr: the tree of interest (ete3 Tree)
    :param to_remove: a method to check is a leaf should be removed.
    :return: void, modifies the initial tree.
    """

    tips = [tip for tip in tr if to_remove(tip)]
    for node in tips:
        if node.is_root():
            return None
        parent = node.up
        parent.remove_child(node)
        # If the parent node has only one child now, merge them.
        if len(parent.children) == 1:
            brother = parent.children[0]
            brother.dist += parent.dist
            if parent.is_root():
                brother.up = None
                tr = brother
            else:
                grandparent = parent.up
                grandparent.remove_child(parent)
                grandparent.add_child(brother)
    return tr


def pasml_annotations2cytoscape_annotation(cat2file, output, sep='\t'):
    logging.info('Combining the data from different columns into {}'.format(output))

    def get_state(name, df):
        row = df.loc[name, :]
        states = df.columns[row]
        return None if len(states) != 1 else states[0]

    cat2df = {cat: pd.read_table(data_path, sep=',', index_col=0, header=0).astype(bool) for (cat, data_path) in
              cat2file.items()}
    df = pd.DataFrame(index=next(iter(cat2df.values())).index, columns=cat2df.keys())
    for cat, cat_df in cat2df.items():
        cat_df.index = cat_df.index.map(str)
        df[cat] = df.index.map(lambda name: get_state(name, cat_df))

    if len(cat2df) == 1:
        cat_df = next(iter(cat2df.values()))
        rsuffix = 'category' if set(df.columns) & set(cat_df.columns) else ''
        df = df.join(cat_df, rsuffix=rsuffix, lsuffix='')

    df.to_csv(output, sep=sep, index_label='Node')


def annotate_tree_with_cyto_metadata(tree, data_path, columns, sep='\t'):
    df = pd.read_table(data_path, sep=sep, index_col=0, header=0)
    df.index = df.index.map(str)
    df.fillna('', inplace=True)
    tree = read_tree(tree) if not isinstance(tree, TreeNode) else tree

    for n in tree.traverse():
        if len(columns) != 1:
            n.add_features(**df.loc[n.name, :].to_dict())
        else:
            category = col_name2cat(columns[0])
            data = df.loc[n.name, :]
            n.add_features(**data[data != False].to_dict())
            # In case category's value was also False (and therefore was not added), let's add it again
            n.add_feature(category, data[category])
    return tree, sorted(df.columns)


def col_name2cat(column):
    """
    Reformats the column string to make sure it contains only numerical or letter characters.
    :param column: str, column name to be reformatted
    :return: str, the column name with illegal characters removed
    """
    column_string = ''.join(s for s in column if s.isalnum())
    return column_string


def read_tree(tree_path):
    for f in (3, 2, 5, 1, 0, 3, 4, 6, 7, 8, 9):
        try:
            return Tree(tree_path, format=f)
        except:
            continue
    raise ValueError('Could not read the tree {}. Is it a valid newick?'.format(tree_path))


def sum_len_values(kv_dict):
    return sum(len(_) for _ in kv_dict.values())


def get_states(n, categories):
    return {cat: getattr(n, cat) for cat in categories if hasattr(n, cat)}


def compress_tree(tree, categories, can_merge_diff_sizes=True, tip_size_threshold=REASONABLE_NUMBER_OF_TIPS):
    for n in tree.traverse('postorder'):
        n.add_feature(TIPS_INSIDE, defaultdict(list))
        n.add_feature(TIPS_BELOW, defaultdict(list))
        if n.is_leaf():
            getattr(n, TIPS_INSIDE)[getattr(n, DATE)].append(n.name)
            getattr(n, TIPS_BELOW)[getattr(n, DATE)].append(n.name)
        else:
            for _ in n:
                getattr(n, TIPS_BELOW)[getattr(_, DATE)].append(_.name)

    collapse_vertically(tree, lambda _: get_states(_, categories))
    remove_mediators(tree, lambda _: get_states(_, categories))

    for n in tree.traverse():
        n.add_feature(NUM_TIPS_INSIDE, sum_len_values(getattr(n, TIPS_INSIDE)))
        n.add_feature(TIPS_INSIDE, [getattr(n, TIPS_INSIDE)])
        n.add_feature(TIPS_BELOW, [getattr(n, TIPS_BELOW)])

    logging.info('Gonna collapse horizontally')
    get_bin = lambda _: _
    # collapse_horizontally(get_bin, tree, lambda _: get_states(_, categories))
    collapse_hor(tree, lambda _: get_states(_, categories), get_bin)

    if can_merge_diff_sizes and len(tree) > tip_size_threshold:
        get_bin = lambda _: int(np.log10(max(1, _)))

        logging.info('Gonna re-collapse horizontally, merging nodes of different sizes')
        # collapse_horizontally(get_bin, tree, lambda _: get_states(_, categories))
        collapse_hor(tree, lambda _: get_states(_, categories), get_bin)

    if len(tree) > tip_size_threshold:
        for n in tree.traverse():
            n.add_feature(NUM_TIPS_INSIDE, sum(sum_len_values(_) for _ in getattr(n, TIPS_INSIDE))
                          / len(getattr(n, TIPS_INSIDE)))

        def get_tsize(n):
            res = getattr(n, NUM_TIPS_INSIDE) * len(getattr(n, TIPS_INSIDE))
            while n.up:
                n = n.up
                res *= len(getattr(n, TIPS_INSIDE))
            return res

        tip_sizes = [get_tsize(_) for _ in tree]
        if len(tip_sizes) > tip_size_threshold:
            threshold = sorted(tip_sizes)[-tip_size_threshold]
            logging.info('Removing tips of size {} or less'.format(threshold))
            remove_small_tips(tree, to_be_removed=lambda _: get_tsize(_) <= threshold)
            remove_mediators(tree, lambda _: get_states(_, categories))

            logging.info('Gonna collapse horizontally one last time')
            # collapse_horizontally(get_bin, tree, lambda _: get_states(_, categories))
            collapse_hor(tree, lambda _: get_states(_, categories), get_bin)
    return tree


def collapse_hor(tree, get_states, tips2bin):
    config_cache = {}

    def get_configuration(n):
        if n not in config_cache:
            config_cache[n] = (len(getattr(n, TIPS_INSIDE)),
                               (tips2bin(getattr(n, NUM_TIPS_INSIDE)),
                                tuple(sorted('{}:{}'.format(k, v) for (k, v) in get_states(n).items())),
                                tuple(sorted([get_configuration(_) for _ in n.children]))))
        return config_cache[n]

    for n in tree.traverse('postorder'):
        config2children = defaultdict(list)
        for _ in n.children:
            config2children[get_configuration(_)[1]].append(_)
        for children in (_ for _ in config2children.values() if len(_) > 1):
            children = sorted(children, key=lambda _: getattr(_, DATE, 0))
            child = children[0]
            for c in children[1:]:
                n.remove_child(c)
            child.add_feature(METACHILD, True)

            tips_inside, tips_below = [], []
            for _ in children:
                tips_inside.extend(getattr(_, TIPS_INSIDE))
                tips_below.extend(getattr(_, TIPS_BELOW))

            child.add_feature(TIPS_INSIDE, tips_inside)
            child.add_feature(TIPS_BELOW, tips_below)
            child.add_feature(NUM_TIPS_INSIDE, sum(sum_len_values(_) for _ in getattr(child, TIPS_INSIDE))
                              / len(getattr(child, TIPS_INSIDE)))
            child.add_feature(DATE, min(getattr(_, DATE) for _ in children))
            if child in config_cache:
                config_cache[child] = (len(getattr(child, TIPS_INSIDE)), config_cache[child][1])


def remove_small_tips(tree, to_be_removed):
    changed = True
    while changed:
        changed = False
        for l in tree.get_leaves():
            parent = l.up
            if parent and to_be_removed(l):
                parent.remove_child(l)
                parent.add_feature(DATE,
                                   min(min(_.keys()) if _ else 200000 for _ in getattr(parent, TIPS_INSIDE))
                                   if getattr(parent, TIPS_INSIDE) else 200000)
                if parent.children:
                    parent.add_feature(DATE, min(getattr(parent, DATE), min(getattr(_, DATE) for _ in parent.children)))
                changed = True


def merge_features(main_node, nodes, features, op, default_value=0):
    for feature in features:
        main_node.add_feature(feature, op([getattr(node, feature, default_value) for node in nodes]))


def collapse_vertically(tree, get_states):
    """
    Collapses a child node into its parent if they are in the same state.
    :param get_states: a function that returns a set of node states
    :param tree: ete3.Tree
    :return: void, modifies the input tree
    """
    for n in tree.traverse('postorder'):
        if n.is_leaf():
            continue

        states = set(get_states(n).items())
        children = list(n.children)
        for child in children:
            # merge the child into this node if their states are the same
            if set(get_states(child).items()) == states:
                for date, tip_names in getattr(child, TIPS_INSIDE).items():
                    getattr(n, TIPS_INSIDE)[date].extend(tip_names)

                n.remove_child(child)
                grandchildren = list(child.children)
                for grandchild in grandchildren:
                    n.add_child(grandchild)


def remove_mediators(tree, get_states):
    """
    Removes intermediate nodes that are just mediators between their parent and child states.
    :param get_states: a function that returns a set of node states
    :param tree: ete3.Tree
    :return: void, modifies the input tree
    """
    for n in tree.traverse('postorder'):
        if getattr(n, METACHILD, False) or n.is_leaf() or len(n.children) > 1 or not n.up:
            continue
        states = get_states(n)
        parent = n.up
        parent_states = get_states(parent)
        child = n.children[0]
        child_states = get_states(child)
        if set(states.keys()) == set(parent_states.keys()) == set(child_states.keys()):
            compatible = next((False for (key, state) in states.items()
                               # if the mediator has this key's state it should be the same as of its parent and child
                               if (state and (state != parent_states[key] or state != child_states[key]))
                               # otherwise it should hesitate between the parent's and the child's one
                               or (not state and parent_states[key] == child_states[key])), True)
        else:
            compatible = set(states.items()) == set(parent_states.items()) | set(child_states.items())
        if compatible:
            n_tips_inside = getattr(n, TIPS_INSIDE)[0] if isinstance(getattr(n, TIPS_INSIDE), list) \
                else getattr(n, TIPS_INSIDE)
            tips_inside = getattr(parent, TIPS_INSIDE)[0] if isinstance(getattr(parent, TIPS_INSIDE), list) \
                else getattr(parent, TIPS_INSIDE)
            for date, tip_names in n_tips_inside.items():
                tips_inside[date].extend(tip_names)

            parent.remove_child(n)
            parent.add_child(child)
