from __future__ import print_function, division

from sympy.combinatorics.fp_groups import FpGroup, FpSubgroup
from sympy.combinatorics.free_groups import FreeGroup, FreeGroupElement
from sympy.combinatorics.perm_groups import PermutationGroup

class GroupHomomorphism(object):
    '''
    A class representing group homomorphisms. Instantiate using `homomorphism()`.

    References
    ==========
    [1] Holt, D., Eick, B. and O'Brien, E. (2005). Handbook of computational group theory.

    '''

    def __init__(self, domain, codomain, images):
        self.domain = domain
        self.codomain = codomain
        self.images = images
        self._inverses = None
        self._kernel = None
        self._image = None

    def _invs(self):
        '''
        Return a dictionary with `{gen: inverse}` where `gen` is a rewriting
        generator of `codomain` (e.g. strong generator for permutation groups)
        and `inverse` is an element of its preimage

        '''
        image = self.image()
        inverses = {}
        for k in list(self.images.keys()):
            v = self.images[k]
            if not (v in inverses
                    or v.is_identity):
                inverses[v] = k
        if isinstance(self.codomain, PermutationGroup):
            gens = image.strong_gens
        else:
            gens = image.generators
        for g in gens:
            if g in inverses or g.is_identity:
                continue
            w = self.domain.identity
            if isinstance(self.codomain, PermutationGroup):
                parts = image._strong_gens_slp[g][::-1]
            else:
                parts = g
            for s in parts:
                if s in inverses:
                    w = w*inverses[s]
                else:
                    w = w*inverses[s**-1]**-1
            inverses[g] = w

        return inverses

    def invert(self, g):
        '''
        Return an element of the preimage of `g` or of each element
        of `g` if `g` is a list.
        NOTE: If the codomain is an FpGroup, the inverse for equal
        elements might not always be the same unless the FpGroup's
        rewriting system is confluent. However, making a system
        confluent can be time-consuming. If it's important, try
        `self.codomain.make_confluent()` first.

        '''
        from sympy.combinatorics import Permutation
        from sympy.combinatorics.free_groups import FreeGroupElement
        if isinstance(g, (Permutation, FreeGroupElement)):
            if isinstance(self.codomain, FpGroup):
                g = self.codomain.reduce(g)
            if self._inverses is None:
                self._inverses = self._invs()
            image = self.image()
            w = self.domain.identity
            if isinstance(self.codomain, PermutationGroup):
                gens = image.generator_product(g)[::-1]
            else:
                gens = g
            # the following can't be "for s in gens:"
            # because that would be equivalent to
            # "for s in gens.array_form:" when g is
            # a FreeGroupElement. On the other hand,
            # when you call gens by index, the generator
            # (or inverse) at position i is returned.
            for i in range(len(gens)):
                s = gens[i]
                if s.is_identity:
                    continue
                if s in self._inverses:
                    w = w*self._inverses[s]
                else:
                    w = w*self._inverses[s**-1]**-1
            return w
        elif isinstance(g, list):
            return [self.invert(e) for e in g]

    def kernel(self):
        '''
        Compute the kernel of `self`.

        '''
        if self._kernel is None:
            self._kernel = self._compute_kernel()
        return self._kernel

    def _compute_kernel(self):
        from sympy import S
        G = self.domain
        G_order = G.order()
        if G_order == S.Infinity:
            raise NotImplementedError(
                "Kernel computation is not implemented for infinite groups")
        gens = []
        if isinstance(G, PermutationGroup):
            K = PermutationGroup(G.identity)
        else:
            K = FpSubgroup(G, gens, normal=True)
        i = self.image().order()
        while K.order()*i != G_order:
            r = G.random()
            k = r*self.invert(self(r))**-1
            if not k in K:
                gens.append(k)
                if isinstance(G, PermutationGroup):
                    K = PermutationGroup(gens)
                else:
                    K = FpSubgroup(G, gens, normal=True)
        return K

    def image(self):
        '''
        Compute the image of `self`.

        '''
        if self._image is None:
            values = list(set(self.images.values()))
            if isinstance(self.codomain, PermutationGroup):
                self._image = self.codomain.subgroup(values)
            else:
                self._image = FpSubgroup(self.codomain, values)
        return self._image

    def _apply(self, elem):
        '''
        Apply `self` to `elem`.

        '''
        if not elem in self.domain:
            if isinstance(elem, (list, tuple)):
                return [self._apply(e) for e in elem]
            raise ValueError("The supplied element doesn't belong to the domain")
        if elem.is_identity:
            return self.codomain.identity
        else:
            images = self.images
            value = self.codomain.identity
            if isinstance(self.domain, PermutationGroup):
                gens = self.domain.generator_product(elem, original=True)
                for g in gens:
                    if g in self.images:
                        value = images[g]*value
                    else:
                        value = images[g**-1]**-1*value
            else:
                i = 0
                for _, p in elem.array_form:
                    if p < 0:
                        g = elem[i]**-1
                    else:
                        g = elem[i]
                    value = value*images[g]**p
                    i += abs(p)
        return value

    def __call__(self, elem):
        return self._apply(elem)

    def is_injective(self):
        '''
        Check if the homomorphism is injective

        '''
        return self.kernel().order() == 1

    def is_surjective(self):
        '''
        Check if the homomorphism is surjective

        '''
        from sympy import S
        im = self.image().order()
        oth = self.codomain.order()
        if im == S.Infinity and oth == S.Infinity:
            return None
        else:
            return im == oth

    def is_isomorphism(self):
        '''
        Check if `self` is an isomorphism.

        '''
        return self.is_injective() and self.is_surjective()

    def is_trivial(self):
        '''
        Check is `self` is a trivial homomorphism, i.e. all elements
        are mapped to the identity.

        '''
        return self.image().order() == 1

    def compose(self, other):
        '''
        Return the composition of `self` and `other`, i.e.
        the homomorphism phi such that for all g in the domain
        of `other`, phi(g) = self(other(g))

        '''
        if not other.image().is_subgroup(self.domain):
            raise ValueError("The image of `other` must be a subgroup of "
                    "the domain of `self`")
        images = {g: self(other(g)) for g in other.images}
        return GroupHomomorphism(other.domain, self.codomain, images)

    def restrict_to(self, H):
        '''
        Return the restriction of the homomorphism to the subgroup `H`
        of the domain.

        '''
        if not isinstance(H, PermutationGroup) or not H.is_subgroup(self.domain):
            raise ValueError("Given H is not a subgroup of the domain")
        domain = H
        images = {g: self(g) for g in H.generators}
        return GroupHomomorphism(domain, self.codomain, images)

    def invert_subgroup(self, H):
        '''
        Return the subgroup of the domain that is the inverse image
        of the subgroup `H` of the homomorphism image

        '''
        if not H.is_subgroup(self.image()):
            raise ValueError("Given H is not a subgroup of the image")
        gens = []
        P = PermutationGroup(self.image().identity)
        for h in H.generators:
            h_i = self.invert(h)
            if h_i not in P:
                gens.append(h_i)
                P = PermutationGroup(gens)
            for k in self.kernel().generators:
                if k*h_i not in P:
                    gens.append(k*h_i)
                    P = PermutationGroup(gens)
        return P

