#!/usr/bin/env python3
import json
import gzip
import time

import brotli
import msgpack
import tabulate
import umsgpack
import zstd

DATA = {
    'a': 'aardvark',
    'map': {str(i): i for i in range(1000)},
    'nums': list(range(1000)),
}

ITERATIONS = 500


def benchmark(to_func, from_func):
    start = time.time()
    for _ in range(ITERATIONS):
        serialized = to_func(DATA)
    to_seconds = time.time() - start
    start = time.time()
    for _ in range(ITERATIONS):
        from_func(serialized)
    from_seconds = time.time() - start
    return to_seconds, from_seconds, len(serialized)


def main():
    results = []

    funcs = {
        'json': [
            lambda x: json.dumps(x).encode(),
            json.loads,
        ],
        'msgpack': [
            msgpack.packb,
            msgpack.unpackb,
        ],
        'umsgpack': [
            umsgpack.packb,
            umsgpack.unpackb,
        ]
    }

    for name, (to_func, from_func) in funcs.items():
        results.append([name, *benchmark(to_func, from_func)])

        for c in [brotli, gzip, zstd]:
            results.append([
                f'{name}-{c.__name__}', *benchmark(
                    lambda d: c.compress(to_func(d)),
                    lambda d: from_func(c.decompress(d)))
            ])

    results = [
        (r[0], f'{r[1]:.2f}', f'{r[2]:.2f}', r[3])
        for r in results
    ]
    results = sorted(results, key=lambda x: x[0])
    print(
        tabulate.tabulate(
            results,
            headers=['format', 'to_seconds', 'from_seconds', 'byte size']))


if __name__ == '__main__':
    main()