finite_set_map_cy.pyx

"""
Data structures for maps between finite sets

This module implements several fast Cython data structures for maps
between two finite set. Those classes are not intended to be used
directly. Instead, such a map should be constructed via its parent,
using the class :class:`~sage.sets.finite_set_maps.FiniteSetMaps`.

EXAMPLES:

To create a map between two sets, one first creates the set of such maps::

    sage: M = FiniteSetMaps(["a", "b"], [3, 4, 5])

The map can then be constructed either from a function::

    sage: f1 = M(lambda c: ord(c)-94); f1
    map: a -> 3, b -> 4

or from a dictionary::

    sage: f2 = M.from_dict({'a':3, 'b':4}); f2
    map: a -> 3, b -> 4

The two created maps are equal::

    sage: f1 == f2
    True

Internally, maps are represented as the list of the ranks of the
images ``f(x)`` in the co-domain, in the order of the domain::

    sage: list(f2)
    [0, 1]

A third fast way to create a map it to use such a list. it should be
kept for internal use::

    sage: f3 = M._from_list_([0, 1]); f3
    map: a -> 3, b -> 4
    sage: f1 == f3
    True

AUTHORS:

- Florent Hivert
"""
#*****************************************************************************
#       Copyright (C) 2010 Florent Hivert <[email protected]>,
#
#  Distributed under the terms of the GNU General Public License (GPL)
#                  http://www.gnu.org/licenses/
#*****************************************************************************

import sage
from sage.structure.list_clone cimport ClonableIntArray
from sage.structure.parent cimport Parent
from sage.structure.element cimport generic_power_c
from sage.sets.set import Set_object_enumerated


cpdef fibers(f, domain):
    r"""
    Returns the fibers of the function ``f`` on the finite set ``domain``

    INPUT:

    - ``f``      -- a function or callable
    - ``domain`` -- a finite iterable

    OUTPUT:

    - a dictionary ``d`` such that ``d[y]`` is the set of all ``x`` in
      ``domain`` such that ``f(x) = y``

    EXAMPLES::

        sage: from sage.sets.finite_set_map_cy import fibers, fibers_args
        sage: fibers(lambda x: 1, [])
        {}
        sage: fibers(lambda x: x^2, [-1, 2, -3, 1, 3, 4])
        {16: {4}, 1: {1, -1}, 4: {2}, 9: {3, -3}}
        sage: fibers(lambda x: 1,   [-1, 2, -3, 1, 3, 4])
        {1: {1, 2, 3, 4, -3, -1}}
        sage: fibers(lambda x: 1, [1,1,1])
        {1: {1}}

    .. seealso:: :func:`fibers_args` if one needs to pass extra
       arguments to ``f``.
    """
    result = {}
    for x in domain:
        y = f(x)
        result.setdefault(y, set()).add(x)
    for x, v in result.iteritems():
        result[x] = Set_object_enumerated(v)
    return result


def fibers_args(f, domain, *args, **opts):
    r"""
    Returns the fibers of the function ``f`` on the finite set ``domain``

    It is the same as :func:`fibers` except that one can pass extra
    argument for ``f`` (with a small overhead)

    EXAMPLES::

        sage: from sage.sets.finite_set_map_cy import fibers_args
        sage: fibers_args(operator.pow, [-1, 2, -3, 1, 3, 4], 2)
        {16: {4}, 1: {1, -1}, 4: {2}, 9: {3, -3}}
    """
    return fibers(lambda x: f(x, *args, **opts), domain)


