import multitask


def setup(N):
    in_queue = multitask.Queue()
    out_queue = init(N-1, in_queue)
    return in_queue, out_queue


def run(N, M, in_queue, out_queue):
    yield send_message(out_queue, 0)
    count = (yield manager(M-1, in_queue, out_queue))
    assert count == (N * M)  # Check for correct number of messages


def init(N, out_queue):
    while N > 0:
        in_queue = multitask.Queue()
        multitask.add(worker(in_queue, out_queue))
        out_queue = in_queue
        N -= 1
    return out_queue


def manager(M, in_queue, out_queue):
    while M > 0:
        yield relay_message(in_queue, out_queue)
        M -= 1
    msg = (yield in_queue.get())
    raise StopIteration(msg)


def worker(in_queue, out_queue):
    while True:
        yield relay_message(in_queue, out_queue)


def send_message(queue, msg):
    yield queue.put(msg + 1)


def relay_message(in_queue, out_queue):
    msg = (yield in_queue.get())
    yield send_message(out_queue, msg)


if __name__ == '__main__':
    from ring_common import get_args, get_times, print_results

    N, M = get_args()

    (in_queue, out_queue), setup_cpu, setup_elapsed = get_times(setup, N)

    multitask.add(run(N, M, in_queue, out_queue))
    run_cpu, run_elapsed = get_times(multitask.run)[1:]

    print_results(N, M, setup_cpu, setup_elapsed, run_cpu, run_elapsed)