import itertools

from itertools import *
from timeit import Timer


# From the itertools recipe
def roundrobin(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    pending = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while pending:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            pending -= 1
            nexts = cycle(islice(nexts, pending))

def bunslow(*iters):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    sentinel = object()
    for tup in itertools.zip_longest(*iters, fillvalue=sentinel):
        yield from (x for x in tup if x is not sentinel)


def terry(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    nexts = cycle(iter(it).__next__ for it in iterables)
    for reduced_len in reversed(range(1, len(iterables))):
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            nexts = cycle(islice(nexts, reduced_len))


from collections import deque

def stackoverflow(*iterators):
    queue = deque(iter(it) for it in iterators)
    while len(queue) > 0:
        iterator = queue.popleft()
        try:
            yield next(iterator)
            queue.append(iterator)
        except StopIteration:
            pass



FUNCTIONS = (roundrobin, bunslow, terry, stackoverflow)

equal_data = [range(500)]*4
unequal_data = [range(500), range(600), range(700), range(400)]
extreme_data = [range(1), range(1), range(1), range(500)]

def trial(data_name, number, repeat):
    print("testing with %s" % data_name)
    template = 'from __main__ import %s as rr, %s as data'
    data = globals()[data_name]
    assert len(data) == 4 and all(type(x) is range for x in data)
    expected = list(roundrobin(*data))
    failed = []
    for func in FUNCTIONS:
        actual = list(func(*data))
        if actual != expected:
            failed.append(func)
        setup = template % (func.__name__, data_name)
        T = Timer('list(rr(*data))', setup=setup)
        t = min(T.repeat(number=number, repeat=repeat))
        print('  %-15s %f' % (func.__name__, t))
    for f in failed:
        print('  *** %s returned incorrect results ***' % f.__name__)


NUMBER = 30
REPEAT = 5

trial('equal_data', NUMBER, REPEAT)
print("*"*50)
trial('unequal_data', NUMBER, REPEAT)
print("*"*50)
trial('extreme_data', NUMBER, REPEAT)