cdef class FiniteSetMap_MN(ClonableIntArray):
    """
    Data structure for maps from ``range(m)`` to ``range(n)``.

    We assume that the parent given as argument is such that:

    - ``m`` is stored in ``self.parent()._m``
    - ``n`` is stored in ``self.parent()._n``
    - the domain is in ``self.parent().domain()``
    - the codomain is in ``self.parent().codomain()``
    """

    def __call__(self, int i):
        """
        Returns the image of ``i`` under the map ``self``

        INPUT:

        - ``i`` -- an int

        OUTPUT:

        an int

        EXAMPLES::

            sage: fs = FiniteSetMaps(4, 3)([1, 0, 2, 1])
            sage: fs(0), fs(1), fs(2), fs(3)
            (1, 0, 2, 1)
        """
        return self._getitem(i)

    # Needed by generic power which refuses to compute 0^0
    def __nonzero__(self):
        """
        Returns whether ``self`` is non zero; this is always ``True``.

        EXAMPLES::

            sage: fs = FiniteSetMaps(4, 3)([1, 0, 2, 1])
            sage: bool(fs)
            True
        """
        return True

    cpdef domain(self):
        """
        Returns the domain of ``self``

        EXAMPLES::

            sage: FiniteSetMaps(4, 3)([1, 0, 2, 1]).domain()
            {0, .., 3}
        """
        return self._parent.domain()

    cpdef codomain(self):
        """
        Returns the codomain of ``self``

        EXAMPLES::

            sage: FiniteSetMaps(4, 3)([1, 0, 2, 1]).codomain()
            {0, .., 2}
        """
        return self._parent.codomain()

    cpdef _setimage(self, int i, int j):
        """
        Set the image of ``i`` as ``j`` in ``self``

        This is a fast internal version where ``i`` and ``j`` are both int's.

        .. warning:: ``self`` must be mutable; otherwise an exception is raised.

        INPUT:

        - ``i``, ``j`` -- two ``int``'s

        OUTPUT: ``None``

        EXAMPLES::

            sage: fs = FiniteSetMaps(4, 3)([1, 0, 2, 1])
            sage: fs2 = copy(fs)
            sage: fs2._setimage(2, 1)
            sage: fs2
            [1, 0, 1, 1]
            sage: with fs.clone() as fs3:
            ...       fs3._setimage(0, 2)
            ...       fs3._setimage(1, 2)
            sage: fs3
            [2, 2, 2, 1]

        TESTS::

            sage: with fs.clone() as fs3:
            ...       fs3._setimage(6, 2)
            Traceback (most recent call last):
            ...
            IndexError: list index out of range
            sage: with fs.clone() as fs3:
            ...       fs3._setimage(1, 4)
            Traceback (most recent call last):
            ...
            AssertionError: Wrong value self(1) = 4
        """
        self._setitem(i, j)

    cpdef _getimage(self, int i):
        """
        Returns the image of ``i`` by ``self``

        This is a fast internal version where ``i`` is an int.

        INPUT:

        - ``i`` -- an ``int``

        EXAMPLES::

            sage: fs = FiniteSetMaps(4, 3)([1, 0, 2, 1])
            sage: fs._getimage(0), fs._getimage(1), fs._getimage(2), fs._getimage(3)
            (1, 0, 2, 1)
        """
        return self._getitem(i)

    cpdef setimage(self, i, j):
        """
        Set the image of ``i`` as ``j`` in ``self``

        .. warning:: ``self`` must be mutable; otherwise an exception is raised.

        INPUT:

        - ``i``, ``j`` -- two ``object``'s

        OUTPUT: ``None``

        .. note:: if you need speed, please use instead :meth:`_setimage`

        EXAMPLES::

            sage: fs = FiniteSetMaps(4, 3)([1, 0, 2, 1])
            sage: fs2 = copy(fs)
            sage: fs2.setimage(2, 1)
            sage: fs2
            [1, 0, 1, 1]
            sage: with fs.clone() as fs3:
            ...       fs3.setimage(0, 2)
            ...       fs3.setimage(1, 2)
            sage: fs3
            [2, 2, 2, 1]
        """
        self._setitem(int(i), int(j))

    cpdef getimage(self, i):
        """
        Returns the image of ``i`` by ``self``

        INPUT:

        - ``i`` -- any object.

        .. note:: if you need speed, please use instead :meth:`_getimage`

        EXAMPLES::

            sage: fs = FiniteSetMaps(4, 3)([1, 0, 2, 1])
            sage: fs.getimage(0), fs.getimage(1), fs.getimage(2), fs.getimage(3)
            (1, 0, 2, 1)
        """
        return self._getitem(int(i))

    cpdef image_set(self):
        """
        Returns the image set of ``self``

        EXAMPLES::

            sage: FiniteSetMaps(4, 3)([1, 0, 2, 1]).image_set()
            {0, 1, 2}
            sage: FiniteSetMaps(4, 3)([1, 0, 0, 1]).image_set()
            {0, 1}
        """
        return Set_object_enumerated(self)

    cpdef fibers(self):
        """
        Returns the fibers of ``self``

        OUTPUT:

            a dictionary ``d`` such that ``d[y]`` is the set of all ``x`` in
            ``domain`` such that ``f(x) = y``

        EXAMPLES::

            sage: FiniteSetMaps(4, 3)([1, 0, 2, 1]).fibers()
            {0: {1}, 1: {0, 3}, 2: {2}}
            sage: F = FiniteSetMaps(["a", "b", "c"])
            sage: F.from_dict({"a": "b", "b": "a", "c": "b"}).fibers()
            {'a': {'b'}, 'b': {'a', 'c'}}
        """
        return fibers(self, self.domain())

    cpdef items(self):
        """
        The items of ``self``

        Return the list of the ordered pairs ``(x, self(x))``

        EXAMPLES::

            sage: FiniteSetMaps(4, 3)([1, 0, 2, 1]).items()
            [(0, 1), (1, 0), (2, 2), (3, 1)]
        """
        return [(i, self._getimage(i)) for i in self.domain()]

    cpdef check(self):
        """
        Performs checks on ``self``

        Check that ``self`` is a proper function and then calls
        ``parent.check_element(self)`` where ``parent`` is the parent
        of ``self``.

        TESTS::

            sage: fs = FiniteSetMaps(3, 2)
            sage: for el in fs: el.check()
            sage: fs([1,1])
            Traceback (most recent call last):
            ...
            AssertionError: Wrong number of values
            sage: fs([0,0,2])
            Traceback (most recent call last):
            ...
            AssertionError: Wrong value self(2) = 2
        """
        cdef int i, m, n
        m = self._parent._m
        n = self._parent._n
        assert self._len == m, "Wrong number of values"
        for i in range(m):
            assert 0 <= self._list[i] < n, "Wrong value self(%i) = %i"%(i, self._list[i])
        if hasattr(self._parent, 'check_element'):
            self._parent.check_element(self)

    cpdef FiniteSetMap_MN _compose_internal_(self, FiniteSetMap_MN other,
                                             Parent resParent):
        """
        TESTS::

            sage: FSM = FiniteSetMaps(3)
            sage: fs1 = FSM([1, 0, 2])
            sage: fs2 = FSM([0, 1, 2])
            sage: el = fs1*fs2; el
            [1, 0, 2]
            sage: el.check()
        """
        cdef FiniteSetMap_MN res = PY_NEW_SAME_TYPE(self)
        res._parent = resParent
        res._alloc_(self._len)
        for i in range(self._len):
            res._list[i] = other._list[self._list[i]]
        res.set_immutable()
        return res


