"""
The multimap data structure associates multiple values to a key.

In this module the multimap is implemented by a dictionary in which each key is
associated to a container, which can be a list, a dictionary, or a set.

These containers are created, accessed, and extended using the usual array
notation, e.g. m["a"] = 1, m["a"] = 2 creates a container for key "a"
containing 1 and 2. An item within a container can be removed using the
"remove"  method.

Requires Python 2.7.
"""

import collections


class Map(object):
    """ Map wraps a dictionary. It is essentially an abstract class from which
    specific multimaps are subclassed. """

    def __init__(self):
        self._dict = {}

    def __repr__(self):
        return "%s(%s)" % (self.__class__.__name__, repr(self._dict))

    __str__ = __repr__

    def __getitem__(self, key):
        return self._dict[key]

    def __setitem__(self, key, value):
        self._dict[key] = value

    def __delitem__(self, key):
        del self._dict[key]

    def insert(self, **values):
        for k, v in values.items():
            self.__setitem__(k, v)

    def remove(self, key, value):
        del self._dict[key]

    def dict(self):
        """ Allows access to internal dictionary, if necessary. Caution: multimaps 
        will break if keys are not associated with proper container."""
        return self._dict


class ListMultimap(Map):
    """ ListMultimap is based on lists and allows multiple instances of same value. """

    def __init__(self, **values):
        self._dict = collections.defaultdict(list)
        self.insert(**values)

    def __setitem__(self, key, value):
        self._dict[key].append(value)

    def remove(self, key, value):
        self._dict[key].remove(value)


class SetMultimap(Map):
    """ SetMultimap is based on sets and prevents multiple instances of same value. """

    def __init__(self):
        self._dict = collections.defaultdict(set)

    def __setitem__(self, key, value):
        self._dict[key].add(value)

    def remove(self, key, value):
        self._dict[key].remove(value)


class DictMultimap(Map):
    """ DictMultimap is based on dicts and allows fast tests for membership. """

    def __init__(self):
        self._dict = collections.defaultdict(dict)

    def __setitem__(self, key, value):
        self._dict[key][value] = True

    def remove(self, key, value):
        del self._dict[key][value]


def test():
    def test_multimap(m):
        print("__________________________________")
        print(m["a"])
        m["a"] = 1
        m["a"] = 2
        m["a"] = 2
        m["a"] = 3
        m["a"] = 4
        m.insert(b='v1', c='v2')
        print(m)
        m.remove("a", 4)
        print(m)
        print(("a" in m.dict()))
        print(m["a"])
        m["a"] = 5
        m["b"] = 6
        print(m)
        del m["b"]
        print(m)
        print((3 in m["a"]))

    test_multimap(ListMultimap())
    test_multimap(SetMultimap())
    test_multimap(DictMultimap())


if __name__ == "__main__":
    test()
