#!/usr/bin/python3

"""Tests different algorithms to count the number of 1 bits in an integer

Python 3.3, Linux x86_64, Intel Core i5 CPU M 450 @ 2.40GHz:

str_count        6.04 µs
func_count      43.54 µs
naive_count     19.24 µs
wegner_count     9.21 µs
lut_count        5.21 µs
"""

import timeit
import time
import random


def str_count(numbers):
    """Uses string operations"""
    return [bin(number).count('1') for number in numbers]


def func_count(numbers):
    """Naive counting using bitwise operations and just one statement"""
    return [sum((number & 1 << bit) >> bit
                for bit
                in range(number.bit_length()))
            for number in numbers]


def naive_count(numbers):
    """Naive counting in a loop"""
    def gen():
        for number in numbers:
            count = 0
            while number:
                count += number & 1
                number >>= 1
            yield count
    return list(gen())


def wegner_count(numbers):
    """Method of Peter Wegner in CACM 3 (1960), 322"""
    def gen():
        for number in numbers:
            count = 0
            while number:
                count += 1
                number &= number - 1
            yield count
    return list(gen())


def _gen_lut():
    """Generates a lookup table for lut_count"""
    lut = [0] * 256
    for index in range(256):
        lut[index] = (index & 1) + lut[index >> 1]
    return tuple(lut)


_lut = _gen_lut()


def lut_count(numbers):
    """Uses a lookup table

    See http://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetTable
    """
    lut = _lut
    def gen():
        for number in numbers:
            count = 0
            while number:
                count += lut[number & 0xff]
                number >>= 8
            yield count
    return list(gen())


def _gen_nums():
    """Creates random integers emphasizing numbers < 1000"""
    distribution = random.expovariate
    lambda_ = 0.001
    return [int(distribution(lambda_)) for _ in range(10000)]


samples = _gen_nums()


def test(methods):
    """Checks that given methods for equality against samples"""
    results = (method(samples) for method in methods)
    trusted = next(results)
    wrongs = (method.__name__ for method, result in zip(methods[1:], results)
              if result != trusted)
    for name in wrongs:
        print("%s is wrong" % name)


def main():
    methods = (str_count, func_count, naive_count, wegner_count, lut_count)
    test(methods)
    methodnames = [method.__name__ for method in methods]
    setups = ("from __main__ import %s, samples" % method
              for method in methodnames)
    statements = ("%s(samples)" % method for method in methodnames)
    number = 10000000 // len(samples)
    times = (timeit.timeit(setup = setup,
                           stmt = statement,
                           number = number,
                           timer = time.process_time,
                       )
             for setup, statement in zip(setups, statements))
    results = ("%s\t%.2f µs" % (method, time) for method, time
               in zip(methodnames, times))
    for result in results:
        print(result)


if __name__ == '__main__':
    main()