cdef class FiniteSetMap_Set(FiniteSetMap_MN):
    """
    Data structure for maps

    We assume that the parent given as argument is such that:

    - the domain is in ``parent.domain()``
    - the codomain is in ``parent.codomain()``
    - ``parent._m`` contains the cardinality of the domain
    - ``parent._n`` contains the cardinality of the codomain
    - ``parent._unrank_domain`` and ``parent._rank_domain`` is a pair of
      reciprocal rank and unrank functions beween the domain and
      ``range(parent._m)``.
    - ``parent._unrank_codomain`` and ``parent._rank_codomain`` is a pair of
      reciprocal rank and unrank functions beween the codomain and
      ``range(parent._n)``.
    """

    def __init__(self, parent, fun, check=True):
        """
        EXAMPLES::

            sage: F = FiniteSetMaps(["a", "b", "c", "d"], ["u", "v", "w"])
            sage: F(lambda x: "v")
            map: a -> v, b -> v, c -> v, d -> v
        """
        # For speed we initialize self by hand.
        # super(FiniteSetMap_Set, self).__init__(parent, [], check=False)
        self._parent = parent
        self._alloc_(int(parent._m))
        for i, el in enumerate(parent.domain().list()):
            self._setitem(i, parent._rank_codomain(fun(el)))
        self.set_immutable()
        if check: self.check()

    from_list = classmethod(FiniteSetMap_Set_from_list)
    from_dict = classmethod(FiniteSetMap_Set_from_dict)

    def __call__(self, i):
        """
        Returns the image of ``i`` under the map ``self``

        INPUT:

        - ``i`` -- an int

        OUTPUT:

        an int

        EXAMPLES::

            sage: F = FiniteSetMaps(["a", "b"], [3, 4, 5])
            sage: fs = F.from_dict({"a": 3, "b": 5})
            sage: fs('a'), fs('b')
            (3, 5)
        """
        parent = self._parent
        return parent._unrank_codomain(self._getitem(parent._rank_domain(i)))

    cpdef image_set(self):
        """
        Returns the image set of ``self``

        EXAMPLES::

            sage: F = FiniteSetMaps(["a", "b", "c"])
            sage: F.from_dict({"a": "b", "b": "a", "c": "b"}).image_set()
            {'a', 'b'}
            sage: F = FiniteSetMaps(["a", "b", "c"])
            sage: F(lambda x: "c").image_set()
            {'c'}
        """
        image_i = self._parent._unrank_codomain
        return Set_object_enumerated([image_i(i) for i in self])

    cpdef setimage(self, i, j):
        """
        Set the image of ``i`` as ``j`` in ``self``

        .. warning:: ``self`` must be mutable otherwise an exception is raised.

        INPUT:

        - ``i``, ``j`` -- two ``object``'s

        OUTPUT: ``None``

        EXAMPLES::

            sage: F = FiniteSetMaps(["a", "b", "c", "d"], ["u", "v", "w"])
            sage: fs = F(lambda x: "v")
            sage: fs2 = copy(fs)
            sage: fs2.setimage("a", "w")
            sage: fs2
            map: a -> w, b -> v, c -> v, d -> v
            sage: with fs.clone() as fs3:
            ...       fs3.setimage("a", "u")
            ...       fs3.setimage("c", "w")
            sage: fs3
            map: a -> u, b -> v, c -> w, d -> v

        TESTS::

            sage: with fs.clone() as fs3:
            ...       fs3.setimage("z", 2)
            Traceback (most recent call last):
            ...
            ValueError: 'z' is not in list

            sage: with fs.clone() as fs3:
            ...       fs3.setimage(1, 4)
            Traceback (most recent call last):
            ...
            ValueError: 1 is not in list
        """
        parent = self._parent
        return self._setitem(parent._rank_domain(i), parent._rank_codomain(j))

    cpdef getimage(self, i):
        """
        Returns the image of ``i`` by ``self``

        INPUT:

        - ``i`` -- an ``int``

        EXAMPLES::

            sage: F = FiniteSetMaps(["a", "b", "c", "d"], ["u", "v", "w"])
            sage: fs = F._from_list_([1, 0, 2, 1])
            sage: map(fs.getimage, ["a", "b", "c", "d"])
            ['v', 'u', 'w', 'v']
        """
        parent = self._parent
        return parent._unrank_codomain(self._getitem(parent._rank_domain(i)))

    cpdef items(self):
        """
        The items of ``self``

        Return the list of the couple ``(x, self(x))``

        EXAMPLES::

            sage: F = FiniteSetMaps(["a", "b", "c"])
            sage: F.from_dict({"a": "b", "b": "a", "c": "b"}).items()
            [('a', 'b'), ('b', 'a'), ('c', 'b')]

        TESTS::

            sage: all(F.from_dict(dict(f.items())) == f for f in F)
            True
        """
        parent = self._parent
        return [(parent._unrank_domain(i),
                 parent._unrank_codomain(self._getitem(i)))
                for i in range(parent._m)]

    def _repr_(self):
        """
        TESTS::

            sage: F = FiniteSetMaps(["a", "b"], [3, 4, 5])
            sage: F._from_list_([0, 2])
            map: a -> 3, b -> 5
        """
        return "map: "+", ".join([("%s -> %s"%(i, self(i))) for i in self.domain()])


