"""Benchmark dict get/store with constant keys to measure
_BINARY_OP_SUBSCR_DICT_KNOWN_HASH and _STORE_SUBSCR_DICT_KNOWN_HASH speedup.

Usage:
  PYTHON_JIT=0 ./python.exe bench_dict_known_hash.py
  PYTHON_JIT=1 ./python.exe bench_dict_known_hash.py
"""

import time
import os
import sys

N = 50_000_000


def bench_get_str(n):
    d = {'a': 1, 'b': 2, 'c': 3}
    x = 0
    for _ in range(n):
        x += d['a']
    return x


def bench_get_int(n):
    d = {1: 10, 2: 20, 3: 30}
    x = 0
    for _ in range(n):
        x += d[1]
    return x


def bench_store_str(n):
    d = {}
    for _ in range(n):
        d['a'] = 1
    return d


def bench_store_int(n):
    d = {}
    for _ in range(n):
        d[1] = 1
    return d


def bench_get_float(n):
    d = {1.5: 1, 2.5: 2, 3.5: 3}
    x = 0
    for _ in range(n):
        x += d[1.5]
    return x


def bench_store_float(n):
    d = {}
    for _ in range(n):
        d[1.5] = 1
    return d


def bench_get_complex(n):
    d = {1+2j: 1, 3+4j: 2}
    x = 0
    for _ in range(n):
        x += d[1+2j]
    return x


def bench_store_complex(n):
    d = {}
    for _ in range(n):
        d[1+2j] = 1
    return d


def bench_get_multi(n):
    d = {'a': 1, 1: 2, b'x': 3, (1, 2): 4}
    x = 0
    for _ in range(n):
        x += d['a'] + d[1] + d[b'x'] + d[(1, 2)]
    return x


class _Key:
    pass

_KEY = _Key()


def bench_get_obj(n):
    d = {_KEY: 1}
    x = 0
    for _ in range(n):
        x += d[_KEY]
    return x


def bench_store_obj(n):
    d = {}
    for _ in range(n):
        d[_KEY] = 1
    return d


def run_avg(name, func, n, runs=3):
    func(1000)  # warmup
    times = []
    for _ in range(runs):
        t0 = time.perf_counter()
        func(n)
        times.append(time.perf_counter() - t0)
    avg = sum(times) / len(times)
    rate = n / avg / 1e6
    print(f"  {name:25s}  {avg:.3f}s  ({rate:.1f}M iter/s)")
    return avg


if __name__ == "__main__":
    jit = "JIT" if os.environ.get("PYTHON_JIT", "1") == "1" else "no JIT"
    print(f"Python {sys.version.split()[0]} ({jit}), avg of 3 runs\n")

    run_avg("dict_get[str]",      bench_get_str,      N)
    run_avg("dict_get[int]",      bench_get_int,      N)
    run_avg("dict_get[float]",    bench_get_float,    N)
    run_avg("dict_get[complex]",  bench_get_complex,  N)
    run_avg("dict_store[str]",    bench_store_str,    N)
    run_avg("dict_store[int]",    bench_store_int,    N)
    run_avg("dict_store[float]",  bench_store_float,  N)
    run_avg("dict_store[complex]",bench_store_complex,N)
    run_avg("dict_get[multi]",    bench_get_multi,    N // 4)
    run_avg("dict_get[obj]",      bench_get_obj,      N)
    run_avg("dict_store[obj]",    bench_store_obj,    N)