import inspect
import itertools

def chained(func):
    def meth(self, *args, **kwargs):
        result = func(self, *args, **kwargs)
        return self if result is None else result

    return meth


class pynq:
    def __init__(self, seq):
        self._seq = seq

    @staticmethod
    def _extractor_func(extractor, kwargs, _ident=lambda x: x):
        '''Return a tuple containing an extractor function and True
        if the function takes more than 1 parameter, else False. This
        second value indicates whether enumeration should take place.'''

        if extractor:
            if len(extractor) > 1:
                raise TypeError('Too many positional args. '
                                'Only 1 allowed.')
            extract = extractor[0]
        elif not kwargs:
            extract=_ident
        elif len(kwargs) != 1:
            raise TypeError('Too many named parameters. Only 1 allowed.')
        else:
            for var, code in kwargs.items():
                extract=eval('lambda {}: {}'.format(var, code))
                break

        sig = inspect.signature(extract)
        return extract, (len(sig.parameters) > 1)

    @staticmethod
    def _select_many(seq, extract):
        '''Deferred execution of select_many'''
        for item in seq:
            for elt in extract(item):
                yield elt

    @chained
    def select_many(self, *extractor, **kwargs):
        '''Flatten a sequence of sequences into a single stream.'''
        extract = pynq._extractor_func(extractor, kwargs)
        self._seq = pynq._select_many(self._seq, extract)

    @staticmethod
    def _sorted(seq, keyfunc):
        for item in sorted(seq, key=keyfunc):
            yield item

    @staticmethod
    def _group_by(seq, extract):
        '''Deferred execution of group_by'''
        for _, group in itertools.groupby(seq, key=extract):
            yield group

    @chained
    def group_by(self, *extractor, **kwargs):
        '''Sort the sequence, and group together items by the extractor
        function.

        >>> import pynq
        >>> z='z'
        >>> a='a'
        >>> L = [(0,z), (1,a), (2,z), (3,a), (4,z), (5,z)]
        >>> it = iter(pynq(L).group_by( k = 'k[1]' ))
        >>> len(next(it))
        2
        >>> len(next(it))
        4
        '''
        extract = pynq._extractor_func(extractor, kwargs)
        self._seq = pynq._sorted(self._seq, extract)
        self._seq = pynq._group_by(self._seq, extract)
        return self