def homomorphism(domain, codomain, gens, images=[], check=True):
    '''
    Create (if possible) a group homomorphism from the group `domain`
    to the group `codomain` defined by the images of the domain's
    generators `gens`. `gens` and `images` can be either lists or tuples
    of equal sizes. If `gens` is a proper subset of the group's generators,
    the unspecified generators will be mapped to the identity. If the
    images are not specified, a trivial homomorphism will be created.

    If the given images of the generators do not define a homomorphism,
    an exception is raised.

    If `check` is `False`, don't check whether the given images actually
    define a homomorphism.

    '''
    if not isinstance(domain, (PermutationGroup, FpGroup, FreeGroup)):
        raise TypeError("The domain must be a group")
    if not isinstance(codomain, (PermutationGroup, FpGroup, FreeGroup)):
        raise TypeError("The codomain must be a group")

    generators = domain.generators
    if any([g not in generators for g in gens]):
        raise ValueError("The supplied generators must be a subset of the domain's generators")
    if any([g not in codomain for g in images]):
        raise ValueError("The images must be elements of the codomain")

    if images and len(images) != len(gens):
        raise ValueError("The number of images must be equal to the number of generators")

    gens = list(gens)
    images = list(images)

    images.extend([codomain.identity]*(len(generators)-len(images)))
    gens.extend([g for g in generators if g not in gens])
    images = dict(zip(gens,images))

    if check and not _check_homomorphism(domain, codomain, images):
        raise ValueError("The given images do not define a homomorphism")
    return GroupHomomorphism(domain, codomain, images)

