# vim: set fileencoding=utf-8 sw=2 ts=2 et :
from __future__ import absolute_import
from __future__ import with_statement

from logging import getLogger

import networkx as NX
import yaml

from systems.collector import Aggregate, CResource
from systems.registry import get_registry
from systems.typesystem import EResource, Transition, ResourceRef

__all__ = ('Realizer', )


LOGGER = getLogger(__name__)


DESC_LIMIT = 64

def describe(thing):
  return '%s' % str(thing)[:DESC_LIMIT]

class CycleError(Exception):
  pass

class Node(object):
  def __init__(self):
    if type(self) == Node:
      raise TypeError

  def __repr__(self):
    return '<%s>' % self

  def __str__(self):
    return type(self).__name__

class CheckPointNode(Node):
  pass

class ExpandableNode(Node):
  def __init__(self, res):
    super(ExpandableNode, self).__init__()
    if type(self) == ExpandableNode:
      # Abstract class
      raise TypeError
    self._res = res

class BeforeExpandableNode(ExpandableNode):
  def __str__(self):
    return 'Before %s' % self._res

class AfterExpandableNode(ExpandableNode):
  def __str__(self):
    return 'After %s' % self._res

class GraphFirstNode(Node, yaml.YAMLObject):
  yaml_tag = u'GraphFirstNode'

class GraphLastNode(Node, yaml.YAMLObject):
  yaml_tag = u'GraphLastNode'

node_types = (CheckPointNode, BeforeExpandableNode, AfterExpandableNode,
    GraphFirstNode, GraphLastNode,
    Transition, Aggregate, CResource, EResource, ResourceRef)

