"""Micro-benchmark for np.sum/prod/max/any dispatcher overhead."""
import numpy as np
import pyperf


class MyArr(np.ndarray):
    pass


def make_setup():
    return {
        "array_4":    np.arange(4, dtype=np.float64),
        "array_20":   np.arange(20, dtype=np.float64),
        "int_20":     np.arange(20, dtype=np.int64),
        "bool_20":    np.ones(20, dtype=bool),
        "arr_10x10":  np.arange(100, dtype=np.float64).reshape(10, 10),
        "arr_2x2f":   np.arange(4, dtype=np.float64).reshape(2, 2),
        "arr_2x2i":   np.arange(4, dtype=np.int64).reshape(2, 2),
        "subclass_20":np.arange(20, dtype=np.float64).view(MyArr),
    }


BENCHES = [
    ("np.sum(array_4)",         "np.sum(a['array_4'])"),
    ("np.sum(array_20)",        "np.sum(a['array_20'])"),
    ("np.prod(array_20)",       "np.prod(a['array_20'])"),
    ("np.max(array_20)",        "np.max(a['array_20'])"),
    ("np.sum(int_20)",          "np.sum(a['int_20'])"),
    ("np.sum(10x10)",           "np.sum(a['arr_10x10'])"),
    ("np.sum(10x10,axis=0)",    "np.sum(a['arr_10x10'], axis=0)"),
    ("np.sum(10x10,axis=1)",    "np.sum(a['arr_10x10'], axis=1)"),
    ("np.any(bool_20)",         "np.any(a['bool_20'])"),
    ("np.sum(2x2_float)",       "np.sum(a['arr_2x2f'])"),
    ("np.sum(2x2_int)",         "np.sum(a['arr_2x2i'])"),
    ("np.sin(array_20)",        "np.sin(a['array_20'])"),
    ("np.sum(subclass_20)",     "np.sum(a['subclass_20'])"),
    # cheap view ops migrated in the mass tuple-spec conversion
    ("np.transpose(10x10)",     "np.transpose(a['arr_10x10'])"),
    ("np.ravel(array_20)",      "np.ravel(a['array_20'])"),
]


def main():
    runner = pyperf.Runner()
    runner.metadata['numpy_version'] = np.__version__
    a = make_setup()
    for name, expr in BENCHES:
        runner.timeit(name=name, stmt=expr, globals={'np': np, 'a': a})


if __name__ == '__main__':
    main()