from __future__ import absolute_import, print_function, division
from collections import defaultdict
from six import iteritems
from theano.gof.graph import list_of_nodes
from theano.compat import cmp

# {{{ http://code.activestate.com/recipes/578231/ (r1)
# Copyright (c) Oren Tirosh 2012
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


def memodict(f):
    """
    Memoization decorator for a function taking a single argument.

    """
    class memodict(defaultdict):
        def __missing__(self, key):
            ret = self[key] = f(key)
            return ret
    return memodict().__getitem__

# end of http://code.activestate.com/recipes/578231/ }}}


def make_depends():
    @memodict
    def depends(pair):
        """
        Returns True if a depends on b.

        """
        a, b = pair
        return (any(bout in a.inputs for bout in b.outputs) or
                any(depends((ainp.owner, b)) for ainp in a.inputs
                    if ainp.owner))
    return depends


def make_dependence_cmp():
    """
    Create a comparator to represent the dependence of nodes in a graph.

    """
    depends = make_depends()

    def dependence(a, b):
        """
        A cmp function for nodes in a graph - does a depend on b?

        Returns
        -------
        int
            Positive number if a depends on b, negative number
            if b depends on a, 0 otherwise.

        """
        if depends((a, b)):
            return 1
        if depends((b, a)):
            return -1
        return 0

    return dependence


def reverse_dict(d):
    """
    Reverses direction of dependence dict.

    Notes
    -----
    dict order is not deterministic. As we iterate on the
    input dict, it makes the output of this function depend on the
    dict order. So this function output order should be considered
    as undeterministic.

    Examples
    --------
    >>> d = {'a': (1, 2), 'b': (2, 3), 'c':()}
    >>> reverse_dict(d)
    {1: ('a',), 2: ('a', 'b'), 3: ('b',)}

    """
    result = {}
    for key in d:
        for val in d[key]:
            result[val] = result.get(val, tuple()) + (key, )
    return result


def _toposort(edges):
    """
    Topological sort algorithm by Kahn [1] - O(nodes + vertices).

    Parameters
    ----------
    edges
        A dict of the form {a: {b, c}} where b and c depend on a.

    Returns
    -------
    L : list
        An ordered list of nodes that satisfy the dependencies of edges.

    Closely follows the wikipedia page [2]

    References
    ----------
    [1] Kahn, Arthur B. (1962), "Topological sorting of large networks",
    Communications of the ACM
    [2] http://en.wikipedia.org/wiki/Toposort#Algorithms

    Examples
    --------
    >>> _toposort({1: {2, 3}, 2: (3, )})
    [1, 2, 3]

    """
    incoming_edges = reverse_dict(edges)
    incoming_edges = dict((k, set(val))
                          for k, val in iteritems(incoming_edges))
    S = set((v for v in edges if v not in incoming_edges))
    L = []

    while S:
        n = S.pop()
        L.append(n)
        for m in edges.get(n, ()):
            assert n in incoming_edges[m]
            incoming_edges[m].remove(n)
            if not incoming_edges[m]:
                S.add(m)
    if any(incoming_edges.get(v, None) for v in edges):
        raise ValueError("Input has cycles")
    return L


def posort(l, *cmps):
    """
    Partially ordered sort with multiple comparators.

    Given a list of comparators, orders the elements in l so that the
    comparators are satisfied as much as possible giving precedence to
    earlier comparators.

    Parameters
    ----------
    l
        An iterable of nodes in a graph.
    cmps
        A sequence of comparator functions that describe which nodes should
        come before which others.

    Returns
    -------
    list
        A list of nodes which satisfy the comparators as much as possible.

    Notes
    -----
    Implemented with _toposort.

    Examples
    --------
    >>> lower_tens = lambda a, b: a/10 - b/10 # prefer lower numbers div 10
    >>> prefer evens = lambda a, b: a%2 - b%2 # prefer even numbers
    >>> posort(list(range(20)), lower_tens, prefer_evens)
    [0, 8, 2, 4, 6, 1, 3, 5, 7, 9, 16, 18, 10, 12, 14, 17, 19, 11, 13, 15]

    """
    comes_before = dict((a, set()) for a in l)
    comes_after = dict((a, set()) for a in l)

    def add_links(a, b):  # b depends on a
        comes_after[a].add(b)
        comes_after[a].update(comes_after[b])
        for c in comes_before[a]:
            comes_after[c].update(comes_after[a])
        comes_before[b].add(a)
        comes_before[b].update(comes_before[a])
        for c in comes_after[b]:
            comes_before[c].update(comes_before[b])

    def check():
        """
        Tests for cycles in manufactured edges.

        """
        for a in l:
            for b in l:
                assert not(b in comes_after[a] and a in comes_after[b])

    for cmp_fn in cmps:
        for a in l:
            for b in l:
                if cmp_fn(a, b) < 0:  # a wants to come before b
                    # if this wouldn't cause a cycle and isn't already known
                    if b not in comes_before[a] and b not in comes_after[a]:
                        add_links(a, b)
    # check() # debug code

    return _toposort(comes_after)


def sort_apply_nodes(inputs, outputs, cmps):
    """
    Order a graph of apply nodes according to a list of comparators.

    The following example sorts first by dependence of nodes (this is a
    topological sort) and then by lexicographical ordering (nodes that start
    with 'E' come before nodes that start with 'I' if there is no dependence.

    Examples
    --------
    >>> from theano.gof.graph import sort_apply_nodes, dependence
    >>> from theano.tensor import matrix, dot
    >>> x = matrix('x')
    >>> y = dot(x*2, x+1)
    >>> str_cmp = lambda a, b: cmp(str(a), str(b)) # lexicographical sort
    >>> sort_apply_nodes([x], [y], cmps=[dependence, str_cmp])
    [Elemwise{add,no_inplace}(x, InplaceDimShuffle{x,x}.0),
     InplaceDimShuffle{x,x}(TensorConstant{2}),
     Elemwise{mul,no_inplace}(x, InplaceDimShuffle{x,x}.0),
     InplaceDimShuffle{x,x}(TensorConstant{1}),
     dot(Elemwise{mul,no_inplace}.0, Elemwise{add,no_inplace}.0)]

    """
    return posort(list_of_nodes(inputs, outputs), *cmps)


def sort_schedule_fn(*cmps):
    """
    Make a schedule function from comparators.

    See Also
    --------
    sort_apply_nodes

    """
    dependence = make_dependence_cmp()
    cmps = (dependence,) + cmps

    def schedule(fgraph):
        """
        Order nodes in a FunctionGraph.

        """
        return sort_apply_nodes(fgraph.inputs, fgraph.outputs, cmps)
    return schedule


def key_to_cmp(key):
    """
    comparator function based on "key" function
    """
    def key_cmp(a, b):
        return cmp(key(a), key(b))
    return key_cmp