class ResourceGraph(yaml.YAMLObject):
  """
  A graph of resources and transitions linked by dependencies.

  Resources are positioned as two sentinels in the transition graph.

  Invariant: directed, acyclic.
  """

  def __init__(self, top=None):
    self._graph = NX.DiGraph()
    self._first = GraphFirstNode()
    self._last = GraphLastNode()
    self._graph.add_edge(self._first, self._last)
    # Contains CResource and EResource, despite the name.
    # Used to enforce max one resource per id.
    self.__expandables = {}
    # Received references, by name.
    self.__received_refs = {}
    # What nodes were processed (meaning expanding or collecting)
    self.__processed = set()
    # Pre-bound args pased by ref. Allow putting extra depends on them.
    if top is not None:
      if not isinstance(top, ResourceGraph):
        raise TypeError(top, ResourceGraph)
      self.__top = top
    else:
      self.__top = self

  yaml_tag = u'!ResourceGraph'

  @classmethod
  def from_yaml(cls, loader, ynode):
    rg = cls()
    # Deep because of aliases and anchors, I think.
    mp = loader.construct_mapping(ynode, deep=True)
    pred_rels = mp['nodes']
    for rel in pred_rels:
      rg._add_node(rel['node'], depends=rel['depends'])
    return rg

  @classmethod
  def to_yaml(cls, dumper, rg):
    # This is incomplete.
    pred_rels = [{'node': node, 'depends': list(depends), }
        for (node, depends) in rg._iter_pred_rels()]
    return dumper.represent_mapping(cls.yaml_tag, {
      'nodes': pred_rels,
      })

  def _iter_node_preds(self, node0):
    return (node
        for node in self._graph.predecessors_iter(node0)
        if node not in (self._first, self._last))

  def _iter_pred_rels(self):
    return ((node, self._iter_node_preds(node))
      for node in self.sorted_nodes()
      if node not in (self._first, self._last))

  def sorted_nodes(self):
    return NX.topological_sort(self._graph)

  def sorted_transitions(self):
    return [n for n in self.sorted_nodes()
        if isinstance(n, Transition)]

  def iter_uncollected_resources(self):
    for nod in self._graph.nodes_iter():
      if isinstance(nod, CResource):
        if not nod in self.__processed:
          yield nod

  def iter_unexpanded_resources(self):
    for nod in self._graph.nodes_iter():
      if isinstance(nod, EResource):
        if not nod in self.__processed:
          yield nod

  def iter_unexpanded_aggregates(self):
    for agg in self._graph.nodes_iter():
      if isinstance(agg, Aggregate):
        if not agg in self.__processed:
          yield agg

  def iter_unprocessed(self):
    for nod in self.iter_uncollected_resources():
      yield nod
    for nod in self.iter_unexpanded_resources():
      yield nod
    for nod in self.iter_unexpanded_aggregates():
      yield nod

  def has_unprocessed(self):
    l = list(self.iter_unprocessed())
    return bool(l) # Tests for non-emptiness

  def require_acyclic(self):
    if not NX.is_directed_acyclic_graph(self._graph):
      # XXX NX doesn't have a 1-line method for listing those cycles
      raise CycleError

  def _add_node(self, node, depends=()):
    if not isinstance(node, node_types):
      raise TypeError(node, node_types)
    self._graph.add_node(node)
    self._graph.add_edge(self._first, node)
    self._graph.add_edge(node, self._last)
    for dep in depends:
      depn = self._intern(dep)
      self._add_node_dep(depn, node)
    return node

  def add_checkpoint(self, depends=()):
    return self._add_node(CheckPointNode(), depends)

  def add_transition(self, transition, depends=()):
    if not isinstance(transition, Transition):
      raise TypeError(transition, Transition)
    return self._add_node(transition, depends)

  def _add_aggregate(self, aggregate, depends=()):
    if not isinstance(aggregate, Aggregate):
      raise TypeError(aggregate, Aggregate)
    return self._add_node(aggregate, depends)

  def add_resource(self, resource, depends=()):
    """
    Add a resource.

    If an identical resource exists, it is returned.
    """

    if not isinstance(resource, (CResource, EResource)):
      raise TypeError(resource, (CResource, EResource))

    if resource.identity in self.__expandables:
      # We have this id already.
      # Either it's the exact same resource, or a KeyError is thrown.
      resource = self._intern(resource)
      # XXX Need to bypass _intern for already expanded.
      # XXX When we use add_to_top, we sometimes have to deal
      # with a resource that's already been expanded.
      # Those are not in the graph anymore. How do we refer to them?
    else:
      self.__expandables[resource.identity] = resource
    # Even if already there, we need to add the depends.
    resource = self._add_node(resource, depends)
    # If already there, notice we aliase it.
    return self.make_ref(resource)

  def make_ref(self, res, depends=()):
    res = self._intern(res)
    if not isinstance(res, (CResource, EResource)):
      raise TypeError(res, (CResource, EResource))
    depends = list(depends)
    depends.append(res)
    return self._add_node(ResourceRef(res), depends)

  def make_alias_ref(self, ref, depends=()):
    ref = self._intern(ref)
    if not isinstance(ref, ResourceRef):
      raise TypeError(ref, ResourceRef)
    depends = list(depends)
    depends.append(ref)
    return self._add_node(ResourceRef(ref.unref), depends)

  def add_to_top(self, res):
    """
    Add a resource to the top ResourceGraph.

    Use it to put things that you don't necessarily
    want to be after the outside dependencies the current graph has.
    """

    ref = self.__top.add_resource(res)
    return self._add_node(ref)

  def _add_node_dep(self, node0, node1):
    if not isinstance(node0, node_types):
      raise TypeError(node0, node_types)
    if not isinstance(node1, node_types):
      raise TypeError(node1, node_types)
    if not self._graph.has_node(node0):
      raise KeyError(node0)
    if not self._graph.has_node(node1):
      raise KeyError(node1)
    if self._graph.has_edge(node0, node1):
      return False
    if node0 == node1:
      # Disallow self-loops to keep acyclic invariant.
      # Also they don't make sense.
      raise ValueError(node0)
    # Invariant check
    rev_path = NX.shortest_path(self._graph, node1, node0)
    if rev_path is not False:
      raise CycleError(rev_path)
    self._graph.add_edge(node0, node1)
    return True

  def _intern(self, thing):
    if not isinstance(thing, node_types):
      raise TypeError
    if thing not in self._graph:
      raise KeyError(thing)
    return thing

  def add_dependency(self, elem0, elem1):
    node0 = self._intern(elem0)
    node1 = self._intern(elem1)
    return self._add_node_dep(node0, node1)

  def _is_direct_rconnect(self, r0, r1):
    s0 = self._intern(r0)
    s1 = self._intern(r1)
    # shortest_path is also a test for connectedness.
    return bool(NX.shortest_path(self._graph, s0, s1))

  def resources_connected(self, r0, r1):
    return self._is_direct_rconnect(r0, r1) \
        or self._is_direct_rconnect(r1, r0)

  def draw(self, fname):
    return self.draw_agraph(fname)

  def draw_agraph(self, fname):
    # XXX pygraphviz has steep dependencies (x11 libs)
    # and recommends (texlive) for a headless box.

    # We duplicate the graph, otherwise networkx / pygraphviz
    # would make a lossy conversion (sometimes refusing to convert), by adding
    # nodes as their string representation. Madness, I know.
    gr2 = NX.create_empty_copy(self._graph, False)
    for node in self._graph.nodes_iter():
      gr2.add_node(id(node))
    for (n0, n1) in self._graph.edges_iter():
      gr2.add_edge(id(n0), id(n1))
    names = dict((id(node), { 'label': describe(node)})
        for node in self._graph.nodes_iter())
    gr2.delete_node(id(self._first))
    gr2.delete_node(id(self._last))
    g = NX.to_agraph(gr2, {
        'graph': {
          'nodesep': '0.2',
          'rankdir': 'TB',
          'ranksep': '0.5',
          },
        'node': {
          'shape': 'box',
          },
        },
        names)
    g.write(fname + '.dot')
    # Dot is good for DAGs.
    g.layout(prog='dot')
    g.draw(fname + '.svg')
    with open(fname + '.yaml', 'w') as f:
      yaml.dump(self, f)
    # Fails with the expanded graph, due to instancemethod
    #yaml.load(yaml.dump(self))

  def draw_matplotlib(self, fname):
    # Pyplot is stateful and awkward to use.
    import matplotlib.pyplot as P
    # Disable hold or it definitely won't work (probably a bug).
    P.hold(False)
    NX.draw(self._graph)
    P.savefig(fname)

  def collect_resources(self, r0s, r1):
    """
    Replace an iterable of resources with one new resource.

    May break the acyclic invariant, caveat emptor.
    """

    # The invariant is kept iff the r0s don't have paths linking them.
    # For our use case (collectors), we could allow paths provided they are
    # internal to r0s. This introduces self-loops that we would then remove.

    for r0 in r0s:
      r0 = self._intern(r0)
      if r0 in self.__processed:
        raise RuntimeError

    if r1 in self._graph:
      raise ValueError(r1)
    r1 = self._add_aggregate(r1)

    for r0 in r0s:
      r0 = self._intern(r0)
      self._move_edges(r0, r1)
      self.__processed.add(r0)
    self.require_acyclic()

  def _move_edges(self, n0, n1):
    if n0 == n1:
      raise RuntimeError
    n0 = self._intern(n0)
    n1 = self._intern(n1)
    # list is used as a temporary
    # add after delete in case of same.
    for pred in list(self._graph.predecessors_iter(n0)):
      self._graph.delete_edge(pred, n0)
      self._graph.add_edge(pred, n1)
    for succ in list(self._graph.successors_iter(n0)):
      self._graph.delete_edge(n0, succ)
      self._graph.add_edge(n1, succ)
    self._graph.delete_node(n0)
    # Can't undo. Invariant will stay broken.

  def _split_node(self, res):
    res = self._intern(res)
    before = self._add_node(BeforeExpandableNode(res))
    after = self._add_node(AfterExpandableNode(res))
    self._graph.add_edge(before, after)
    for pred in list(self._graph.predecessors_iter(res)):
      self._graph.delete_edge(pred, res)
      self._graph.add_edge(pred, before)
    for succ in list(self._graph.successors_iter(res)):
      self._graph.delete_edge(res, succ)
      self._graph.add_edge(after, succ)
    self._graph.delete_node(res)
    return before, after

  def _receive_by_ref(self, name, ref):
    if name in self.__received_refs:
      raise RuntimeError(name, ref)
    ref = self._add_node(ref)
    self.__received_refs[name] = ref
    return ref

  def _pass_by_ref(self, subgraph, name, ref):
    # The origin/value distinction is important
    # for aliased arguments (two refs, same val).
    ref = self._intern(ref)
    if not isinstance(ref, ResourceRef):
      raise TypeError(ref, ResourceRef)

    subgraph._receive_by_ref(name, ref)

  def expand_resource(self, res):
    """
    Replace res by a small resource graph.

    The resource_graph is inserted in the main graph
    between the sentinels that represent the resource.
    """

    res = self._intern(res)

    # We're processing from the outside in.
    if res in self.__processed:
      raise RuntimeError

    resource_graph = ResourceGraph(self.__top)
    if isinstance(res, EResource):
      for (name, ref) in res.iter_passed_by_ref():
        # ref will be present in both graphs.
        self._pass_by_ref(resource_graph, name, ref)
    elif isinstance(res, Aggregate):
      pass
    else:
      raise TypeError(res)
    res.expand_into(resource_graph)
    # We expand from the outside in
    if bool(resource_graph.__processed):
      raise RuntimeError

    # Do not skip sentinels.
    for n in resource_graph._graph.nodes_iter():
      self._add_node(n)
    for (n0, n1) in resource_graph._graph.edges_iter():
      self._add_node_dep(n0, n1)

    for (id1, res1) in resource_graph.__expandables.iteritems():
      # We expand from the outside in.
      assert res1 not in self.__processed
      if id1 in self.__expandables:
        # Pass by reference if you must use the same resource
        # in different contexts.
        raise RuntimeError('ResourceBase collision.', res, res1)
      else:
        self.__expandables[id1] = res1

    before, after = self._split_node(res)
    self.__processed.add(res)
    self._move_edges(resource_graph._first, before)
    self._move_edges(resource_graph._last, after)
    # What may break the invariant:
    # Passing a ref to res, and making res depend on ref.
    # ref ends up on both sides of ref.before.
    self.require_acyclic()


