from __future__ import absolute_import

from functools import reduce
from operator import mul
try:
    from StringIO import StringIO
except ImportError:
    # Python3 compatibility
    from io import StringIO

import pytest
from conftest import skipif_yask

import numpy as np

from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio
from devito.logger import logger, logging


@silencio(log_level='DEBUG')
@skipif_yask
@pytest.mark.parametrize("shape,expected", [
    ((30, 30), 17),
    ((30, 30, 30), 21)
])
def test_at_is_actually_working(shape, expected):
    """
    Check that autotuning is actually running when switched on,
    in both 2D and 3D operators.
    """
    grid = Grid(shape=shape)

    buffer = StringIO()
    temporary_handler = logging.StreamHandler(buffer)
    logger.addHandler(temporary_handler)

    infield = Function(name='infield', grid=grid)
    infield.data[:] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape)
    outfield = Function(name='outfield', grid=grid)
    stencil = Eq(outfield.indexify(), outfield.indexify() + infield.indexify()*3.0)
    op = Operator(stencil, dle=('blocking', {'blockinner': True, 'blockalways': True}))

    # Expected 3 AT attempts for the given shape
    op(infield=infield, outfield=outfield, autotune=True)
    out = [i for i in buffer.getvalue().split('\n') if 'took' in i]
    assert len(out) == 4

    # Now try the same with aggressive autotuning, which tries 9 more cases
    configuration['autotuning'] = 'aggressive'
    op(infield=infield, outfield=outfield, autotune=True)
    out = [i for i in buffer.getvalue().split('\n') if 'took' in i]
    assert len(out) == expected
    configuration['autotuning'] = configuration._defaults['autotuning']

    logger.removeHandler(temporary_handler)

    temporary_handler.flush()
    temporary_handler.close()
    buffer.flush()
    buffer.close()


@silencio(log_level='DEBUG')
@skipif_yask
def test_timesteps_per_at_run():
    """
    Check that each autotuning run (ie with a given block shape) takes
    ``autotuning.core.options['at_squeezer']`` timesteps, for an operator
    performing the increment ``a[t + timeorder, ...] = f(a[t, ...], ...)``.
    """
    from devito.core.autotuning import options

    buffer = StringIO()
    temporary_handler = logging.StreamHandler(buffer)
    logger.addHandler(temporary_handler)

    shape = (30, 30, 30)
    grid = Grid(shape=shape)
    x, y, z = grid.dimensions
    t = grid.stepping_dim

    # Function
    infield = Function(name='infield', grid=grid)
    infield.data[:] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape)
    outfield = Function(name='outfield', grid=grid)
    stencil = Eq(outfield.indexify(), outfield.indexify() + infield.indexify()*3.0)
    op = Operator(stencil, dle=('blocking', {'blockalways': True}))
    op(infield=infield, outfield=outfield, autotune=True)
    out = [i for i in buffer.getvalue().split('\n') if 'took' in i]
    assert len(out) == 4
    assert all('in 1 timesteps' in i for i in out)
    buffer.truncate(0)

    # TimeFunction with increasing time order; increasing the time order
    # shouldn't affect how many iterations the autotuner is gonna run
    for to in [1, 2, 4]:
        infield = TimeFunction(name='infield', grid=grid, time_order=to)
        infield.data[:] = np.arange(reduce(mul, infield.shape),
                                    dtype=np.int32).reshape(infield.shape)
        outfield = TimeFunction(name='outfield', grid=grid, time_order=to)
        stencil = Eq(outfield.indexed[t + to, x, y, z],
                     outfield.indexify() + infield.indexify()*3.0)
        op = Operator(stencil, dle=('blocking', {'blockalways': True}))
        op(infield=infield, outfield=outfield, t=2, autotune=True)
        out = [i for i in buffer.getvalue().split('\n') if 'took' in i]
        assert len(out) == 4
        assert all('in %d timesteps' % options['at_squeezer'] in i for i in out)
        buffer.truncate(0)

    logger.removeHandler(temporary_handler)

    temporary_handler.flush()
    temporary_handler.close()
    buffer.flush()
    buffer.close()

