A calculator for permutation

permute.py

from functools import total_ordering
from typing import *

_PermutationRingTermSubType = Union['Permutation', 'PermutationRingTerm']
_PermutationRingSubType = Union['Permutation', 'PermutationRingTerm', 'PermutationRing']
_CoType = Union[int, float]  # type of coefficient in group ring
_MultiplierType = Union[_PermutationRingSubType, _CoType]


def _check_co_type(_o: Any) -> bool:
    """
    check whether a variable is a proper coefficient type
    """
    return hasattr(_o, '__float__')


@total_ordering
class Permutation(object):
    def __init__(self, _map: Dict[int, int]):
        _map = {k: v for k, v in _map.items() if k != v}  # drop redundant map
        self.map: Dict[int, int] = _map

    def __mul__(self, other: _MultiplierType) -> _PermutationRingSubType:
        return _multiply(self, other)

    def __rmul__(self, other: _CoType) -> 'PermutationRingTerm':
        return _multiply(self, other)

    def __add__(self, other: _PermutationRingSubType) -> 'PermutationRing':
        return _add(self, other)

    def __neg__(self) -> 'PermutationRingTerm':
        return PermutationRingTerm(self, -1)

    def __sub__(self, other: _PermutationRingSubType) -> 'PermutationRing':
        return self + (- other)

    def __getitem__(self, i: int) -> int:
        return self.map.get(i, i)

    def __repr__(self) -> str:
        def repr_cycle(cycle):
            return f'({", ".join(map(str, cycle))})'
        if self.cycles():
            return f'{" ".join(map(repr_cycle, self.cycles()))}'
        else:
            return '(1)'

    def val(self) -> List[Tuple[int, int]]:
        # return a value for comparing
        return sorted(zip(self.map.keys(), self.map.values()))

    def __eq__(self, other: 'Permutation') -> bool:
        return self.val() == other.val()

    def __lt__(self, other: 'Permutation') -> bool:
        return self.val() < other.val()

    def inv(self) -> 'Permutation':
        """
        get the inverse of a permutation
        :return: the inverse permutation
        """
        return Permutation(_map={v: k for k, v in self.map.items()})

    def cycles(self) -> List[List[int]]:
        """
        return the cycles in a permutation
        :return: a list of integer list, representing elements in the cycle
        """
        elements = set(self.map.keys())
        cycles = []
        while elements:
            head = min(elements)
            elements.remove(head)
            cur = self[head]
            if head == cur:
                continue
            cycle: List[int] = [head]
            while cur != head:
                elements.remove(cur)
                cycle.append(cur)
                cur = self[cur]
            cycles.append(cycle)
        return cycles


def c(*nums: int) -> Permutation:
    """
    build a cyclic permutation from a list of integers
    :param nums: some distinct positive integers in the cycle
    :return: corresponding cyclic permutation
    """
    length = len(nums)
    _map = {nums[i]: nums[(i + 1) % length] for i in range(length)}
    return Permutation(_map)


class PermutationRingTerm:
    def __init__(self, permutation: 'Permutation', co: Optional[float]):
        """
        A class representing a term in a element in group ring
        :param permutation:
        :param co: coefficient
        """
        self.co: _CoType = co or 1
        self.permutation: 'Permutation' = permutation

    def __add__(self, other: _PermutationRingSubType) -> 'PermutationRing':
        return _add(self, other)

    def __neg__(self) -> 'PermutationRingTerm':
        return PermutationRingTerm(self.permutation, -self.co)

    def __sub__(self, other: _PermutationRingSubType) -> 'PermutationRing':
        return self + (- other)

    def __mul__(self, other: _MultiplierType) -> Union['PermutationRingTerm', 'PermutationRing']:
        return _multiply(self, other)

    def __rmul__(self, other: _CoType) -> 'PermutationRingTerm':
        return _multiply(self, other)

    def __repr__(self) -> str:
        return str(PermutationRing.build(self))

    @classmethod
    def build(cls, _from: _PermutationRingTermSubType) -> 'PermutationRingTerm':
        # build an object from low-level object
        if isinstance(_from, Permutation):
            return PermutationRingTerm(_from, 1)
        elif isinstance(_from, PermutationRingTerm):
            return _from
        else:
            raise TypeError


