"""Deligne's partition category (Schwartz / Harman-Snowden basis) and S_t.

A morphism m -> n is a dict {partition: coeff}; a partition is a canonical
tuple of sorted tuples of point labels. Bottom points are 0..m-1 (left to
right), top points are m..m+n-1 (left to right).

Composition: blocks are merged through the middle; every block lying entirely
in the middle contributes a factor t.  The coefficient ring is whatever T is
(Fraction, or a flint/sympy element).
"""
from fractions import Fraction
from itertools import combinations


def canon(blocks):
    return tuple(sorted(tuple(sorted(b)) for b in blocks if len(b) > 0))


class Cat:
    def __init__(self, t):
        self.t = t
        self.one = t ** 0
        self.zero = t * 0
        self._cache = {}

    # ---------- basic morphisms -------------------------------------------
    def mor(self, m, n, d=None):
        return Mor(self, m, n, d or {})

    def part(self, m, n, blocks, c=None):
        return Mor(self, m, n, {canon(blocks): self.one if c is None else c})

    def id(self, k):
        return self.part(k, k, [(i, k + i) for i in range(k)])

    # ---------- operations ------------------------------------------------
    def compose(self, g, f):
        """g o f  (f first)."""
        assert f.n == g.m, (f.m, f.n, g.m, g.n)
        m, n, p = f.m, f.n, g.n
        out = {}
        t = self.t
        for pf, cf in f.d.items():
            for pg, cg in g.d.items():
                key = (pf, pg, m, n, p)
                r = self._cache.get(key)
                if r is None:
                    r = self._comp_parts(pf, pg, m, n, p)
                    self._cache[key] = r
                part, k = r
                v = cf * cg * (t ** k)
                out[part] = out.get(part, self.zero) + v
        return Mor(self, m, p, out).clean()

    @staticmethod
    def _comp_parts(pf, pg, m, n, p):
        # global labels: f bottom 0..m-1 ; middle m..m+n-1 ; g top m+n..m+n+p-1
        parent = list(range(m + n + p))

        def find(x):
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x

        def union(a, b):
            ra, rb = find(a), find(b)
            if ra != rb:
                parent[ra] = rb
        for b in pf:  # f labels: bottom i -> i ; top j (label m+j) -> m+j
            for x in b[1:]:
                union(b[0], x)
        for b in pg:  # g labels: bottom j (label j) -> m+j ; top k (label n+k) -> m+n+k
            g0 = [(m + x) for x in b]
            for x in g0[1:]:
                union(g0[0], x)
        groups = {}
        for x in range(m + n + p):
            groups.setdefault(find(x), []).append(x)
        k = 0
        blocks = []
        for gr in groups.values():
            outer = [x if x < m else x - n for x in gr if x < m or x >= m + n]
            if outer:
                blocks.append(outer)
            else:
                k += 1
        return canon(blocks), k

    def tensor(self, f, g):
        m1, n1, m2, n2 = f.m, f.n, g.m, g.n
        m, n = m1 + m2, n1 + n2
        out = {}

        def rf(x):  # f label -> new label
            return x if x < m1 else m + (x - m1)

        def rg(x):
            return m1 + x if x < m2 else m + n1 + (x - m2)
        for pf, cf in f.d.items():
            for pg, cg in g.d.items():
                blocks = [[rf(x) for x in b] for b in pf] + [[rg(x) for x in b] for b in pg]
                key = canon(blocks)
                out[key] = out.get(key, self.zero) + cf * cg
        return Mor(self, m, n, out).clean()


class Mor:
    def __init__(self, C, m, n, d):
        self.C, self.m, self.n, self.d = C, m, n, d

    def clean(self):
        self.d = {k: v for k, v in self.d.items() if v != 0}
        return self

    def __add__(self, o):
        assert (self.m, self.n) == (o.m, o.n)
        d = dict(self.d)
        for k, v in o.d.items():
            d[k] = d.get(k, self.C.zero) + v
        return Mor(self.C, self.m, self.n, d).clean()

    def __neg__(self):
        return Mor(self.C, self.m, self.n, {k: -v for k, v in self.d.items()})

    def __sub__(self, o):
        return self + (-o)

    def scale(self, c):
        return Mor(self.C, self.m, self.n, {k: c * v for k, v in self.d.items()}).clean()

    def __rmul__(self, c):
        return self.scale(c)

    def __mul__(self, other):  # composition: self o other
        if isinstance(other, Mor):
            return self.C.compose(self, other)
        return self.scale(other)

    def __matmul__(self, other):  # tensor product
        return self.C.tensor(self, other)

    def iszero(self):
        return len(self.d) == 0

    def __repr__(self):
        return f"Mor({self.m}->{self.n}, {len(self.d)} terms)"


# ---------------------------------------------------------------------------
# cyclic labelling helpers: the cyclic order of the boundary of Hom(m,n) is
# b0 b1 ... b_{m-1} t_{n-1} ... t_0   (counterclockwise starting bottom left)
def cyc_order(m, n):
    return list(range(m)) + [m + j for j in reversed(range(n))]


def to_cyclic(part, m, n):
    """partition with Hom(m,n) labels -> partition of cyclic positions 0..m+n-1"""
    order = cyc_order(m, n)
    pos = {lab: i for i, lab in enumerate(order)}
    return canon([[pos[x] for x in b] for b in part])


def from_cyclic(cpart, m, n):
    order = cyc_order(m, n)
    return canon([[order[i] for i in b] for b in cpart])


def rotate(f, m2, n2, k=1):
    """Rotate the (m+n)-box f by k clicks (cyclic position i -> i+k) and read it
    as a morphism m2 -> n2."""
    N = f.m + f.n
    assert m2 + n2 == N
    out = {}
    for p, c in f.d.items():
        cp = to_cyclic(p, f.m, f.n)
        cp2 = canon([[(i + k) % N for i in b] for b in cp])
        out[from_cyclic(cp2, m2, n2)] = c
    return Mor(f.C, m2, n2, out)


def has_singleton(part):
    return any(len(b) == 1 for b in part)


def set_partitions(lst):
    if not lst:
        yield []
        return
    first, rest = lst[0], lst[1:]
    for p in set_partitions(rest):
        yield [[first]] + p
        for i in range(len(p)):
            yield p[:i] + [[first] + p[i]] + p[i + 1:]


def sf_partitions(N):
    """singleton-free partitions of range(N), canonical, sorted"""
    return sorted(canon(p) for p in set_partitions(list(range(N))) if all(len(b) > 1 for b in p))