# -*- coding: utf-8 -*-

"""
Implements the free rigid monoidal category.

The objects are given by the free pregroup, the arrows by planar diagrams.

>>> unit, s, n = Ty(), Ty('s'), Ty('n')
>>> t = n.r @ s @ n.l
>>> assert t @ unit == t == unit @ t
>>> assert t.l.r == t == t.r.l
>>> left_snake, right_snake = Id(n.r).transpose_l(), Id(n.l).transpose_r()
>>> assert left_snake.normal_form() == Id(n) == right_snake.normal_form()
"""

from discopy import cat, monoidal, messages
from discopy.cat import AxiomError


class Ob(cat.Ob):
    """
    Implements simple pregroup types: basic types and their iterated adjoints.

    >>> a = Ob('a')
    >>> assert a.l.r == a.r.l == a and a != a.l.l != a.r.r
    """
    @property
    def z(self):
        """ Winding number """
        return self._z

    @property
    def l(self):
        """ Left adjoint """
        return Ob(self.name, self.z - 1)

    @property
    def r(self):
        """ Right adjoint """
        return Ob(self.name, self.z + 1)

    def __init__(self, name, z=0):
        if not isinstance(z, int):
            raise TypeError(messages.type_err(int, z))
        self._z = z
        super().__init__(name)

    def __eq__(self, other):
        if not isinstance(other, Ob):
            if isinstance(other, cat.Ob):
                return self.z == 0 and self.name == other.name
            return False
        return (self.name, self.z) == (other.name, other.z)

    def __hash__(self):
        return hash(self.name if not self.z else (self.name, self.z))

    def __repr__(self):
        return "Ob({}{})".format(
            repr(self.name), ", z=" + repr(self.z) if self.z else '')

    def __str__(self):
        return str(self.name) + (
            - self.z * '.l' if self.z < 0 else self.z * '.r')


class Ty(monoidal.Ty, Ob):
    """ Implements pregroup types as lists of simple types.

    >>> s, n = Ty('s'), Ty('n')
    >>> assert n.l.r == n == n.r.l
    >>> assert (s @ n).l == n.l @ s.l and (s @ n).r == n.r @ s.r
    """
    @property
    def l(self):
        """ Left adjoint. """
        return Ty(*[x.l for x in self.objects[::-1]])

    @property
    def r(self):
        """ Right adjoint. """
        return Ty(*[x.r for x in self.objects[::-1]])

    def tensor(self, other):
        return Ty(*super().tensor(other))

    def __init__(self, *t):
        t = [x if isinstance(x, Ob)
             else Ob(x.name) if isinstance(x, cat.Ob)
             else Ob(x) for x in t]
        monoidal.Ty.__init__(self, *t)
        Ob.__init__(self, str(self))

    def __getitem__(self, key):
        if isinstance(key, slice):
            return Ty(*super().__getitem__(key))
        return super().__getitem__(key)

    def __repr__(self):
        return "Ty({})".format(', '.join(
            repr(x if x.z else x.name) for x in self.objects))


class PRO(monoidal.PRO, Ty):
    """
    Objects of the free rigid monoidal category generated by 1.
    """
    @property
    def l(self):
        """
        >>> assert PRO(2).l == PRO(2)
        """
        return self

    @property
    def r(self):
        return self