class PermutationRing:
    def __init__(self, terms: List['PermutationRingTerm']):
        self.terms = sorted(terms, key=lambda term: term.permutation)

    def trim(self) -> 'PermutationRing':
        result_terms: List['PermutationRingTerm'] = []
        cur_permutation: Optional['Permutation'] = None
        cur_co: float = 0
        for term in self.terms:
            if cur_permutation is None:
                cur_permutation = term.permutation
                cur_co = term.co
            elif cur_permutation is not None and term.permutation == cur_permutation:
                cur_co += term.co
            else:
                result_terms.append(PermutationRingTerm(cur_permutation, cur_co))
                cur_permutation = term.permutation
                cur_co = term.co
        result_terms.append(PermutationRingTerm(cur_permutation, cur_co))
        return PermutationRing(result_terms)

    def __mul__(self, other: _MultiplierType) -> 'PermutationRing':
        return _multiply(self, other)

    def __rmul__(self, other: _CoType) -> 'PermutationRing':
        return _multiply(self, other)

    def __add__(self, other: _PermutationRingSubType) -> 'PermutationRing':
        return _add(self, other)

    def __neg__(self) -> 'PermutationRing':
        return PermutationRing([
            -t for t in self.terms
        ])

    def __sub__(self, other: _PermutationRingSubType) -> 'PermutationRing':
        return self + (- other)

    def __repr__(self) -> str:
        res: List[str] = []
        for term in self.terms:
            if term.co == 0:
                res += f' + {term.permutation}'
            elif term.co > 0:
                res += f' + {term.co} {term.permutation}'
            elif term.co < 0:
                res += f' - {- term.co} {term.permutation}'
        res_str = ''.join(res)
        if res_str.startswith(' + '):
            return res_str[3:]
        else:
            return res_str

    @classmethod
    def build(cls, terms: _PermutationRingSubType) -> 'PermutationRing':
        # build an object from low-level object
        if isinstance(terms, Permutation):
            terms = [1 * terms]
        elif isinstance(terms, PermutationRingTerm):
            terms = [terms]
        elif isinstance(terms, PermutationRing):
            terms = terms.terms
        elif isinstance(terms, List):
            terms = sorted(terms, key=lambda term: term.permutation)
        return PermutationRing(terms)


def _multiply(_a: _MultiplierType, _b: _MultiplierType) -> _PermutationRingSubType:
    if isinstance(_a, Permutation):
        if isinstance(_b, Permutation):
            b_r = _b.inv()
            values = _b.map.values()
            _map = _b.map.copy()
            for k, v in _a.map.items():
                if k in values:
                    _map[b_r[k]] = v
                else:
                    _map[k] = v
            return Permutation(_map=_map)
        elif isinstance(_b, PermutationRingTerm):
            return PermutationRingTerm(_a * _b.permutation, _b.co)
        elif isinstance(_b, PermutationRing):
            return PermutationRing([
                _a * t for t in _b.terms
            ])
        elif _check_co_type(_b):
            return PermutationRingTerm(_a, _b)
        else:
            raise TypeError

    elif isinstance(_a, PermutationRingTerm):
        if isinstance(_b, Permutation):
            return PermutationRingTerm(_a.permutation * _b, _a.co)
        elif isinstance(_b, PermutationRingTerm):
            return PermutationRingTerm(_a.permutation * _b.permutation, _a.co * _b.co)
        elif isinstance(_b, PermutationRing):
            return PermutationRing([
                _a * t for t in _b.terms
            ])
        elif _check_co_type(_b):
            return PermutationRingTerm(_a.permutation, _a.co * _b)
        else:
            raise TypeError

    elif isinstance(_a, PermutationRing):
        if isinstance(_b, Permutation) or isinstance(_b, PermutationRingTerm) or _check_co_type(_b):
            return PermutationRing([
                t * _b for t in _a.terms
            ])
        if isinstance(_b, PermutationRing):
            return PermutationRing([
                t1 * t2 for t1 in _a.terms for t2 in _b.terms
            ]).trim()

    else:
        raise TypeError


def _add(_a: _PermutationRingSubType, _b: _PermutationRingSubType) -> PermutationRing:
    _a = PermutationRing.build(_a)
    _b = PermutationRing.build(_b)
    return PermutationRing(_a.terms + _b.terms).trim()


if __name__ == "__main__":
    x = c(1, 2) + 2 * c(3, 5) - c() - 3 * c(1)
    y = c() + 3 * c(1, 2)
    z = c(1, 10) + c(1, 2)
    print(x * y)  # - 1 (1) - 11 (1, 2) + 6 (1, 2) (3, 5) + 2 (3, 5)