"""The category S_t (restriction of Deligne's category to the image X of P)."""
from deligne import Cat, Mor, canon, rotate, sf_partitions, has_singleton, to_cyclic, from_cyclic
from fractions import Fraction


class St:
    def __init__(self, t, vscale=None):
        C = Cat(t)
        self.C, self.t = C, t
        one = C.one
        self.one = one
        self.zero = C.zero
        a = t  # counit o unit
        self.P = C.part(1, 1, [(0, 1)]) - C.part(1, 1, [(0,), (1,)], one / a)
        P2 = self.P @ self.P
        self.P2 = P2
        Y = C.part(2, 1, [(0, 1, 2)])
        Ys = C.part(1, 2, [(0, 1, 2)])
        v = one if vscale is None else vscale
        self.tau = (self.P * Y * P2).scale(v)
        self.taus = (P2 * Ys * self.P).scale(v)
        self.cup = P2 * C.part(0, 2, [(0, 1)])
        self.cap = C.part(2, 0, [(0, 1)]) * P2
        self.c = P2 * C.part(2, 2, [(0, 3), (1, 2)])
        self.i1 = self.P
        self._bases = {}

    def id(self, k):
        r = self.C.mor(0, 0, {(): self.one})
        for _ in range(k):
            r = r @ self.P
        return r

    def Pk(self, k):
        return self.id(k)

    # coordinates ----------------------------------------------------------
    def basis_parts(self, m, n):
        key = (m, n)
        if key not in self._bases:
            cps = sf_partitions(m + n)  # cyclic-position partitions
            self._bases[key] = [from_cyclic(cp, m, n) for cp in cps]
        return self._bases[key]

    def coords(self, f):
        """coordinates of an S_t morphism f (assumed P-restricted) in the basis
        Res(B_pi), pi singleton free."""
        B = self.basis_parts(f.m, f.n)
        return [f.d.get(p, self.zero) for p in B]

    def basis_elt(self, m, n, i):
        p = self.basis_parts(m, n)[i]
        return self.id(n) * self.C.mor(m, n, {p: self.one}) * self.id(m)

    def rot(self, f, m2, n2, k=1):
        return rotate(f, m2, n2, k)