# Copyright 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Support for graphs."""

import collections

import common_util


class Node(object):
  """A node in a Graph.

  Nodes are identified within a graph using object identity.
  """
  def __init__(self):
    """Create a new node."""
    self.cost = 0

  def ToJsonDict(self):
    return common_util.SerializeAttributesToJsonDict({}, self, ['cost'])

  @classmethod
  def FromJsonDict(cls, json_dict):
    return common_util.DeserializeAttributesFromJsonDict(
        json_dict, cls(), ['cost'])


class Edge(object):
  """Represents an edge in a graph."""
  def __init__(self, from_node, to_node):
    """Creates an Edge.

    Args:
      from_node: (Node) Start node.
      to_node: (Node) End node.
    """
    self.from_node = from_node
    self.to_node = to_node
    self.cost = 0

  def ToJsonDict(self):
    return common_util.SerializeAttributesToJsonDict(
        {}, self, ['from_node', 'to_node', 'cost'])

  @classmethod
  def FromJsonDict(cls, json_dict):
    result = cls(None, None)
    return common_util.DeserializeAttributesFromJsonDict(
        json_dict, result, ['from_node', 'to_node', 'cost'])


class DirectedGraph(object):
  """Directed graph.

  A graph is identified by a list of nodes and a list of edges. It does not need
  to be acyclic, but then some methods will fail.
  """
  __GRAPH_NODE_INDEX = '__graph_node_index'
  __TO_NODE_INDEX = '__to_node_index'
  __FROM_NODE_INDEX = '__from_node_index'

  def __init__(self, nodes, edges):
    """Builds a graph from a set of node and edges.

    Note that the edges referencing a node not in the provided list are dropped.

    Args:
      nodes: ([Node]) Sequence of Nodes.
      edges: ([Edge]) Sequence of Edges.
    """
    self._nodes = set(nodes)
    self._edges = set(filter(
        lambda e: e.from_node in self._nodes and e.to_node in self._nodes,
        edges))
    assert all(isinstance(node, Node) for node in self._nodes)
    assert all(isinstance(edge, Edge) for edge in self._edges)
    self._in_edges = {n: [] for n in self._nodes}
    self._out_edges = {n: [] for n in self._nodes}
    for edge in self._edges:
      self._out_edges[edge.from_node].append(edge)
      self._in_edges[edge.to_node].append(edge)

  def OutEdges(self, node):
    """Returns a list of edges starting from a node.
    """
    return self._out_edges[node]

  def InEdges(self, node):
    """Returns a list of edges ending at a node."""
    return self._in_edges[node]

  def Nodes(self):
    """Returns the set of nodes of this graph."""
    return self._nodes

  def Edges(self):
    """Returns the set of edges of this graph."""
    return self._edges

  def RootNodes(self):
    """Returns an iterable of nodes that have no incoming edges."""
    return filter(lambda n: not self.InEdges(n), self._nodes)

  def UpdateEdge(self, edge, new_from_node, new_to_node):
    """Updates an edge.

    Args:
      edge:
      new_from_node:
      new_to_node:
    """
    assert edge in self._edges
    assert new_from_node in self._nodes
    assert new_to_node in self._nodes
    self._in_edges[edge.to_node].remove(edge)
    self._out_edges[edge.from_node].remove(edge)
    edge.from_node = new_from_node
    edge.to_node = new_to_node
    # TODO(lizeb): Check for duplicate edges?
    self._in_edges[edge.to_node].append(edge)
    self._out_edges[edge.from_node].append(edge)

  def TopologicalSort(self, roots=None):
    """Returns a list of nodes, in topological order.

      Args:
        roots: ([Node]) If set, the topological sort will only consider nodes
                        reachable from this list of sources.
    """
    sorted_nodes = []
    if roots is None:
      nodes_subset = self._nodes
    else:
      nodes_subset = self.ReachableNodes(roots)
    remaining_in_edges = {n: 0 for n in nodes_subset}
    for edge in self._edges:
      if edge.from_node in nodes_subset and edge.to_node in nodes_subset:
        remaining_in_edges[edge.to_node] += 1
    sources = [node for (node, count) in remaining_in_edges.items()
               if count == 0]
    while sources:
      node = sources.pop(0)
      sorted_nodes.append(node)
      for e in self.OutEdges(node):
        successor = e.to_node
        if successor not in nodes_subset:
          continue
        assert remaining_in_edges[successor] > 0
        remaining_in_edges[successor] -= 1
        if remaining_in_edges[successor] == 0:
          sources.append(successor)
    return sorted_nodes

  def ReachableNodes(self, roots, should_stop=lambda n: False):
    """Returns a list of nodes from a set of root nodes.

    Args:
      roots: ([Node]) List of roots to start from.
      should_stop: (callable) Returns True when a node should stop the
                   exploration and be skipped.
    """
    return self._ExploreFrom(
        roots, lambda n: (e.to_node for e in self.OutEdges(n)),
        should_stop=should_stop)

  def AncestorNodes(self, descendants):
    """Returns a set of nodes that are ancestors of a set of nodes.

    This is not quite the opposite of ReachableNodes, because (in a tree) it
    will not include |descendants|.

    Args:
      descendants: ([Node]) List of nodes to start from.

    """
    return set(self._ExploreFrom(
        descendants,
        lambda n: (e.from_node for e in self.InEdges(n)))) - set(descendants)

  def Cost(self, roots=None, path_list=None, costs_out=None):
    """Compute the cost of the graph.

    Args:
      roots: ([Node]) If set, only compute the cost of the paths reachable
             from this list of nodes.
      path_list: if not None, gets a list of nodes in the longest path.
      costs_out: if not None, gets a vector of node costs by node.

    Returns:
      Cost of the longest path.
    """
    if not self._nodes:
     return 0
    costs = {n: 0 for n in self._nodes}
    for node in self.TopologicalSort(roots):
      cost = 0
      if self.InEdges(node):
        cost = max([costs[e.from_node] + e.cost for e in self.InEdges(node)])
      costs[node] = cost + node.cost
    max_cost = max(costs.values())
    if costs_out is not None:
      del costs_out[:]
      costs_out.extend(costs)
    if path_list is not None:
      del path_list[:]
      node = (i for i in self._nodes if costs[i] == max_cost).next()
      path_list.append(node)
      while self.InEdges(node):
        predecessors = [e.from_node for e in self.InEdges(node)]
        node = reduce(
            lambda costliest_node, next_node:
            next_node if costs[next_node] > costs[costliest_node]
            else costliest_node, predecessors)
        path_list.insert(0, node)
    return max_cost

  def ToJsonDict(self):
    node_dicts = []
    node_to_index = {node: index for (index, node) in enumerate(self._nodes)}
    for (node, index) in node_to_index.items():
      node_dict = node.ToJsonDict()
      assert self.__GRAPH_NODE_INDEX not in node_dict
      node_dict.update({self.__GRAPH_NODE_INDEX: index})
      node_dicts.append(node_dict)
    edge_dicts = []
    for edge in self._edges:
      edge_dict = edge.ToJsonDict()
      assert self.__TO_NODE_INDEX not in edge_dict
      assert self.__FROM_NODE_INDEX not in edge_dict
      edge_dict.update({self.__TO_NODE_INDEX: node_to_index[edge.to_node],
                        self.__FROM_NODE_INDEX: node_to_index[edge.from_node]})
      edge_dicts.append(edge_dict)
    return {'nodes': node_dicts, 'edges': edge_dicts}

  @classmethod
  def FromJsonDict(cls, json_dict, node_class, edge_class):
    """Returns an instance from a dict.

    Note that the classes of the nodes and edges need to be specified here.
    This is done to reduce the likelihood of error.
    """
    index_to_node = {
        node_dict[cls.__GRAPH_NODE_INDEX]: node_class.FromJsonDict(node_dict)
        for node_dict in json_dict['nodes']}
    edges = []
    for edge_dict in json_dict['edges']:
      edge = edge_class.FromJsonDict(edge_dict)
      edge.from_node = index_to_node[edge_dict[cls.__FROM_NODE_INDEX]]
      edge.to_node = index_to_node[edge_dict[cls.__TO_NODE_INDEX]]
      edges.append(edge)
    result = DirectedGraph(index_to_node.values(), edges)
    return result

  def _ExploreFrom(self, initial, expand, should_stop=lambda n: False):
    """Explore from a set of nodes.

    Args:
      initial: ([Node]) List of nodes to start from.
      expand: (callable) Given a node, return an iterator of nodes to explore
        from that node.
      should_stop: (callable) Returns True when a node should stop the
                   exploration and be skipped.
    """
    visited = set()
    fifo = collections.deque([n for n in initial if not should_stop(n)])
    while fifo:
      node = fifo.pop()
      if should_stop(node):
        continue
      visited.add(node)
      for n in expand(node):
        if n not in visited and not should_stop(n):
          visited.add(n)
          fifo.appendleft(n)
    return list(visited)