cpdef FiniteSetMap_Set FiniteSetMap_Set_from_list(cls, Parent parent, list lst):
    """
    Creates a ``FiniteSetMap`` from a list

    .. warning:: no check is performed !

    TESTS::

        sage: from sage.sets.finite_set_map_cy import FiniteSetMap_Set_from_list as from_list
        sage: F = FiniteSetMaps(["a", "b"], [3, 4, 5])
        sage: f = from_list(F.element_class, F, [0,2]); f.check(); f
        map: a -> 3, b -> 5
        sage: f.parent() is F
        True
        sage: f.is_immutable()
        True
    """
    cdef FiniteSetMap_MN res
    res = PY_NEW(cls)
    super(FiniteSetMap_MN, res).__init__(parent, lst)
    return res

cpdef FiniteSetMap_Set FiniteSetMap_Set_from_dict(cls, Parent parent, dict d):
    """
    Creates a ``FiniteSetMap`` from a dictionary

    .. warning:: no check is performed !

    TESTS::

        sage: from sage.sets.finite_set_map_cy import FiniteSetMap_Set_from_dict as from_dict
        sage: F = FiniteSetMaps(["a", "b"], [3, 4, 5])
        sage: f = from_dict(F.element_class, F, {"a": 3, "b": 5}); f.check(); f
        map: a -> 3, b -> 5
        sage: f.parent() is F
        True
        sage: f.is_immutable()
        True
    """
    cdef FiniteSetMap_Set res
    res = PY_NEW(cls)
    res.__init__(parent, d.__getitem__)
    return res