def _check_homomorphism(domain, codomain, images):
    if hasattr(domain, 'relators'):
        rels = domain.relators
    else:
        gens = domain.presentation().generators
        rels = domain.presentation().relators
    identity = codomain.identity

    def _image(r):
        if r.is_identity:
            return identity
        else:
            w = identity
            r_arr = r.array_form
            i = 0
            j = 0
            # i is the index for r and j is for
            # r_arr. r_arr[j] is the tuple (sym, p)
            # where sym is the generator symbol
            # and p is the power to which it is
            # raised while r[i] is a generator
            # (not just its symbol) or the inverse of
            # a generator - hence the need for
            # both indices
            while i < len(r):
                power = r_arr[j][1]
                if isinstance(domain, PermutationGroup):
                    s = domain.generators[gens.index(r[i])]
                else:
                    s = r[i]
                if s in images:
                    w = w*images[s]**power
                else:
                    w = w*images[s**-1]**power
                i += abs(power)
                j += 1
            return w

    for r in rels:
        if isinstance(codomain, FpGroup):
            s = codomain.equals(_image(r), identity)
            if s is None:
                # only try to make the rewriting system
                # confluent when it can't determine the
                # truth of equality otherwise
                success = codomain.make_confluent()
                s = codomain.equals(_image(r), identity)
                if s in None and not success:
                    raise RuntimeError("Can't determine if the images "
                        "define a homomorphism. Try increasing "
                        "the maximum number of rewriting rules "
                        "(group._rewriting_system.set_max(new_value); "
                        "the current value is stored in group._rewriting"
                        "_system.maxeqns)")
        else:
            s = _image(r).is_identity
        if not s:
            return False
    return True

def orbit_homomorphism(group, omega):
    '''
    Return the homomorphism induced by the action of the permutation
    group `group` on the set `omega` that is closed under the action.

    '''
    from sympy.combinatorics import Permutation
    from sympy.combinatorics.named_groups import SymmetricGroup
    codomain = SymmetricGroup(len(omega))
    identity = codomain.identity
    omega = list(omega)
    images = {g: identity*Permutation([omega.index(o^g) for o in omega]) for g in group.generators}
    group._schreier_sims(base=omega)
    H = GroupHomomorphism(group, codomain, images)
    if len(group.basic_stabilizers) > len(omega):
        H._kernel = group.basic_stabilizers[len(omega)]
    else:
        H._kernel = PermutationGroup([group.identity])
    return H

def block_homomorphism(group, blocks):
    '''
    Return the homomorphism induced by the action of the permutation
    group `group` on the block system `blocks`. The latter should be
    of the same form as returned by the `minimal_block` method for
    permutation groups, namely a list of length `group.degree` where
    the i-th entry is a representative of the block i belongs to.

    '''
    from sympy.combinatorics import Permutation
    from sympy.combinatorics.named_groups import SymmetricGroup

    n = len(blocks)

    # number the blocks; m is the total number,
    # b is such that b[i] is the number of the block i belongs to,
    # p is the list of length m such that p[i] is the representative
    # of the i-th block
    m = 0
    p = []
    b = [None]*n
    for i in range(n):
        if blocks[i] == i:
            p.append(i)
            b[i] = m
            m += 1
    for i in range(n):
        b[i] = b[blocks[i]]

    codomain = SymmetricGroup(m)
    # the list corresponding to the identity permutation in codomain
    identity = range(m)
    images = {g: Permutation([b[p[i]^g] for i in identity]) for g in group.generators}
    H = GroupHomomorphism(group, codomain, images)
    return H
