"""
Composition species
"""
from __future__ import absolute_import
#*****************************************************************************
#       Copyright (C) 2008 Mike Hansen <mhansen@gmail.com>,
#
#  Distributed under the terms of the GNU General Public License (GPL)
#
#    This code is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    General Public License for more details.
#
#  The full text of the GPL is available at:
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************
from .species import GenericCombinatorialSpecies
from .structure import GenericSpeciesStructure
from .partition_species import PartitionSpecies
from sage.misc.cachefunc import cached_function
from sage.structure.unique_representation import UniqueRepresentation

class CompositionSpeciesStructure(GenericSpeciesStructure):
    def __init__(self, parent, labels, pi, f, gs):
        """
        TESTS::

            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: a = L.structures(['a','b','c']).random_element()
            sage: a == loads(dumps(a))
            True
        """
        self._partition = pi
        GenericSpeciesStructure.__init__(self, parent, labels, [f, gs])

    def __repr__(self):
        """
        EXAMPLES::

            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: L.structures(['a','b','c']).random_element()
            F-structure: {{'a', 'b', 'c'}}; G-structures: (('a', 'b', 'c'),)
        """
        f, gs = self._list
        return "F-structure: %s; G-structures: %s"%(repr(f), repr(gs))

    def transport(self, perm):
        """
        EXAMPLES::

            sage: p = PermutationGroupElement((2,3))
            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: S = L.structures(['a','b','c']).list()
            sage: a = S[2]; a
            F-structure: {{'a', 'c'}, {'b'}}; G-structures: (('a', 'c'), ('b'))
            sage: a.transport(p)
            F-structure: {{'a', 'b'}, {'c'}}; G-structures: (('a', 'c'), ('b'))
        """
        f, gs = self._list
        pi = self._partition.transport(perm)
        f = f.change_labels(pi._list)
        g = [g.change_labels(part) for g,part in zip(gs, pi)]
        return self.__class__(self, self._labels, pi, f, gs)

    def change_labels(self, labels):
        """
        Return a relabelled structure.

        INPUT:

        - ``labels``, a list of labels.

        OUTPUT:

        A structure with the i-th label of self replaced with the i-th
        label of the list.

        EXAMPLES::

            sage: p = PermutationGroupElement((2,3))
            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: S = L.structures(['a','b','c']).list()
            sage: a = S[2]; a
            F-structure: {{'a', 'c'}, {'b'}}; G-structures: (('a', 'c'), ('b'))
            sage: a.change_labels([1,2,3])
            F-structure: {{1, 3}, {2}}; G-structures: [(1, 3), (2)]
        """
        f, gs = self._list
        pi = self._partition.change_labels(labels)
        f = f.change_labels(list(pi))
        g = [g.change_labels(part) for g,part in zip(gs, pi)]
        return self.__class__(self, labels, pi, f, g)


