#!/usr/bin/python3
# -*- coding: utf-8 -*-

class Relation:
    """Mutable eagerly materialized binary relation class.

    For example, Relation([(1,2), (1,3)])[1] == {2, 3}"""
    def __init__(self, kvpairs=()):
        self.kids = {}
        for k, v in kvpairs:
            self.add(k, v)

    def __getitem__(self, key):
        "Get set of values associated with key."
        return set(self.kids.get(key, ()))

    def add(self, key, value):
        "Add a value associated with a key."
        if key not in self.kids:
            self.kids[key] = set()
        self.kids[key].add(value)

    def __contains__(self, kvpair):
        "Determine whether a value is associated with a key."
        k, v = kvpair
        return v in self[k]

    def keys(self):
        "Iterate over keys."
        return iter(self.kids)

    def __iter__(self):
        "Iterate over keys, like iter(dict)."
        return self.keys()

    def pairs(self):
        "Return a fresh set of key-value pairs."
        return {(k, v) for k in self for v in self[k]}

    def __repr__(self):
        return 'Relation({})'.format(', '.join(repr(pair)
                                               for pair in self.pairs()))

    def __eq__(self, other):
        return self.pairs() == other.pairs()

    def inverse(self):
        "Compute the inverse relation."
        return Relation((v, k) for k, v in self.pairs())

    def tclosure(self):
        "Transitive closure.  R⁺, not R*, which is .rtclosure()."
        closure = Relation()
        tovisit = self.pairs()
        while tovisit:
            k, v = pair = tovisit.pop()
            if pair not in closure:
                closure.add(k, v)
                for next in self[v]:
                    tovisit.add((k, next))

        return closure

    def __or__(self, other):
        "Union."
        return Relation((k, v) for r in [self, other] for k, v in r.pairs())

    def __add__(self, other):
        "Composition.  Operator chosen for harmony with Kleene algebra."
        return Relation((k, vv) for k, v in self.pairs() for vv in other[v])

    def __mul__(self, other):
        "Relational cartesian product.  Operator chosen for harmony with Kleene algebra."
        return Relation((k, (v1, v2)) for k in set(self.keys()) | set(other.keys())
                        for v1 in self[k] for v2 in other[k])

    def rclosure(self):
        "Reflexive closure R^."
        return self | Relation((k, k) for k in self)

    def rtclosure(self):
        "Reflexive transitive closure R*; the Kleene algebra * operator."
        return self.rclosure() | self.tclosure()


def _ok(a, b): assert a == b, (a, b)

_r = Relation([(1,2), (1,3)])
_ok(set(_r[1]), {2, 3})
_ok(_r[0], set())
_ok(_r.inverse(), Relation([(2, 1), (3, 1)]))
_ok(_r.rclosure(), Relation([(1, 1), (1, 2), (1, 3)]))
_ok(_r.rtclosure(), Relation([(1, 1), (1, 2), (1, 3)]))

_q = Relation([(1, 2), (2, 3)])
_ok(_q.tclosure().pairs(), {(1, 2), (1, 3), (2, 3)})
_ok(_q.rclosure().pairs(), {(1, 1), (1, 2), (2, 2), (2, 3)})
_ok(_q.rtclosure().pairs(), {(1, 1), (1, 2), (1, 3), (2, 2), (2, 3)})

# Edges of a cube, with coordinates represented in bits; or the
# transitive reduction of the ⊂ relation on the powerset of a set of
# three items.
_s = Relation([(0, 1), (0, 2), (0, 4),
               (1, 3), (2, 3), (1, 5), (2, 6), (4, 5), (4, 6),
               (3, 7), (5, 7), (6, 7)])
_ok(_s[0], {1, 2, 4})
_ok(_s.rclosure()[0], {0, 1, 2, 4})
_ok(_s.tclosure()[3], {7})
_ok(_s.tclosure()[1], {3, 5, 7})

_ok(_s[2], {3, 6})
_ok(_s.tclosure()[2], {3, 6, 7})
_ok(_s.rclosure()[2], {2, 3, 6})
_ok(_s.rtclosure()[2], {2, 3, 6, 7})

_ok(_s.tclosure()[0], set(range(1, 8)))
_ok(_s.tclosure(), _s | _s + _s | _s + _s + _s)
_ok((_s + _s)[0], {3, 5, 6})
_ok(_s.inverse()[5], {1, 4})
_ok((_s | _s.inverse())[3], {1, 2, 7})
_ok((_s * _s.inverse())[3], {(7, 1), (7, 2)})
_ok((_r * _s)[1], {(2, 3), (2, 5), (3, 3), (3, 5)})