# Copyright (C) 2005, 2008 Canonical Ltd
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA


def max_distance(node, ancestors, distances, root_descendants):
    """Calculate the max distance to an ancestor.  
    Return None if not all possible ancestors have known distances"""
    best = None
    if node in distances:
        best = distances[node]
    for ancestor in ancestors[node]:
        # skip ancestors we will never traverse:
        if root_descendants is not None and ancestor not in root_descendants:
            continue
        # An ancestor which is not listed in ancestors will never be in
        # distances, so we pretend it never existed.
        if ancestor not in ancestors:
            continue
        if ancestor not in distances:
            return None
        if best is None or distances[ancestor]+1 > best:
            best = distances[ancestor] + 1
    return best

    
def node_distances(graph, ancestors, start, root_descendants=None):
    """Produce a list of nodes, sorted by distance from a start node.
    This is an algorithm devised by Aaron Bentley, because applying Dijkstra
    backwards seemed too complicated.

    For each node, we walk its descendants.  If all the descendant's ancestors
    have a max-distance-to-start, (excluding ones that can never reach start),
    we calculate their max-distance-to-start, and schedule their descendants.

    So when a node's last parent acquires a distance, it will acquire a
    distance on the next iteration.

    Once we know the max distances for all nodes, we can return a list sorted
    by distance, farthest first.
    """
    distances = {start: 0}
    lines = set([start])
    while len(lines) > 0:
        new_lines = set()
        for line in lines:
            line_descendants = graph[line]
            assert line not in line_descendants, "%s refers to itself" % line
            for descendant in line_descendants:
                distance = max_distance(descendant, ancestors, distances,
                                        root_descendants)
                if distance is None:
                    continue
                distances[descendant] = distance
                new_lines.add(descendant)
        lines = new_lines
    return distances

def nodes_by_distance(distances):
    """Return a list of nodes sorted by distance"""
    def by_distance(n):
        return distances[n],n

    node_list = distances.keys()
    node_list.sort(key=by_distance, reverse=True)
    return node_list

def select_farthest(distances, common):
    """Return the farthest common node, or None if no node qualifies."""
    node_list = nodes_by_distance(distances)
    for node in node_list:
        if node in common:
            return node

def all_descendants(descendants, start):
    """Produce a set of all descendants of the start node.
    The input is a map of node->list of descendants for a graph encompassing
    start.
    """
    result = set()
    lines = set([start])
    while len(lines) > 0:
        new_lines = set()
        for line in lines:
            if line not in descendants:
                continue
            for descendant in descendants[line]:
                if descendant in result:
                    continue
                result.add(descendant)
                new_lines.add(descendant)
        lines = new_lines
    return result


class Graph(object):
    """A graph object which can memoise and cache results for performance."""

    def __init__(self):
        super(Graph, self).__init__()
        self.roots = set([])
        self.ghosts = set([])
        self._graph_ancestors = {}
        self._graph_descendants = {}

    def add_ghost(self, node_id):
        """Add a ghost to the graph."""
        self.ghosts.add(node_id)
        self._ensure_descendant(node_id)

    def add_node(self, node_id, parent_ids):
        """Add node_id to the graph with parent_ids as its parents."""
        if len(parent_ids) == 0:
            self.roots.add(node_id)
        self._graph_ancestors[node_id] = list(parent_ids)
        self._ensure_descendant(node_id)
        for parent in parent_ids:
            self._ensure_descendant(parent)
            self._graph_descendants[parent][node_id] = 1
        
    def _ensure_descendant(self, node_id):
        """Ensure that a descendant lookup for node_id will work."""
        if not node_id in self._graph_descendants:
            self._graph_descendants[node_id] = {}

    def get_ancestors(self):
        """Return a dictionary of graph node:ancestor_list entries."""
        return dict(self._graph_ancestors.items())

    def get_ancestry(self, node_id, topo_sorted=True):
        """Return the inclusive ancestors of node_id in topological order."""
        # maybe optimise this ?
        from bzrlib.tsort import topo_sort
        result = {}
        pending = set([node_id])
        while len(pending):
            current = pending.pop()
            parents = self._graph_ancestors[current]
            parents = [parent for parent in parents if parent not in self.ghosts]
            result[current] = parents
            for parent in parents:
                if parent not in result and parent not in pending:
                    pending.add(parent)
        if not topo_sorted:
            return result.keys()
        return topo_sort(result.items())

    def get_descendants(self):
        """Return a dictionary of graph node:child_node:distance entries."""
        return dict(self._graph_descendants.items())