class Diagram(monoidal.Diagram):
    """ Implements diagrams in the free rigid monoidal category.

    >>> I, n, s = Ty(), Ty('n'), Ty('s')
    >>> Alice, jokes = Box('Alice', I, n), Box('jokes', I, n.r @ s)
    >>> boxes, offsets = [Alice, jokes, Cup(n, n.r)], [0, 1, 0]
    >>> print(Diagram(Alice.dom @ jokes.dom, s, boxes, offsets))
    Alice >> Id(n) @ jokes >> Cup(n, n.r) @ Id(s)
    """
    @staticmethod
    def _upgrade(diagram):
        """
        Takes a monoidal.Diagram and returns a rigid.Diagram.
        """
        return Diagram(Ty(*diagram.dom), Ty(*diagram.cod),
                       diagram.boxes, diagram.offsets, layers=diagram.layers)

    @staticmethod
    def id(x):
        return Id(x)

    def then(self, other):
        return self._upgrade(super().then(other))

    def tensor(self, other):
        return self._upgrade(super().tensor(other))

    def interchange(self, i, j, left=False):
        return self._upgrade(super().interchange(i, j, left=left))

    def flatten(self):
        """
        >>> x = Ty('x')
        >>> f = Box('f', x, x)
        >>> assert isinstance((3 * f).foliation().flatten(), Diagram)
        """
        return self._upgrade(super().flatten())

    def foliate(self, yield_slices=False):
        """
        >>> x = Ty('x')
        >>> f = Box('f', x, x)
        >>> gen = (f @ Id(x) >> (f @ f)).foliate()
        >>> print(next(gen))
        f @ Id(x) >> Id(x) @ f >> f @ Id(x)
        """
        for diagram in super().foliate(yield_slices=yield_slices):
            if isinstance(diagram, cat.Arrow):
                yield self._upgrade(diagram)
            else:
                yield [self._upgrade(diagram[i]) for i in range(len(diagram))]

    def foliation(self):
        """
        >>> x = Ty('x')
        >>> f = Box('f', x, x)
        >>> assert isinstance((f @ Id(x) >> (f @ f)).foliation(), Diagram)
        """
        return self._upgrade(super().foliation())

    def __getitem__(self, key):
        """
        >>> n, s = Ty('n'), Ty('s')
        >>> cup, cap = Cup(n, n.r), Cap(n.r, n)
        >>> f, g, h = Box('f', n, n), Box('g', s @ n, n), Box('h', n, n @ s)
        >>> diagram = g @ cap >> f[::-1] @ Id(n.r) @ f >> cup @ h
        >>> assert isinstance(diagram[1:3], Diagram)
        >>> assert diagram[3:] == Id(n @ n.r) @ f >> cup @ h
        >>> assert diagram[:2] >> diagram[2:] == diagram
        """
        return self._upgrade(super().__getitem__(key))

    @staticmethod
    def cups(left, right):
        """ Constructs nested cups witnessing adjointness of x and y

        >>> a, b = Ty('a'), Ty('b')
        >>> assert Diagram.cups(a, a.r) == Cup(a, a.r)
        >>> assert Diagram.cups(a @ b, (a @ b).r) ==\\
        ...     Id(a) @ Cup(b, b.r) @ Id(a.r) >> Cup(a, a.r)
        """
        if not isinstance(left, Ty):
            raise TypeError(messages.type_err(Ty, left))
        if not isinstance(right, Ty):
            raise TypeError(messages.type_err(Ty, right))
        cups = Id(left @ right)
        for i in range(len(left)):
            j = len(left) - i - 1
            cup = Cup(left[j:j + 1], right[i:i + 1])
            cups = cups >> Id(left[:j]) @ cup @ Id(right[i + 1:])
        return cups

    @staticmethod
    def caps(left, right):
        """ Constructs nested cups witnessing adjointness of x and y

        >>> a, b = Ty('a'), Ty('b')
        >>> assert Diagram.caps(a, a.l) == Cap(a, a.l)
        >>> assert Diagram.caps(a @ b, (a @ b).l) == (Cap(a, a.l)
        ...                 >> Id(a) @ Cap(b, b.l) @ Id(a.l))
        """
        if not isinstance(left, Ty):
            raise TypeError(messages.type_err(Ty, left))
        if not isinstance(right, Ty):
            raise TypeError(messages.type_err(Ty, right))
        caps = Id(left @ right)
        for i in range(len(left)):
            j = len(left) - i - 1
            cap = Cap(left[j:j + 1], right[i:i + 1])
            caps = caps << Id(left[:j]) @ cap @ Id(right[i + 1:])
        return caps

    def transpose_r(self):
        """
        >>> a, b = Ty('a'), Ty('b')
        >>> double_snake = Id(a @ b).transpose_r()
        >>> two_snakes = Id(b).transpose_r() @ Id(a).transpose_r()
        >>> double_snake == two_snakes
        False
        >>> *_, two_snakes_nf = monoidal.Diagram.normalize(two_snakes)
        >>> assert double_snake == two_snakes_nf
        >>> f = Box('f', a, b)
        """
        return Diagram.caps(self.dom.r, self.dom) @ Id(self.cod.r)\
            >> Id(self.dom.r) @ self @ Id(self.cod.r)\
            >> Id(self.dom.r) @ Diagram.cups(self.cod, self.cod.r)

    def transpose_l(self):
        """
        >>> a, b = Ty('a'), Ty('b')
        >>> double_snake = Id(a @ b).transpose_l()
        >>> two_snakes = Id(b).transpose_l() @ Id(a).transpose_l()
        >>> double_snake == two_snakes
        False
        >>> *_, two_snakes_nf = monoidal.Diagram.normalize(
        ...     two_snakes, left=True)
        >>> assert double_snake == two_snakes_nf
        >>> f = Box('f', a, b)
        """
        return Id(self.cod.l) @ Diagram.caps(self.dom, self.dom.l)\
            >> Id(self.cod.l) @ self @ Id(self.dom.l)\
            >> Diagram.cups(self.cod.l, self.cod) @ Id(self.dom.l)

    def normalize(self, left=False):
        """
        Return a generator which yields normalization steps.

        >>> n, s = Ty('n'), Ty('s')
        >>> cup, cap = Cup(n, n.r), Cap(n.r, n)
        >>> f, g, h = Box('f', n, n), Box('g', s @ n, n), Box('h', n, n @ s)
        >>> diagram = g @ cap >> f[::-1] @ Id(n.r) @ f >> cup @ h
        >>> for d in diagram.normalize(): print(d)  # doctest: +ELLIPSIS
        g... >> Cup(n, n.r) @ Id(n)...
        g >> f[::-1] >> Id(n) @ Cap(n.r, n) >> Cup(n, n.r) @ Id(n) >> f >> h
        g >> f[::-1] >> f >> h
        """
        def follow_wire(diagram, i, j):
            """
            Given a diagram, the index of a box i and the offset j of an output
            wire, returns (i, j, obstructions) where:
            - i is the index of the box which takes this wire as input, or
            len(diagram) if it is connected to the bottom boundary.
            - j is the offset of the wire at its bottom end.
            - obstructions is a pair of lists of indices for the boxes on
            the left and right of the wire we followed.
            """
            left_obstruction, right_obstruction = [], []
            while i < len(diagram) - 1:
                i += 1
                box, off = diagram.boxes[i], diagram.offsets[i]
                if off <= j < off + len(box.dom):
                    return i, j, (left_obstruction, right_obstruction)
                if off <= j:
                    j += len(box.cod) - len(box.dom)
                    left_obstruction.append(i)
                else:
                    right_obstruction.append(i)
            return len(diagram), j, (left_obstruction, right_obstruction)

        def find_snake(diagram):
            """
            Given a diagram, returns (cup, cap, obstructions, left_snake)
            if there is a yankable pair, otherwise returns None.
            """
            for cap in range(len(diagram)):
                if not isinstance(diagram.boxes[cap], Cap):
                    continue
                for left_snake, wire in [(True, diagram.offsets[cap]),
                                         (False, diagram.offsets[cap] + 1)]:
                    cup, wire, obstructions =\
                        follow_wire(diagram, cap, wire)
                    not_yankable =\
                        cup == len(diagram)\
                        or not isinstance(diagram.boxes[cup], Cup)\
                        or left_snake and diagram.offsets[cup] + 1 != wire\
                        or not left_snake and diagram.offsets[cup] != wire
                    if not_yankable:
                        continue
                    return cup, cap, obstructions, left_snake
            return None

        def unsnake(diagram, cup, cap, obstructions, left_snake=False):
            """
            Given a diagram and the indices for a cup and cap pair
            and a pair of lists of obstructions on the left and right,
            returns a new diagram with the snake removed.

            A left snake is one of the form Id @ Cap >> Cup @ Id.
            A right snake is one of the form Cap @ Id >> Id @ Cup.
            """
            left_obstruction, right_obstruction = obstructions
            if left_snake:
                for box in left_obstruction:
                    diagram = diagram.interchange(box, cap)
                    yield diagram
                    for i, right_box in enumerate(right_obstruction):
                        if right_box < box:
                            right_obstruction[i] += 1
                    cap += 1
                for box in right_obstruction[::-1]:
                    diagram = diagram.interchange(box, cup)
                    yield diagram
                    cup -= 1
            else:
                for box in left_obstruction[::-1]:
                    diagram = diagram.interchange(box, cup)
                    yield diagram
                    for i, right_box in enumerate(right_obstruction):
                        if right_box > box:
                            right_obstruction[i] -= 1
                    cup -= 1
                for box in right_obstruction:
                    diagram = diagram.interchange(box, cap)
                    yield diagram
                    cap += 1
            boxes = diagram.boxes[:cap] + diagram.boxes[cup + 1:]
            offsets = diagram.offsets[:cap] + diagram.offsets[cup + 1:]
            layers = diagram.layers[:cap] >> diagram.layers[cup + 1:]
            yield Diagram(diagram.dom, diagram.cod, boxes, offsets, layers)

        diagram = self
        while True:
            yankable = find_snake(diagram)
            if yankable is None:
                break
            for _diagram in unsnake(diagram, *yankable):
                yield _diagram
                diagram = _diagram
        for _diagram in monoidal.Diagram.normalize(diagram, left=left):
            yield _diagram

    def normal_form(self, left=False):
        """
        Implements the normalisation of rigid monoidal categories,
        see arxiv:1601.05372, definition 2.12.
        """
        return super().normal_form(normalize=Diagram.normalize, left=left)


