import numpy as np
import bottleneck as bn
from .autotimeit import autotimeit

__all__ = ['bench']


def bench(shapes=[(100,), (1000, 1000), (1000, 1000), (1000, 1000),
                  (1000, 1000)],
          axes=[0, 0, 0, 1, 1],
          nans=[False, False, True, False, True],
          dtype='float64',
          order='C',
          functions=None):
    """
    Bottleneck benchmark.

    Parameters
    ----------
    shapes : list, optional
        A list of tuple shapes of input arrays to use in the benchmark.
    axes : list, optional
        List of axes along which to perform the calculations that are being
        benchmarked.
    nans : list, optional
        A list of the bools (True or False), one for each tuple in the
        `shapes` list, that tells whether the input arrays should be randomly
        filled with one-fifth NaNs.
    dtype : str, optional
        Data type string such as 'float64', which is the default.
    order : {'C', 'F'}, optional
        Whether to store multidimensional data in C- or Fortran-contiguous
        (row- or column-wise) order in memory.
    functions : {list, None}, optional
        A list of strings specifying which functions to include in the
        benchmark. By default (None) all functions are included in the
        benchmark.

    Returns
    -------
    A benchmark report is printed to stdout.

    """

    if len(shapes) != len(nans):
        raise ValueError("`shapes` and `nans` must have the same length")
    if len(shapes) != len(axes):
        raise ValueError("`shapes` and `axes` must have the same length")

    # Header
    print('Bottleneck performance benchmark')
    print("    Bottleneck %s; Numpy %s" % (bn.__version__, np.__version__))
    print("    Speed is NumPy time divided by Bottleneck time")
    print("    NaN means approx one-fifth NaNs; %s used" % str(dtype))

    print('')
    header = [" "*11]
    for nan in nans:
        if nan:
            header.append("NaN".center(11))
        else:
            header.append("no NaN".center(11))
    print("".join(header))
    header = ["".join(str(shape).split(" ")).center(11) for shape in shapes]
    header = [" "*12] + header
    print("".join(header))
    header = ["".join(("axis=" + str(axis)).split(" ")).center(11)
              for axis in axes]
    header = [" "*12] + header
    print("".join(header))

    suite = benchsuite(shapes, dtype, nans, axes, order, functions)
    for test in suite:
        name = test["name"].ljust(12)
        fmt = name + "%7.1f" + "%11.1f"*(len(shapes) - 1)
        speed = timer(test['statements'], test['setups'])
        print(fmt % tuple(speed))


def timer(statements, setups):
    speed = []
    if len(statements) != 2:
        raise ValueError("Two statements needed.")
    for setup in setups:
        with np.errstate(invalid='ignore'):
            t0 = autotimeit(statements[0], setup)
            t1 = autotimeit(statements[1], setup)
        speed.append(t1 / t0)
    return speed


def getarray(shape, dtype, nans, order):
    a = np.arange(np.prod(shape), dtype=dtype)
    if nans and issubclass(a.dtype.type, np.inexact):
        a[::5] = np.nan
    rs = np.random.RandomState(shape)
    rs.shuffle(a)
    return np.array(a.reshape(*shape), order=order)


def benchsuite(shapes, dtype, nans, axes, order, functions):

    suite = []

    def getsetups(setup, shapes, nans, axes, dtype, order):
        template = """
        from bottleneck.benchmark.bench import getarray
        a = getarray(%s, '%s', %s, '%s')
        axis=%s
        %s"""
        setups = []
        for shape, axis, nan in zip(shapes, axes, nans):
            s = template % (str(shape), str(dtype), str(nan), str(order),
                            str(axis), setup)
            s = '\n'.join([line.strip() for line in s.split('\n')])
            setups.append(s)
        return setups

    # non-moving window functions
    funcs = bn.get_functions("reduce", as_string=True)
    funcs += ['rankdata', 'nanrankdata']
    for func in funcs:
        if functions is not None and func not in functions:
            continue
        run = {}
        run['name'] = func
        run['statements'] = ["bn_func(a, axis)", "sl_func(a, axis)"]
        setup = """
            from bottleneck import %s as bn_func
            try: from numpy import %s as sl_func
            except ImportError: from bottleneck.slow import %s as sl_func
            if "%s" == "median": from bottleneck.slow import median as sl_func
        """ % (func, func, func, func)
        run['setups'] = getsetups(setup, shapes, nans, axes, dtype, order)
        suite.append(run)

    # partition, argpartition
    funcs = ['partition', 'argpartition']
    for func in funcs:
        if functions is not None and func not in functions:
            continue
        run = {}
        run['name'] = func
        run['statements'] = ["bn_func(a, n, axis)",
                             "sl_func(a, n, axis)"]
        setup = """
            from bottleneck import %s as bn_func
            from bottleneck.slow import %s as sl_func
            if axis is None: n = a.size
            else: n = a.shape[axis] - 1
            n = max(n / 2, 0)
        """ % (func, func)
        run['setups'] = getsetups(setup, shapes, nans, axes, dtype, order)
        suite.append(run)

    # replace, push
    funcs = ['replace', 'push']
    for func in funcs:
        if functions is not None and func not in functions:
            continue
        run = {}
        run['name'] = func
        if func == 'replace':
            run['statements'] = ["bn_func(a, nan, 0)",
                                 "slow_func(a, nan, 0)"]
        elif func == 'push':
            run['statements'] = ["bn_func(a, 5, axis)",
                                 "slow_func(a, 5, axis)"]
        else:
            raise ValueError('Unknow function name')
        setup = """
            from numpy import nan
            from bottleneck import %s as bn_func
            from bottleneck.slow import %s as slow_func
        """ % (func, func)
        run['setups'] = getsetups(setup, shapes, nans, axes, dtype, order)
        suite.append(run)

    # moving window functions
    funcs = bn.get_functions('move', as_string=True)
    for func in funcs:
        if functions is not None and func not in functions:
            continue
        run = {}
        run['name'] = func
        run['statements'] = ["bn_func(a, w, 1, axis)",
                             "sw_func(a, w, 1, axis)"]
        setup = """
            from bottleneck.slow.move import %s as sw_func
            from bottleneck import %s as bn_func
            w = a.shape[axis] // 5
        """ % (func, func)
        run['setups'] = getsetups(setup, shapes, nans, axes, dtype, order)
        suite.append(run)

    return suite