# This file should be kept compatible with both Python 2.6 and Python >= 3.0.

from __future__ import division
from __future__ import print_function

"""
ccbench, a Python concurrency benchmark.
"""

import time
import os
import sys
import itertools
import threading
import subprocess
import socket
from optparse import OptionParser, SUPPRESS_HELP
import platform

# Compatibility
try:
    xrange
except NameError:
    xrange = range

try:
    map = itertools.imap
except AttributeError:
    pass


THROUGHPUT_DURATION = 2.0

LATENCY_PING_INTERVAL = 0.1
LATENCY_DURATION = 2.0

BANDWIDTH_PACKET_SIZE = 1024
BANDWIDTH_DURATION = 2.0


def task_pidigits():
    """Pi calculation (Python)"""
    _map = map
    _count = itertools.count
    _islice = itertools.islice

    def calc_ndigits(n):
        # From http://shootout.alioth.debian.org/
        def gen_x():
            return _map(lambda k: (k, 4*k + 2, 0, 2*k + 1), _count(1))

        def compose(a, b):
            aq, ar, as_, at = a
            bq, br, bs, bt = b
            return (aq * bq,
                    aq * br + ar * bt,
                    as_ * bq + at * bs,
                    as_ * br + at * bt)

        def extract(z, j):
            q, r, s, t = z
            return (q*j + r) // (s*j + t)

        def pi_digits():
            z = (1, 0, 0, 1)
            x = gen_x()
            while 1:
                y = extract(z, 3)
                while y != extract(z, 4):
                    z = compose(z, next(x))
                    y = extract(z, 3)
                z = compose((10, -10*y, 0, 1), z)
                yield y

        return list(_islice(pi_digits(), n))

    return calc_ndigits, (50, )

def task_regex():
    """regular expression (C)"""
    # XXX this task gives horrendous latency results.
    import re
    # Taken from the `inspect` module
    pat = re.compile(r'^(\s*def\s)|(.*(? 0:
        # Warm up
        func(*args)

        results = []
        loop = TimedLoop(func, args)
        ready = []
        ready_cond = threading.Condition()

        def run():
            with ready_cond:
                ready.append(None)
                ready_cond.notify()
            with start_cond:
                while not started:
                    start_cond.wait()
            loop(start_time, duration * 1.5, end_event, do_yield=False)

        for i in range(nthreads):
            threads.append(threading.Thread(target=run))
        for t in threads:
            t.setDaemon(True)
            t.start()
        # Wait for threads to be ready
        with ready_cond:
            while len(ready) < nthreads:
                ready_cond.wait()

    # Run the client and wait for the first ping(s) to arrive before
    # unblocking the background threads.
    chunks = []
    process = run_latency_client(addr=sock.getsockname(),
                                 nb_pings=nb_pings, interval=interval)
    s = _recv(sock, 4096)
    _time = time.time

    with start_cond:
        start_time = _time()
        started = True
        start_cond.notify(nthreads)

    while LAT_END not in s:
        s = _recv(sock, 4096)
        t = _time()
        chunks.append((t, s))

    # Tell the background threads to stop.
    end_event.append(None)
    for t in threads:
        t.join()
    process.wait()
    sock.close()

    for recv_time, chunk in chunks:
        # NOTE: it is assumed that a line sent by a client wasn't received
        # in two chunks because the lines are very small.
        for line in chunk.splitlines():
            line = line.strip()
            if line and line != LAT_END:
                send_time = eval(line)
                assert isinstance(send_time, float)
                results.append((send_time, recv_time))

    return results

def run_latency_tests(max_threads):
    for task in latency_tasks:
        print("Background CPU task:", task.__doc__)
        print()
        func, args = task()
        nthreads = 0
        while nthreads  0:
            # Warm up
            func(*args)

            results = []
            loop = TimedLoop(func, args)
            ready = []
            ready_cond = threading.Condition()

            def run():
                with ready_cond:
                    ready.append(None)
                    ready_cond.notify()
                with start_cond:
                    while not started:
                        start_cond.wait()
                loop(start_time, duration * 1.5, end_event, do_yield=False)

            for i in range(nthreads):
                threads.append(threading.Thread(target=run))
            for t in threads:
                t.setDaemon(True)
                t.start()
            # Wait for threads to be ready
            with ready_cond:
                while len(ready) < nthreads:
                    ready_cond.wait()

        # Run the client and wait for the first packet to arrive before
        # unblocking the background threads.
        process = run_bandwidth_client(addr=addr,
                                       packet_size=packet_size,
                                       duration=duration)
        _time = time.time
        # This will also wait for the parent to be ready
        s = _recv(sock, packet_size)
        remote_addr = eval(s.partition('#')[0])

        with start_cond:
            start_time = _time()
            started = True
            start_cond.notify(nthreads)

        n = 0
        first_time = None
        while not end_event and BW_END not in s:
            _sendto(sock, s, remote_addr)
            s = _recv(sock, packet_size)
            if first_time is None:
                first_time = _time()
            n += 1
        end_time = _time()

    end_event.append(None)
    for t in threads:
        t.join()
    process.kill()

    return (n - 1) / (end_time - first_time)

def run_bandwidth_tests(max_threads):
    for task in bandwidth_tasks:
        print("Background CPU task:", task.__doc__)
        print()
        func, args = task()
        nthreads = 0
        baseline_speed = None
        while nthreads