class Box(monoidal.Box, Diagram):
    """ Implements generators of rigid monoidal diagrams.

    >>> a, b = Ty('a'), Ty('b')
    >>> Box('f', a, b.l @ b, data={42})
    Box('f', Ty('a'), Ty(Ob('b', z=-1), 'b'), data={42})
    """
    def __init__(self, name, dom, cod, data=None, _dagger=False):
        """
        >>> a, b = Ty('a'), Ty('b')
        >>> Box('f', a, b.l @ b)
        Box('f', Ty('a'), Ty(Ob('b', z=-1), 'b'))
        """
        monoidal.Box.__init__(self, name, dom, cod, data=data, _dagger=_dagger)
        Diagram.__init__(self, dom, cod, [self], [0], layers=self.layers)


class Id(Diagram):
    """ Define an identity arrow in a free rigid category

    >>> t = Ty('a', 'b', 'c')
    >>> assert Id(t) == Diagram(t, t, [], [])
    """
    def __init__(self, t):
        super().__init__(t, t, [], [], layers=cat.Id(t))

    def __repr__(self):
        """
        >>> Id(Ty('n'))
        Id(Ty('n'))
        """
        return "Id({})".format(repr(self.dom))

    def __str__(self):
        """
        >>> n = Ty('n')
        >>> print(Id(n))
        Id(n)
        """
        return "Id({})".format(str(self.dom))


