"""
Benchmarks for ``@vectorize`` ufuncs.
"""

import numpy as np


def setup():
    global mul, rel_diff

    from numba import vectorize

    @vectorize(["float32(float32, float32)",
                "float64(float64, float64)",
                "complex64(complex64, complex64)",
                "complex128(complex128, complex128)"])
    def mul(x, y):
        return x * y

    @vectorize(["float32(float32, float32)",
                "float64(float64, float64)"])
    def rel_diff(x, y):
        # XXX for float32 performance, we should write `np.float32(2)`, but
        # that's not the natural way to write this code...
        return 2 * (x - y) / (x + y)


class Vectorize:

    n = 10000
    dtypes = ('float32', 'float64', 'complex64', 'complex128')

    def setup(self):
        setup()
        self.samples = {}
        self.out = {}
        for dtype in self.dtypes:
            self.samples[dtype] = np.linspace(0.1, 1, self.n, dtype=dtype)
            self.out[dtype] = np.zeros(self.n, dtype=dtype)

    def _binary_func(fname, dtype, name):
        def f(self):
            func = globals()[fname]
            func(self.samples[dtype], self.samples[dtype], self.out[dtype])
        f.__name__ = name
        return f

    for dtype in dtypes:
        fname = 'time_mul_%s' % dtype
        locals()[fname] = _binary_func('mul', dtype, fname)

    time_rel_diff_float32 = _binary_func('rel_diff', 'float32', 'time_rel_diff_float32')
    time_rel_diff_float64 = _binary_func('rel_diff', 'float64', 'time_rel_diff_float64')

    del _binary_func