class Realizer(object):
  """
  A graph of realizables linked by dependencies.
  """

  def __init__(self, expandable):
    self.__resources = ResourceGraph()
    self.__expandable = expandable
    self.__state = 'init'

  def require_state(self, state):
    """
    Raise an exception if we are not in the required state.
    """

    if self.__state != state:
      raise RuntimeError(u'Realizer state should be «%s»' % state)

  def ensure_frozen(self):
    """
    Build the finished dependency graph.

    Merge identical realizables, collect what can be.
    """

    if self.__state == 'frozen':
      return
    # Order is important
    self.require_state('init')
    self.__expandable.expand_into(self.__resources)
    #self.__resources.draw('/tmp/freezing')
    self._expand()
    #self.__resources.draw('/tmp/pre-collect')
    self._collect()
    self._expand_aggregates()
    assert not bool(list(self.__resources.iter_unprocessed()))
    self.__state = 'frozen'
    #self.__resources.draw('/tmp/frozen')

  def _collect(self):
    # Collects compatible nodes into merged nodes.

    def can_merge(part0, part1):
      for n0 in part0:
        for n1 in part1:
          if self.__resources.resources_connected(n0, n1):
            return False
      return True

    def possibly_merge(partition):
      # Merge once if possible. Return true if did merge.
      e = dict(enumerate(partition))
      n = len(partition)
      # Loop over the triangle of unordered pairs
      for i in xrange(n):
        for j in xrange(i + 1, n):
          part0, part1 = e[i], e[j]
          if can_merge(part0, part1):
            partition.add(part0.union(part1))
            partition.remove(part0)
            partition.remove(part1)
            return True
      return False

    reg = get_registry()
    for collector in reg.collectors:
      # Pre-partition is made of parts acceptable for the collector.
      pre_partition = collector.partition(
          [r for r in self.__resources.iter_uncollected_resources()
            if collector.filter(r)])
      for part in pre_partition:
        # Collector parts are split again, the sub-parts are merged
        # when dependencies allow.
        # Not a particularly efficient algorithm, just simple.
        # Gives one solution among many possibilities.
        partition = set(frozenset((r, ))
            for r in part
            for part in pre_partition)
        while possibly_merge(partition):
          pass

        # Let the collector handle the rest
        for part in partition:
          if not bool(part):
            # Test for emptiness.
            # Aggregate even singletons.
            continue
          merged = collector.collect(part)
          self.__resources.collect_resources(part, merged)
    assert not bool(list(self.__resources.iter_uncollected_resources()))

  def _expand(self):
    # Poor man's recursion
    while True:
      fresh = set(r
          for r in self.__resources.iter_unexpanded_resources())
      if bool(fresh) == False: # Test for emptiness
        break
      for r in fresh:
        self.__resources.expand_resource(r)
    assert not bool(list(self.__resources.iter_unexpanded_resources()))

  def _expand_aggregates(self):
    for a in list(self.__resources.iter_unexpanded_aggregates()):
      self.__resources.expand_resource(a)
    assert not bool(list(self.__resources.iter_unexpanded_aggregates()))
    # Enforce the rule that aggregates can only expand into transitions.
    if self.__resources.has_unprocessed():
      raise RuntimeError(list(self.__resources.iter_unprocessed()))

  def realize(self):
    """
    Realize all realizables and transitions in dependency order.
    """

    self.ensure_frozen()
    for t in self.__resources.sorted_transitions():
      t.realize()
    self.__state = 'realized'