class Cup(Box):
    """ Defines cups for simple types.

    >>> n = Ty('n')
    >>> Cup(n, n.r)
    Cup(Ty('n'), Ty(Ob('n', z=1)))
    """
    def __init__(self, x, y):
        if not isinstance(x, Ty):
            raise TypeError(messages.type_err(Ty, x))
        if not isinstance(y, Ty):
            raise TypeError(messages.type_err(Ty, y))
        if x.r != y and x != y.r:
            raise AxiomError(messages.are_not_adjoints(x, y))
        if len(x) != 1 or len(y) != 1:
            raise ValueError(messages.cup_vs_cups(x, y))
        if x == y.r:
            raise NotImplementedError(messages.pivotal_not_implemented())
        super().__init__('CUP', x @ y, Ty())

    def dagger(self):
        raise NotImplementedError(messages.pivotal_not_implemented())

    def __repr__(self):
        """
        >>> n = Ty('n')
        >>> assert repr(Cup(n, n.r)) == "Cup(Ty('n'), Ty(Ob('n', z=1)))"
        """
        return "Cup({}, {})".format(repr(self.dom[:1]), repr(self.dom[1:]))

    def __str__(self):
        """
        >>> n = Ty('n')
        >>> assert str(Cup(n, n.r)) == "Cup(n, n.r)"
        """
        return "Cup({}, {})".format(self.dom[:1], self.dom[1:])