class CompositionSpecies(GenericCombinatorialSpecies, UniqueRepresentation):
    def __init__(self, F, G, min=None, max=None, weight=None):
        """
        Returns the composition of two species.

        EXAMPLES::

            sage: E = species.SetSpecies()
            sage: C = species.CycleSpecies()
            sage: S = E(C)
            sage: S.generating_series().coefficients(5)
            [1, 1, 1, 1, 1]
            sage: E(C) is S
            True

        TESTS::

            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: c = L.generating_series().coefficients(3)
            sage: L._check() #False due to isomorphism types not being implemented
            False
            sage: L == loads(dumps(L))
            True
        """
        self._F = F
        self._G = G
        self._name = "Composition of (%s) and (%s)"%(F, G)
        self._state_info = [F, G]
        GenericCombinatorialSpecies.__init__(self, min=None, max=None, weight=None)

    _default_structure_class = CompositionSpeciesStructure

    def _structures(self, structure_class, labels):
        """
        EXAMPLES::

            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: L.structures(['a','b','c']).list()
            [F-structure: {{'a', 'b', 'c'}}; G-structures: (('a', 'b', 'c'),),
             F-structure: {{'a', 'b', 'c'}}; G-structures: (('a', 'c', 'b'),),
             F-structure: {{'a', 'c'}, {'b'}}; G-structures: (('a', 'c'), ('b')),
             F-structure: {{'a', 'b'}, {'c'}}; G-structures: (('a', 'b'), ('c')),
             F-structure: {{'b', 'c'}, {'a'}}; G-structures: (('b', 'c'), ('a')),
             F-structure: {{'a'}, {'b'}, {'c'}}; G-structures: (('a'), ('b'), ('c'))]

        TESTS::

            sage: a = _[2]
            sage: f, gs = a._list
            sage: f
            {{'a', 'c'}, {'b'}}
            sage: f.parent()
            Set species
            sage: f._list
            [1, 2]
            sage: f._labels
            [{'a', 'c'}, {'b'}]
            sage: [g.parent() for g in gs]
            [Cyclic permutation species, Cyclic permutation species]
            sage: [g._labels for g in gs]
            [['a', 'c'], ['b']]
            sage: [g._list for g in gs]
            [[1, 2], [1]]
        """
        from itertools import product
        P = PartitionSpecies()
        for pi in P.structures(labels):
            # The labels of the G-structures will be just be the things
            # in labels
            gs = product(*[self._G.structures(part.label_subset()) for part in pi])

            # The labels of the F-structure will be set objects
            fs = self._F.structures(list(pi))
            for f, gg in product(fs, gs):
                yield structure_class(self, labels, pi, f, gg)

    def _isotypes(self, structure_class, labels):
        """
        EXAMPLES::

            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: L.isotypes(['a','b','c']).list()
            Traceback (most recent call last):
            ...
            NotImplementedError
        """
        raise NotImplementedError


    def _gs(self, series_ring, base_ring):
        """
        EXAMPLES::

            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: L.generating_series().coefficients(5)
            [1, 1, 1, 1, 1]
        """
        return self._F.generating_series(base_ring)(self._G.generating_series(base_ring))

    def _itgs(self, series_ring, base_ring):
        """
        EXAMPLES::

            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: L.isotype_generating_series().coefficients(10)
            [1, 1, 2, 3, 5, 7, 11, 15, 22, 30]
        """
        cis = self.cycle_index_series(base_ring)
        return cis.isotype_generating_series()

    def _cis(self, series_ring, base_ring):
        """
        EXAMPLES::

            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: L.cycle_index_series().coefficients(5)
            [p[],
             p[1],
             p[1, 1] + p[2],
             p[1, 1, 1] + p[2, 1] + p[3],
             p[1, 1, 1, 1] + p[2, 1, 1] + p[2, 2] + p[3, 1] + p[4]]

        Here we (indirectly) check that the cycle index series for
        permutations weighted by the number of cycles is correctly
        computed.

        ::

            sage: t = QQ['t'].gen()
            sage: E = species.SetSpecies()
            sage: C = species.CycleSpecies(weight=t)
            sage: S = E(C)
            sage: S.isotype_generating_series().coefficients(5) #indirect
            [1, t, t^2 + t, t^3 + t^2 + t, t^4 + t^3 + 2*t^2 + t]

        We do the same thing with set partitions weighed by the number of
        blocks.

        ::

            sage: t = QQ['t'].gen()
            sage: E = species.SetSpecies()
            sage: E_t = species.SetSpecies(min=1,weight=t)
            sage: Par = E(E_t)
            sage: Par.isotype_generating_series().coefficients(5)
            [1, t, t^2 + t, t^3 + t^2 + t, t^4 + t^3 + 2*t^2 + t]
        """
        f_cis = self._F.cycle_index_series(base_ring)
        g_cis = self._G.cycle_index_series(base_ring)

        #If G is a weighted species, then we can't use the default
        #algorithm for the composition of the cycle index series
        #since we must raise the weighting to the power.
        if self._G.is_weighted():
            return f_cis.weighted_composition(self._G)
        return f_cis(g_cis)

    def weight_ring(self):
        """
        Returns the weight ring for this species. This is determined by
        asking Sage's coercion model what the result is when you multiply
        (and add) elements of the weight rings for each of the operands.

        EXAMPLES::

            sage: E = species.SetSpecies(); C = species.CycleSpecies()
            sage: L = E(C)
            sage: L.weight_ring()
            Rational Field
        """
        return self._common_parent([self._F.weight_ring(), self._G.weight_ring()])


CompositionSpecies_class = CompositionSpecies