[((410, 437), 'devito.silencio', 'silencio', ([], {'log_level': '"""DEBUG"""'}), "(log_level='DEBUG')\n", (418, 437), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((452, 531), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""shape,expected"""', '[((30, 30), 17), ((30, 30, 30), 21)]'], {}), "('shape,expected', [((30, 30), 17), ((30, 30, 30), 21)])\n", (475, 531), False, 'import pytest\n'), ((1914, 1941), 'devito.silencio', 'silencio', ([], {'log_level': '"""DEBUG"""'}), "(log_level='DEBUG')\n", (1922, 1941), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((716, 733), 'devito.Grid', 'Grid', ([], {'shape': 'shape'}), '(shape=shape)\n', (720, 733), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((748, 758), 'io.StringIO', 'StringIO', ([], {}), '()\n', (756, 758), False, 'from io import StringIO\n'), ((783, 812), 'devito.logger.logging.StreamHandler', 'logging.StreamHandler', (['buffer'], {}), '(buffer)\n', (804, 812), False, 'from devito.logger import logger, logging\n'), ((817, 853), 'devito.logger.logger.addHandler', 'logger.addHandler', (['temporary_handler'], {}), '(temporary_handler)\n', (834, 853), False, 'from devito.logger import logger, logging\n'), ((869, 904), 'devito.Function', 'Function', ([], {'name': '"""infield"""', 'grid': 'grid'}), "(name='infield', grid=grid)\n", (877, 904), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((1003, 1039), 'devito.Function', 'Function', ([], {'name': '"""outfield"""', 'grid': 'grid'}), "(name='outfield', grid=grid)\n", (1011, 1039), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((1133, 1211), 'devito.Operator', 'Operator', (['stencil'], {'dle': "('blocking', {'blockinner': True, 'blockalways': True})"}), "(stencil, dle=('blocking', {'blockinner': True, 'blockalways': True}))\n", (1141, 1211), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((1772, 1811), 'devito.logger.logger.removeHandler', 'logger.removeHandler', (['temporary_handler'], {}), '(temporary_handler)\n', (1792, 1811), False, 'from devito.logger import logger, logging\n'), ((2286, 2296), 'io.StringIO', 'StringIO', ([], {}), '()\n', (2294, 2296), False, 'from io import StringIO\n'), ((2321, 2350), 'devito.logger.logging.StreamHandler', 'logging.StreamHandler', (['buffer'], {}), '(buffer)\n', (2342, 2350), False, 'from devito.logger import logger, logging\n'), ((2355, 2391), 'devito.logger.logger.addHandler', 'logger.addHandler', (['temporary_handler'], {}), '(temporary_handler)\n', (2372, 2391), False, 'from devito.logger import logger, logging\n'), ((2429, 2446), 'devito.Grid', 'Grid', ([], {'shape': 'shape'}), '(shape=shape)\n', (2433, 2446), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((2533, 2568), 'devito.Function', 'Function', ([], {'name': '"""infield"""', 'grid': 'grid'}), "(name='infield', grid=grid)\n", (2541, 2568), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((2667, 2703), 'devito.Function', 'Function', ([], {'name': '"""outfield"""', 'grid': 'grid'}), "(name='outfield', grid=grid)\n", (2675, 2703), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((2797, 2855), 'devito.Operator', 'Operator', (['stencil'], {'dle': "('blocking', {'blockalways': True})"}), "(stencil, dle=('blocking', {'blockalways': True}))\n", (2805, 2855), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((4013, 4052), 'devito.logger.logger.removeHandler', 'logger.removeHandler', (['temporary_handler'], {}), '(temporary_handler)\n', (4033, 4052), False, 'from devito.logger import logger, logging\n'), ((3268, 3322), 'devito.TimeFunction', 'TimeFunction', ([], {'name': '"""infield"""', 'grid': 'grid', 'time_order': 'to'}), "(name='infield', grid=grid, time_order=to)\n", (3280, 3322), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((3481, 3536), 'devito.TimeFunction', 'TimeFunction', ([], {'name': '"""outfield"""', 'grid': 'grid', 'time_order': 'to'}), "(name='outfield', grid=grid, time_order=to)\n", (3493, 3536), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((3673, 3731), 'devito.Operator', 'Operator', (['stencil'], {'dle': "('blocking', {'blockalways': True})"}), "(stencil, dle=('blocking', {'blockalways': True}))\n", (3681, 3731), False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((937, 955), 'functools.reduce', 'reduce', (['mul', 'shape'], {}), '(mul, shape)\n', (943, 955), False, 'from functools import reduce\n'), ((2601, 2619), 'functools.reduce', 'reduce', (['mul', 'shape'], {}), '(mul, shape)\n', (2607, 2619), False, 'from functools import reduce\n'), ((3359, 3385), 'functools.reduce', 'reduce', (['mul', 'infield.shape'], {}), '(mul, infield.shape)\n', (3365, 3385), False, 'from functools import reduce\n')]