class Cap(Box):
    """ Defines cups for simple types.

    >>> n = Ty('n')
    >>> print(Cap(n, n.l).cod)
    n @ n.l
    >>> print(Cap(n.l, n.l.l).cod)
    n.l @ n.l.l
    """
    def __init__(self, x, y):
        if not isinstance(x, Ty):
            raise TypeError(messages.type_err(Ty, x))
        if not isinstance(y, Ty):
            raise TypeError(messages.type_err(Ty, y))
        if x != y.r and x.r != y:
            raise AxiomError(messages.are_not_adjoints(x, y))
        if len(x) != 1 or len(y) != 1:
            raise ValueError(messages.cap_vs_caps(x, y))
        if x.r == y:
            raise NotImplementedError(messages.pivotal_not_implemented())
        super().__init__('CAP', Ty(), x @ y)

    def dagger(self):
        raise NotImplementedError(messages.pivotal_not_implemented())

    def __repr__(self):
        """
        >>> n = Ty('n')
        >>> Cap(n, n.l)
        Cap(Ty('n'), Ty(Ob('n', z=-1)))
        """
        return "Cap({}, {})".format(repr(self.cod[:1]), repr(self.cod[1:]))

    def __str__(self):
        """
        >>> n = Ty('n')
        >>> print(Cap(n, n.l))
        Cap(n, n.l)
        """
        return "Cap({}, {})".format(self.cod[:1], self.cod[1:])


class Functor(monoidal.Functor):
    """
    Implements rigid monoidal functors, i.e. preserving cups and caps.

    >>> s, n = Ty('s'), Ty('n')
    >>> Alice, Bob = Box("Alice", Ty(), n), Box("Bob", Ty(), n)
    >>> loves = Box('loves', Ty(), n.r @ s @ n.l)
    >>> love_box = Box('loves', n @ n, s)
    >>> ob = {s: s, n: n}
    >>> ar = {Alice: Alice, Bob: Bob}
    >>> ar.update({loves: Cap(n.r, n) @ Cap(n, n.l)
    ...                   >> Id(n.r) @ love_box @ Id(n.l)})
    >>> F = Functor(ob, ar)
    >>> sentence = Alice @ loves @ Bob >> Cup(n, n.r) @ Id(s) @ Cup(n.l, n)
    >>> assert F(sentence).normal_form() == Alice >> Id(n) @ Bob >> love_box
    """
    def __init__(self, ob, ar, ob_factory=Ty, ar_factory=Diagram):
        """
        >>> F = Functor({Ty('x'): Ty('y')}, {})
        >>> F(Id(Ty('x')))
        Id(Ty('y'))
        """
        super().__init__(ob, ar, ob_factory=ob_factory, ar_factory=ar_factory)

    def __call__(self, diagram):
        """
        >>> x, y, z = Ty('x'), Ty('y'), Ty('z')
        >>> f, g = Box('f', x, y), Box('g', y, z)
        >>> F = Functor({x: y, y: z}, {f: g})
        >>> assert F(f.transpose_l()) == F(f).transpose_l()
        >>> assert F(f.transpose_r()) == F(f).transpose_r()
        """
        if isinstance(diagram, Ty):
            return sum([self(b) for b in diagram.objects], self.ob_factory())
        if isinstance(diagram, Ob) and not diagram.z:
            return self.ob[Ty(diagram.name)]
        if isinstance(diagram, Ob):
            result = self(Ob(diagram.name, z=0))
            if diagram.z < 0:
                for _ in range(-diagram.z):
                    result = result.l
            elif diagram.z > 0:
                for _ in range(diagram.z):
                    result = result.r
            return result
        if isinstance(diagram, Cup):
            return self.ar_factory.cups(
                self(diagram.dom[0]), self(diagram.dom[1]))
        if isinstance(diagram, Cap):
            return self.ar_factory.caps(
                self(diagram.cod[0]), self(diagram.cod[1]))
        if isinstance(diagram, Diagram):
            return super().__call__(diagram)
        raise TypeError(messages.type_err(Diagram, diagram))