cdef class FiniteSetEndoMap_N(FiniteSetMap_MN):
    """
    Maps from ``range(n)`` to itself.

    .. seealso:: :class:`FiniteSetMap_MN` for assumptions on the parent

    TESTS::

        sage: fs = FiniteSetMaps(3)([1, 0, 2])
        sage: TestSuite(fs).run()
    """

    def __mul__(FiniteSetEndoMap_N self, FiniteSetEndoMap_N other):
        """
        TESTS::

            sage: F = FiniteSetMaps(3)
            sage: F([1, 0, 2]) * F([2, 1, 0])
            [2, 0, 1]
            sage: F = FiniteSetMaps(3, action="right")
            sage: F([1, 0, 2]) * F([2, 1, 0])
            [1, 2, 0]
        """
        assert(self._parent is other._parent), "Parent mismatch"
        if self._parent._action == "right":
            return self._compose_internal_(other, self._parent)
        else:
            return other._compose_internal_(self, self._parent)

    def __pow__(self, n, dummy):
        """
        Return the ``n``-th power of ``self``.

        INPUT::

        - ``n`` -- a positive integer
        - ``dummy`` -- not used; must be set to ``None`` (for compatibility only).

        EXAMPLES::

            sage: fs = FiniteSetMaps(3)([1,0,2])
            sage: fs^2
            [0, 1, 2]
            sage: fs^0
            [0, 1, 2]
            sage: fs.__pow__(2)
            [0, 1, 2]
        """
        if dummy is not None:
            raise RuntimeError, "__pow__ dummy argument not used"
        return generic_power_c(self, n, self.parent().one())


cdef class FiniteSetEndoMap_Set(FiniteSetMap_Set):
    """
    Maps from a set to itself

    .. seealso:: :class:`FiniteSetMap_Set` for assumptions on the parent

    TESTS::

        sage: F = FiniteSetMaps(["a", "b", "c"])
        sage: f = F.from_dict({"a": "b", "b": "a", "c": "b"}); f
        map: a -> b, b -> a, c -> b
        sage: TestSuite(f).run()
    """

    def __mul__(FiniteSetEndoMap_Set self, FiniteSetEndoMap_Set other):
        """
        TESTS::

            sage: F = FiniteSetMaps(["a", "b", "c"])
            sage: f = F.from_dict({"a": "b", "b": "a", "c": "b"}); f
            map: a -> b, b -> a, c -> b
            sage: g = F.from_dict({"a": "c", "b": "c", "c": "a"}); g
            map: a -> c, b -> c, c -> a
            sage: f * g
            map: a -> b, b -> b, c -> b
            sage: g * f
            map: a -> c, b -> c, c -> c
        """
        assert(self._parent is other._parent), "Parent mismatch"
        if self._parent._action == "right":
            return self._compose_internal_(other, self._parent)
        else:
            return other._compose_internal_(self, self._parent)

    def __pow__(self, n, dummy):
        """
        Return the ``n``-th power of self.

        INPUT::

        - ``n`` -- a positive integer
        - ``dummy`` -- not used; must be set to None (for compatibility only).

        EXAMPLES::

            sage: F = FiniteSetMaps(["a", "b", "c"])
            sage: f = F.from_dict({"a": "b", "b": "a", "c": "b"}); f
            map: a -> b, b -> a, c -> b
            sage: f^2
            map: a -> a, b -> b, c -> a
            sage: f^3 == f
            True
        """
        if dummy is not None:
            raise RuntimeError, "__pow__ dummy argument not used"
        return generic_power_c(self, n, self.parent().one())