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()

[((20, 1, 20, 28), 'devito.silencio', 'silencio', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((22, 1, 25, 2), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(22, 25, 22, 41): '"""shape,expected"""', (22, 43, 25, 1): '[((30, 30), 17), ((30, 30, 30), 21)]'}, {}), "('shape,expected', [((30, 30), 17), ((30, 30, 30), 21)])", False, 'import pytest\n'), ((63, 1, 63, 28), 'devito.silencio', 'silencio', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((31, 11, 31, 28), 'devito.Grid', 'Grid', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((33, 13, 33, 23), 'io.StringIO', 'StringIO', ({}, {}), '()', False, 'from io import StringIO\n'), ((34, 24, 34, 53), 'devito.logger.logging.StreamHandler', 'logging.StreamHandler', ({(34, 46, 34, 52): 'buffer'}, {}), '(buffer)', False, 'from devito.logger import logger, logging\n'), ((35, 4, 35, 40), 'devito.logger.logger.addHandler', 'logger.addHandler', ({(35, 22, 35, 39): 'temporary_handler'}, {}), '(temporary_handler)', False, 'from devito.logger import logger, logging\n'), ((37, 14, 37, 49), 'devito.Function', 'Function', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((39, 15, 39, 51), 'devito.Function', 'Function', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((41, 9, 41, 87), 'devito.Operator', 'Operator', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((55, 4, 55, 43), 'devito.logger.logger.removeHandler', 'logger.removeHandler', ({(55, 25, 55, 42): 'temporary_handler'}, {}), '(temporary_handler)', False, 'from devito.logger import logger, logging\n'), ((73, 13, 73, 23), 'io.StringIO', 'StringIO', ({}, {}), '()', False, 'from io import StringIO\n'), ((74, 24, 74, 53), 'devito.logger.logging.StreamHandler', 'logging.StreamHandler', ({(74, 46, 74, 52): 'buffer'}, {}), '(buffer)', False, 'from devito.logger import logger, logging\n'), ((75, 4, 75, 40), 'devito.logger.logger.addHandler', 'logger.addHandler', ({(75, 22, 75, 39): 'temporary_handler'}, {}), '(temporary_handler)', False, 'from devito.logger import logger, logging\n'), ((78, 11, 78, 28), 'devito.Grid', 'Grid', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((83, 14, 83, 49), 'devito.Function', 'Function', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((85, 15, 85, 51), 'devito.Function', 'Function', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((87, 9, 87, 67), 'devito.Operator', 'Operator', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((110, 4, 110, 43), 'devito.logger.logger.removeHandler', 'logger.removeHandler', ({(110, 25, 110, 42): 'temporary_handler'}, {}), '(temporary_handler)', False, 'from devito.logger import logger, logging\n'), ((97, 18, 97, 72), 'devito.TimeFunction', 'TimeFunction', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((100, 19, 100, 74), 'devito.TimeFunction', 'TimeFunction', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((103, 13, 103, 71), 'devito.Operator', 'Operator', (), '', False, 'from devito import Grid, Function, TimeFunction, Eq, Operator, configuration, silencio\n'), ((38, 32, 38, 50), 'functools.reduce', 'reduce', ({(38, 39, 38, 42): 'mul', (38, 44, 38, 49): 'shape'}, {}), '(mul, shape)', False, 'from functools import reduce\n'), ((84, 32, 84, 50), 'functools.reduce', 'reduce', ({(84, 39, 84, 42): 'mul', (84, 44, 84, 49): 'shape'}, {}), '(mul, shape)', False, 'from functools import reduce\n'), ((98, 36, 98, 62), 'functools.reduce', 'reduce', ({(98, 43, 98, 46): 'mul', (98, 48, 98, 61): 'infield.shape'}, {}), '(mul, infield.shape)', False, 'from functools import reduce\n')]