import cloudpickle as pickle
import pytest
import numpy as np

from .conftest import skipif
from submodules.devito.conftest import check_array, get_arrays
from tests.common import init_chessboard
from devito import Eq, Grid, Function, TimeFunction, Operator, sin, cos
from devito.exceptions import CompilationError, InvalidOperator
from devito.finite_differences import Coefficient, Substitutions, Weights
from devito.ir import KernelLaunch, FindNodes, FindSymbols
from devito.symbolics import RoundUp
from devito.tools import filter_sorted, flatten
from devito.types import SpaceDimension
from devitopro import *  # noqa
from devitopro.types.enriched import NoLayers


# *** Preset optimization modes

def opt0(extra=None):
    if configuration['language'] in ('cuda', 'hip'):
        opt_options = {'deriv-schedule': 'basic',
                       'deriv-uniformize': False,
                       'cire-minmem': True}
    else:
        opt_options = {'deriv-schedule': 'basic',
                       'openmp': False,
                       'cire-minmem': True}
    opt_options.update(extra or {})
    return ('noop', opt_options)


def opt1(extra=None):
    if configuration['language'] in ('cuda', 'hip'):
        opt_options = {'deriv-schedule': 'basic',
                       'deriv-uniformize': False,
                       'gpu-opt': True,
                       'cire-minmem': True}
    else:
        opt_options = {'deriv-schedule': 'basic',
                       'expand': False,
                       'blocklevels': 0,
                       'openmp': False,
                       'cire-minmem': True}
    opt_options.update(extra or {})
    return ('advanced', opt_options)


def get_n_arrays(op, n):
    arrays = get_arrays(op)
    assert len(arrays) == n
    dimensions = filter_sorted(flatten(a.dimensions for a in arrays),
                               lambda i: i.name)
    sizes = [i.symbolic_size for i in dimensions]
    return arrays, sizes


def get_weights(op):
    functions = FindSymbols().visit(op)
    return sorted({f for f in functions if isinstance(f, Weights)},
                  key=lambda i: i.name)


def check_array_wrap(array, exp_halo, exp_shape, rotate=False):
    # Take PRO's autopadding into account and then resort to OSS' check_array
    assert configuration['autopadding']

    mmts = configuration['platform'].max_mem_trans_size(array.dtype)

    exp_shape = list(exp_shape)
    exp_shape[-1] = RoundUp(exp_shape[-1], mmts)

    check_array(array, exp_halo, tuple(exp_shape), rotate)


@skipif('device')
class Test2Pass:

    @skipif('sycl', reason="Produces compilation failures on SYCL")
    def test_mixed_derivs_v0(self):
        grid = Grid(shape=(10, 10, 10))

        u = TimeFunction(name='u', grid=grid, space_order=4)
        u1 = TimeFunction(name='u', grid=grid, space_order=4)

        eqns = Eq(u.forward, u.laplace + u.dx.dy + 1.)

        op0 = Operator(eqns, opt=opt0())
        op1 = Operator(eqns, opt=opt1({'cire-schedule': 1}))
        op2 = Operator(eqns, opt=opt1({'cire-schedule': 2}))

        # Three schedules available!
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 3}))

        # Check generated code
        arrays, (ys, zs) = get_n_arrays(op1, 1)
        check_array_wrap(arrays[0], ((2, 2), (0, 0)), (ys+4, zs))
        arrays, (ys, zs) = get_n_arrays(op2, 2)
        check_array_wrap(arrays[0], ((2, 2), (0, 0)), (ys+4, zs))
        check_array_wrap(arrays[1], ((2, 2), (0, 0)), (ys+4, zs))

        op0.apply(time_M=5)
        op2.apply(time_M=5, u=u1)

        assert np.allclose(u.data, u1.data, rtol=1e-4)

    @skipif('sycl', reason="Produces compilation failures on SYCL")
    def test_mixed_derivs_v1(self):
        grid = Grid(shape=(10, 10, 10))

        f = Function(name='f', grid=grid, space_order=4)
        u = TimeFunction(name='u', grid=grid, space_order=4)
        u1 = TimeFunction(name='u', grid=grid, space_order=4)

        eqns = Eq(u.forward, f*u.dx.dy + 1.)

        op0 = Operator(eqns, opt=opt0())
        op1 = Operator(eqns, opt=opt1({'cire-schedule': 1}))

        # Two schedules available!
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 2}))

        # Check generated code
        arrays, (ys, zs) = get_n_arrays(op1, 1)
        check_array_wrap(arrays[0], ((2, 2), (0, 0)), (ys+4, zs))

        op0.apply(time_M=5)
        op1.apply(time_M=5, u=u1)

        assert np.allclose(u.data, u1.data, rtol=1e-4)

    def test_same_inner_deriv_v0(self):
        grid = Grid(shape=(10, 10, 10))

        u = TimeFunction(name='u', grid=grid, space_order=8)

        eqns = Eq(u.forward, 2*u.dx.dx + 3*u.dx.dy + 4*u.dx.dz)

        op = Operator(eqns, opt=opt1())

        # Two schedules available!
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 3}))

        # Check generated code
        arrays, (ys, zs) = get_n_arrays(op, 2)
        check_array_wrap(arrays[0], ((4, 4), (4, 4)), (ys+8, zs+8))
        assert op._profiler._sections['section0'].sops == 93

    def test_same_inner_deriv_v1(self):
        grid = Grid(shape=(10, 10, 10))

        u = TimeFunction(name='u', grid=grid, space_order=8)

        eqns = Eq(u.forward, 2*(u.dx + u.dy).dx + 3*(u.dx + u.dy).dy)

        op = Operator(eqns, opt=opt1())

        # Three schedules available!
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 3}))

        # Check generated code
        arrays, (xs, ys, zs) = get_n_arrays(op, 1)
        check_array_wrap(arrays[0], ((4, 4), (4, 4), (0, 0)), (xs+8, ys+8, zs))
        assert op._profiler._sections['section0'].sops == 76

    def test_same_inner_deriv_v2(self):
        grid = Grid(shape=(10, 10, 10))

        f = Function(name='f', grid=grid, space_order=8)
        u = TimeFunction(name='u', grid=grid, space_order=8)

        eqns = Eq(u.forward, 2*(u.dx + u.dy).dx + 3*(u.dx + u.dy).dy + f*u.dx)

        op = Operator(eqns, opt=opt1())

        # Four schedules available!
        # * []
        # * [u.dx, u.dy]
        # * [u.dx + u.dy]
        # * [u.dx + u.dy, u.dx]
        Operator(eqns, opt=opt1({'cire-schedule': 3}))
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 4}))

        # Check generated code
        # By default, [u.dx + u.dy] is picked
        arrays, (xs, ys, zs) = get_n_arrays(op, 1)
        check_array_wrap(arrays[0], ((4, 4), (4, 4), (0, 0)), (xs+8, ys+8, zs))
        assert op._profiler._sections['section0'].sops == 96

        # We also check some of the other schedules
        op1 = Operator(eqns, opt=opt1({'cire-schedule': 1}))  # [u.dx, u.dy]
        arrays, (xs, ys, zs) = get_n_arrays(op1, 2)
        check_array_wrap(arrays[0], ((4, 4), (4, 4), (0, 0)), (xs+8, ys+8, zs))
        check_array_wrap(arrays[1], ((4, 4), (4, 4), (0, 0)), (xs+8, ys+8, zs))
        assert op1._profiler._sections['section0'].sops == 78

        op2 = Operator(eqns, opt=opt1({'cire-schedule': -1}))  # [u.dx + u.dy, u.dx]
        arrays, (xs, ys, zs) = get_n_arrays(op2, 2)
        check_array_wrap(arrays[0], ((4, 4), (4, 4), (0, 0)), (xs+8, ys+8, zs))
        check_array_wrap(arrays[1], ((4, 4), (4, 4), (0, 0)), (xs+8, ys+8, zs))
        assert op2._profiler._sections['section0'].sops == 78

    @skipif('sycl', reason="Produces compilation failures on SYCL")
    def test_tti_zhang_like(self):
        grid = Grid(shape=(31, 17, 25))

        so = 8
        so1 = so // 2

        e = Function(name='e', grid=grid, space_order=so)
        h = Function(name='h', grid=grid, space_order=so)
        p0 = TimeFunction(name='p0', grid=grid, space_order=so)

        e.data[:] = .02
        h.data[:] = .04

        def g3(field):
            return (sin(e) * cos(h) * field.dx(fd_order=so1) +
                    sin(e) * sin(h) * field.dy(fd_order=so1) +
                    cos(e) * field.dz(fd_order=so1))

        def g3t(field):
            return ((sin(e) * cos(h) * field).dx(fd_order=so1).T +
                    (sin(e) * sin(h) * field).dy(fd_order=so1).T +
                    (cos(e) * field).dz(fd_order=so1).T)

        eqns = [Eq(p0.forward, g3t(g3(p0)) + p0.laplace + 1.)]

        op0 = Operator(eqns, opt=opt0())
        init_chessboard(p0)
        op0.apply(time_M=3)

        for i in [0, 2, 5]:
            p1 = TimeFunction(name='p0', grid=grid, space_order=so)

            op1 = Operator(eqns, opt=opt1({'cire-schedule': i}))

            # Check generated code
            if i == 0:
                arrays, (xs, ys, zs) = get_n_arrays(op1, 3 + 0)  # sin/cos temps
                assert op1._profiler._sections['section1'].sops == 630
            elif i == 2:
                arrays, (xs, ys, zs) = get_n_arrays(op1, 3 + 1)
                assert op1._profiler._sections['section1'].sops == 140
            elif i == 5:
                arrays, (xs, ys, zs) = get_n_arrays(op1, 3 + 2)
            all(check_array_wrap(a, ((2, 2), (2, 2), (2, 2)), (xs+4, ys+4, zs+4))
                for a in arrays[3:])

            init_chessboard(p1)
            op1.apply(time_M=3, p0=p1)

            assert np.allclose(p0.data, p1.data, rtol=1.e-4)

        # 8 schedules available!
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 8}))

    @pytest.mark.parametrize('expr', [
        'sf*cg+u.dx2 + sf*sg*u.dy2',
        'sf*sf*cg*cg*u.dx2 + sf*sf*sg*sg*u.dy2',
    ])
    def test_noop_variant_only(self, expr):
        grid = Grid(shape=(10, 10, 10))

        f = Function(name='f', grid=grid)
        g = Function(name='g', grid=grid)
        u = TimeFunction(name='u', grid=grid, space_order=4)

        cf = cos(f)  # noqa
        sf = sin(f)  # noqa
        cg = cos(g)  # noqa
        sg = sin(g)  # noqa

        eqn = Eq(u.forward, eval(expr) + 1.)

        # NOTE: we use `cire-mingain` to prevent hoisting of sin and cos
        op = Operator(eqn, opt=opt1({'cire-mingain': 200}))

        # Only the 1-pass variant available, aka the "noop" variant
        op1 = Operator(eqn, opt=opt1({'cire-mingain': 200,
                                      'cire-schedule': -1}))
        assert str(op) == str(op1)

        # Check generated code
        arrays, _ = get_n_arrays(op, 0)

    def test_deep_unmixed_derivs(self):
        grid = Grid(shape=(10, 10, 10))

        f = Function(name='f', grid=grid, space_order=4)
        g = Function(name='g', grid=grid, space_order=4)
        u = TimeFunction(name='u', grid=grid, space_order=4)

        cf = cos(f)
        sf = sin(f)
        cg = cos(g)
        sg = sin(g)

        eqns = Eq(u.forward, (sf*cg*cg*u.dx2 +
                              sf*sg*sg*u.dy2 +
                              cf*cf*u.dz2 +
                              u.dx.dy + 1.))

        # NOTE: we use `cire-mingain` to prevent hoisting of sin and cos
        op = Operator(eqns, opt=opt1({'cire-mingain': 200,
                                      'cire-schedule': 1}))

        # Three schedules available!
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 3}))

        # Check generated code
        arrays, (ys, zs) = get_n_arrays(op, 1)
        check_array_wrap(arrays[0], ((2, 2), (0, 0)), (ys+4, zs))

        op1 = Operator(eqns, opt=opt1({'cire-mingain': 200,
                                       'cire-schedule': 2}))
        arrays, (ys, zs) = get_n_arrays(op1, 4)
        all(check_array_wrap(a, ((2, 2), (0, 0)), (ys+4, zs)) for a in arrays)

    @skipif('sycl', reason="Produces compilation failures on SYCL")
    def test_redundancies_outside_of_derivs_v0(self):
        grid = Grid(shape=(20, 20, 20))

        f = Function(name='f', grid=grid, space_order=8)
        g = Function(name='g', grid=grid, space_order=8)
        u = TimeFunction(name='u', grid=grid, space_order=8)

        cf = cos(f)
        sf = sin(f)
        cg = cos(g)
        sg = sin(g)

        eqns = Eq(u.forward, (sf*sf*cg*cg*u.dx2 +
                              sf*sf*sg*sg*u.dy2 +
                              cf*cf*u.dz2 +
                              2.*sf*sf*sg*cg*u.dy.dx +
                              2.*cf*sf*sg*u.dydz +
                              2.*cf*sf*cg*u.dxdz) + u.laplace + 1.)

        # NOTE: we use `cire-mingain` to prevent hoisting of sin and cos
        op = Operator(eqns, opt=opt1({'cire-mingain': 200}))

        # Check generated code
        arrays, (xs, ys, zs) = get_n_arrays(op, 2)
        all(check_array_wrap(a, ((4, 4), (0, 0), (4, 4)), (xs+8, ys, zs+8))
            for a in arrays)

        # Three schedules available!
        # []  (1 pass)
        # [u.dx, u.dy]  (the two inner derivs)
        # [u.dx2, u.dy2, u.dz2, u.dx, u.dy]
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 3}))

        # Also check the other 2-pass variant
        op1 = Operator(eqns, opt=opt1({'cire-mingain': 200,
                                       'cire-schedule': -1}))
        arrays, (xs, ys, zs) = get_n_arrays(op1, 6)
        all(check_array_wrap(a, ((4, 4), (0, 0), (4, 4)), (xs+8, ys, zs+8))
            for a in arrays)

        # We also check the output since op1 is a tricky bastard
        f.data[:] = .02
        g.data[:] = .04
        init_chessboard(u)

        u1 = TimeFunction(name='u', grid=grid, space_order=8)
        init_chessboard(u1)

        op.apply(time_M=5)
        op1.apply(time_M=5, u=u1)

        assert np.allclose(u.data, u1.data, rtol=1.e-4)

    def test_multiple_eqns(self):
        grid = Grid(shape=(10, 10, 10))

        u = TimeFunction(name='u', grid=grid, space_order=4)
        v = TimeFunction(name='v', grid=grid, space_order=4)

        expr = (v.dx2 + v.dz2 + v.dxdz)

        eqns = [Eq(u.forward, expr + 1.),
                Eq(v.forward, expr + 1.)]

        op = Operator(eqns, opt=opt1({'cire-schedule': -1}))

        # Check generated code
        arrays, (zs,) = get_n_arrays(op, 2)
        all(check_array_wrap(a, ((2, 2),), (zs+4,)) for a in arrays)
        assert op._profiler._sections['section0'].sops == 43

    def test_huge_eqns(self):
        grid = Grid(shape=(10, 10, 10))

        f = Function(name='f', grid=grid, space_order=8)
        g = Function(name='g', grid=grid, space_order=8)
        h = Function(name='h', grid=grid, space_order=8)
        u = TimeFunction(name='u', grid=grid, space_order=8)
        v = TimeFunction(name='v', grid=grid, space_order=8)

        sf = sin(f)
        cg = cos(g)
        sg = sin(g)

        expr0 = (cg*cg*u.dx2 + sg*sg*u.dy2 + u.dz2 +
                 sf*sg*cg*u.dy.dx + sf*sg*u.dydz + sf*cg*u.dxdz)
        expr1 = (cg*cg*v.dx2 + sg*sg*v.dy2 + v.dz2 +
                 sf*sg*cg*v.dy.dx + sf*sg*v.dydz + sf*cg*v.dxdz)

        eqns = [Eq(u.forward, -expr0 + h*expr0 + -expr1 + 1.)]

        op = Operator(eqns, opt=opt1())

        # Three schedules available!
        Operator(eqns, opt=opt1({'cire-schedule': 2}))
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 3}))

        # Check generated code
        arrays, (xs, ys, zs) = get_n_arrays(op, 3 + 4)
        all(check_array_wrap(a, ((4, 4), (0, 0), (4, 4)), (xs+8, ys, zs+8))
            for a in arrays[3:7])
        assert op._profiler._sections['section1'].sops == 335

    def test_largest_redundancy_v0(self):
        grid = Grid(shape=(31, 17, 25))

        f = Function(name='f', grid=grid, space_order=12)
        g = Function(name='g', grid=grid, space_order=12)
        h = Function(name='h', grid=grid, space_order=12)
        u = TimeFunction(name='u', grid=grid, space_order=12)

        sf = sin(f)

        eqns = [Eq(u.forward, (f*g + sf*h*u.dy).dx + 1.)]

        op = Operator(eqns, opt=opt1())

        # Three schedules available!
        with pytest.raises(CompilationError):
            Operator(eqns, opt=opt1({'cire-schedule': 3}))

        # Check generated code
        get_n_arrays(op, 1 + 1)
        assert op._profiler._sections['section1'].sops == 57

    def test_tti_sa_like(self):
        grid = Grid(shape=(31, 17, 25))

        eqns = tti_sa_eqns(grid)

        op = Operator(eqns, subs=grid.spacing_map, opt=opt1())

        # Check generated code
        get_n_arrays(op, 4 + 6)
        assert op._profiler._sections['section1'].sops == 190

    @pytest.mark.parametrize('cire_mode', [
        'x:center,automatic',
        'automatic:3,automatic:3'
    ])
    def test_automatic_pass_unsupported(self, cire_mode):
        grid = Grid(shape=(8, 8, 8))
        u = TimeFunction(name='u', grid=grid, space_order=4)
        with pytest.raises(InvalidOperator, match='automatic'):
            Operator(Eq(u.forward, u + u.dx), opt=opt1({'cire-mode': cire_mode}))

    def test_ordered_passes(self):
        grid = Grid(shape=(8, 8, 8))
        u = TimeFunction(name='u', grid=grid, space_order=4)
        v = TimeFunction(name='v', grid=grid, space_order=4)
        op = Operator(Eq(u.forward, u + v.dx.dy.dz), name='op',
                      language=configuration['language'],
                      opt=opt1({'cire-mode': 'x:center,y:center'}))
        op.cfunction
        get_n_arrays(op, 2)


@skipif('nocudalike')
class TestDerivSelector:

    @pytest.mark.parametrize('cire_mode,expected', [
        ('x', 1),
        ('y', 1),
        ('z', 1),
        ('x2', 1),
        ('x|x2', 2),
        ('x|y2|z2|z', 4),
        ('x2:center', 1),
        ('x1:center', 1),
        ('x2:stag', 0),
        ('x2:interp', 0),
        ('x|y|z:interp', 0),
        ('automatic:3', 2),
        ('automatic:3,automatic:3', 2),
        ('automatic:200,automatic:200', 0)
    ])
    def test_basic(self, cire_mode, expected):
        grid = Grid(shape=(10, 10, 10))

        f = Function(name='f', grid=grid, space_order=4)
        g = f.func(name='g')
        u = TimeFunction(name='u', grid=grid, space_order=4)

        expr = g*u.dx2 + g*u.dy2 + g*u.dz2 + f*u.dx + f*u.dy + f*u.dz
        eqns = Eq(u.forward, expr + 1.)

        op = Operator(eqns, opt=opt1({'cire-mode': cire_mode}))

        op.cfunction

        # Check generated code
        get_n_arrays(op.body, expected)
        assert len(FindNodes(KernelLaunch).visit(op)) == (2 if expected else 1)

    @pytest.mark.parametrize('cire_mode', ['xfoo', 'xa', 'a'])
    def test_borked_value(self, cire_mode):
        grid = Grid(shape=(10, 10, 10))

        u = TimeFunction(name='u', grid=grid, space_order=4)

        eqns = Eq(u.forward, u.laplace + 1.)

        with pytest.raises(InvalidOperator):
            Operator(eqns, opt=opt1({'cire-mode': cire_mode}))

    @pytest.mark.parametrize('cire_mode', [
        'x|y|a', 'a|x', '(a):2', '1pass,1pass', '1pass,x:stag', 'x:stag,1pass'
    ])
    def test_borked_selector(self, cire_mode):
        grid = Grid(shape=(10, 10, 10))

        u = TimeFunction(name='u', grid=grid, space_order=4)

        eqns = Eq(u.forward, u.laplace + 1.)

        with pytest.raises(ValueError):
            Operator(eqns, opt=opt1({'cire-mode': cire_mode}))

    @pytest.mark.parametrize('cire_mode,expected', [
        # The eqn has `u.dx.dy` but flip() brings the `.dy` on the
        # inside, hence we expect 1 temporary array anyway
        ('y', 1),
        # The u.dx part is expected to be hoisted
        ('x', 1),
        # Try catching the whole .dy.dx mixed derivative
        ('xy', 1),
        # Mixed derivative are NOT order sensitive
        ('yx', 1),
        ('xy:center', 1),
        ('yx:center', 1),
        ('xy:stag', 0),
        ('xy:interp', 0),
    ])
    def test_mixed(self, cire_mode, expected):
        grid = Grid(shape=(10, 10, 10))

        u = TimeFunction(name='u', grid=grid, space_order=4)

        eqns = Eq(u.forward, u.dz2 + u.dx.dy + 1.)

        op = Operator(eqns, opt=opt1({'cire-mode': cire_mode}))

        op.cfunction

        # Check generated code
        get_n_arrays(op, expected)
        assert len(FindNodes(KernelLaunch).visit(op)) == (2 if expected else 1)

    @pytest.mark.parametrize('cire_mode,exp_kernels,exp_arrays', [
        (('z2', 'x'), 4, 3),
        (('z2', 'automatic'), 3, 1)
    ])
    def test_multiple_values(self, cire_mode, exp_kernels, exp_arrays):
        grid = Grid(shape=(10, 10, 10))

        u = TimeFunction(name='u', grid=grid, space_order=4)
        v = u.func(name='v')

        eqns = [Eq(u.forward, u.dz2 + u.dx.dy + 1.),
                Eq(v.forward, u.forward.dx + v.dy.dz + v.dx.dz)]

        op = Operator(eqns, opt=opt1({'cire-mode': cire_mode}))

        op.cfunction

        # Check generated code
        get_n_arrays(op.body, exp_arrays)
        assert len(FindNodes(KernelLaunch).visit(op)) == exp_kernels

    def test_product(self):
        grid = Grid(shape=(10, 10, 10))

        u = TimeFunction(name='u', grid=grid, space_order=4)
        v = u.func(name='v')

        eqns = [Eq(u.forward, u.dz2 + u.dx*v.dx + 1.)]

        op = Operator(eqns, opt=opt1({'cire-mode': 'x'}))

        op.cfunction

        # Check generated code
        get_n_arrays(op.body, 2)
        assert len(FindNodes(KernelLaunch).visit(op)) == 2

    @pytest.mark.parametrize('cire_mode,selected', [
        ('x:stag', (1,)),
        ('x:center', (0,)),
        ('x:interp', (2, 5)),
        ('x', (0, 1, 2, 3, 4, 5))
    ])
    def test_kinds(self, cire_mode, selected):
        grid = Grid(shape=(31, 17, 25))
        x, _, _ = grid.dimensions
        fields = [TimeFunction(name=f'f{i}', grid=grid, space_order=16)
                  for i in range(6)]
        out = TimeFunction(name='out', grid=grid, space_order=16)
        out1 = out.func(name='out1')

        exprs = [fields[0].dx,
                 fields[1].dx(x0={x: x + x.spacing/2}),
                 fields[2].diff(x, deriv_order=0, fd_order=16,
                                x0={x: x + x.spacing/2}),
                 fields[3].dx(weights=(-2, 1, 1)),
                 fields[4].dx(x0={x: x + x.spacing/4}),
                 fields[5].diff(x, deriv_order=0, fd_order=8,
                                x0={x: x + x.spacing/2})]

        eqn = Eq(out.forward, out + sum(exprs))
        op0 = Operator(eqn, name='op0', opt=opt0())
        op1 = Operator(eqn, name='op1', opt=opt1({'cire-mode': cire_mode}))

        producer, _ = FindNodes(KernelLaunch).visit(op1)
        assert tuple(i for i, f in enumerate(fields)
                     if f.dmap in producer.arguments) == selected

        for f in fields:
            init_chessboard(f)

        op0.apply(time_M=2)
        op1.apply(time_M=2, out=out1)

        assert np.allclose(out.data, out1.data, rtol=1.e-4, atol=1.e-5)

    @pytest.mark.parametrize('nested', [False, True])
    @pytest.mark.parametrize('cire_mode,noutputs,expected', [
        ('x', 1, 1),
        ('x|y', 1, 1),
        ('x|y', 3, 2),
        ('x:center', 1, 1),
        ('x|y:center', 1, 1),
        ('x|y:center', 3, 2),
        ('x|y:stag', 1, 0),
        ('x|y:interp', 1, 0)
    ])
    def test_sums(self, cire_mode, noutputs, expected, nested):
        grid = Grid(shape=(31, 17, 25))
        _, _, z = grid.dimensions

        u = TimeFunction(name='u', grid=grid, space_order=8)
        v = TimeFunction(name='v', grid=grid, space_order=8)
        c = Function(name='c', grid=grid)
        outputs = [TimeFunction(name=f'f{i}', grid=grid, space_order=8)
                   for i in range(noutputs)]
        eqns = []
        for i, f in enumerate(outputs, 1):
            expr = u.dx + i*v.dy
            if nested:
                expr = expr.diff(z, deriv_order=0, fd_order=8,
                                 x0={z: z + z.spacing/2})
            else:
                expr = c*expr
            eqns.append(Eq(f.forward, f + expr))

        op = Operator(eqns, opt=opt1({'cire-mode': cire_mode}))

        # Reusing two derivatives across three different sums is cheaper than
        # materializing all three sums. A partial selector must still be honored
        get_n_arrays(op.body, expected)
        if expected:
            producer = FindNodes(KernelLaunch).visit(op)[0]
            assert (v.dmap in producer.arguments) == (
                cire_mode in ('x|y', 'x|y:center'))

    @pytest.mark.parametrize('cire_mode,exp_arrays,exp_consumer', [
        (('1pass', '1pass'), 0, 0),
        (('x|y|z,x|y|z', '1pass'), 2, 2),
        (('1pass', 'x|y|z,x|y|z'), 2, 0),
        (('x|y|z', 'x|y|z,x|y|z'), 3, 1),
        (('x|y|z,x|y|z', 'x|y|z'), 3, 2),
        (('x|y|z,x|y|z',), 4, 2),
    ])
    def test_passes_multiple_values(self, cire_mode, exp_arrays, exp_consumer):
        grid = Grid(shape=(31, 17, 25))
        _, y, z = grid.dimensions

        u, v, w = [TimeFunction(name=n, grid=grid, space_order=8) for n in 'uvw']
        eqns = []
        for f, e in ((u, v.dx + w.dy), (v, u.forward.dx + w.dy)):
            for d in (y, z):
                e = e.diff(d, deriv_order=0, fd_order=8,
                           x0={d: d + d.spacing/2})
            eqns.append(Eq(f.forward, f + e))

        op = Operator(eqns, opt=opt1({'cire-mode': cire_mode}))

        op.cfunction

        get_n_arrays(op.body, exp_arrays)
        launches = FindNodes(KernelLaunch).visit(op)
        assert all(u.dmap not in i.arguments for i in launches[:exp_consumer])
        assert u.dmap in launches[exp_consumer].arguments

    @pytest.mark.parametrize('dims,cire_mode,schedule,exp_arrays', [
        ('z', 'x|y|z', None, 6),
        ('z', 'x|y|z', 0, 9),
        ('xy', 'x|y|z', None, 6),
        ('xy', 'x|y|z,x|y|z', None, 9),
        ('xy', 'x|y|z,x|y|z,x|y|z', None, 12),
        ('xy', 'x|y|z:stag,x|y|z:center,x|y|z:interp', None, 9),
        ('xy', 'x|y|z:center,x|y|z:interp,x|y|z:interp', None, 12)
    ])
    def test_strain_like(self, dims, cire_mode, schedule, exp_arrays):
        grid = Grid(shape=(31, 17, 25))
        dimensions = [d for d in grid.dimensions if d.name in dims]

        v_x, v_y, v_z = [TimeFunction(name=f'v_{d}', grid=grid, space_order=8)
                         for d in grid.dimensions]
        tau = [TimeFunction(name=f'tau{i}', grid=grid, space_order=8)
               for i in range(6)]
        tau1 = [f.func(name=f'{f.name}1') for f in tau]

        strain = (v_x.dx, v_y.dy, v_z.dz, (v_x.dy + v_y.dx)/2,
                  (v_x.dz + v_z.dx)/2, (v_y.dz + v_z.dy)/2)
        eqns = []
        for i, (f, e) in enumerate(zip(tau, strain)):
            if i >= 3:
                for d in dimensions:
                    e = e.diff(d, deriv_order=0, fd_order=8,
                               x0={d: d + d.spacing/2})
            eqns.append(Eq(f.forward, f + e))

        op0 = Operator(eqns, name='op0', opt=opt0())
        op1 = Operator(eqns, name='op1', opt=opt1({'cire-mode': cire_mode,
                                                  'cire-schedule': schedule}))

        # The cost model favors six combined components over nine individual
        # derivatives, which remain available as the first schedule. Further
        # passes add three intermediates per interpolation stage, with nothing
        # left to extract once the consumer is purely pointwise
        get_n_arrays(op1.body, exp_arrays)

        for f in (v_x, v_y, v_z):
            init_chessboard(f)

        op0.apply(time_M=2)
        op1.apply(time_M=2, **{f.name: f1 for f, f1 in zip(tau, tau1)})

        for f, f1 in zip(tau, tau1):
            assert np.allclose(f.data, f1.data, rtol=1.e-4, atol=1.e-5)


@skipif('sycl', reason="Most of these tests produce compilation failures on SYCL")
class TestDerivSchedule:

    @skipif('device')
    @pytest.mark.parametrize('name,value', [
        ('deriv-uniformize', (False, True)),
        ('deriv-schedule', ('basic', 'smart')),
        ('deriv-unroll', (False, 'full')),
    ])
    def test_local_options_unsupported(self, name, value):
        grid = Grid(shape=(11, 11))

        u = TimeFunction(name='u', grid=grid, space_order=4)

        with pytest.raises(InvalidOperator, match="Per-Kernel derivative options"):
            Operator(Eq(u.forward, u.dx + 1), opt=opt1({name: value}))

    @pytest.mark.parametrize('unroll', [False, 'full'])
    def test_smart_v0(self, unroll):
        grid = Grid(shape=(11, 11))

        u = TimeFunction(name='u', grid=grid, space_order=4)
        u1 = TimeFunction(name='u', grid=grid, space_order=4)

        eqn = Eq(u.forward, u.dx2 + u.dx + 1.)

        op0 = Operator(eqn, opt=opt0())
        op1 = Operator(eqn, opt=opt1({'deriv-schedule': 'smart',
                                      'deriv-unroll': unroll}))

        # Check generated code
        weights = get_weights(op1)
        assert len(weights) == 2
        assert all(len(i.weights) == 3 for i in weights)

        op0.apply(time_M=5)
        op1.apply(time_M=5, u=u1)

        assert np.allclose(u.data, u1.data, rtol=1e-4)

    @pytest.mark.parametrize('unroll', [False, 'inner', 'full'])
    def test_smart_v1(self, unroll):
        grid = Grid(shape=(11, 11, 11))

        u = TimeFunction(name='u', grid=grid, space_order=4)
        u1 = TimeFunction(name='u', grid=grid, space_order=4)

        eqn = Eq(u.forward, u.dy.dz + 1.)

        op0 = Operator(eqn, opt=opt0())
        op1 = Operator(eqn, opt=opt1({'deriv-schedule': 'smart',
                                      'deriv-unroll': unroll}))

        # Check generated code
        weights = get_weights(op1)
        assert len(weights) == 2
        assert all(len(i.weights) == 3 for i in weights)

        op0.apply(time_M=5)
        op1.apply(time_M=5, u=u1)

        assert np.allclose(u.data, u1.data, rtol=1e-3)

    @pytest.mark.parametrize('unroll', [False, 'full'])
    @pytest.mark.parametrize('expr,exp_n,exp_len', [
        ('u.dy.dx + u.dy.dz', 3, 5),
        ('u.dx2 + u.dy2 + u.dz2 + u.dy.dx + u.dy.dz + f*u.dx.dz', 6, 5),
        ('u.dx.dy + u.dy.dz + u.dx.dz', 3, 5),
    ])
    def test_complete(self, unroll, expr, exp_n, exp_len):
        grid = Grid(shape=(31, 17, 25))
        nt = 5

        f = Function(name='f', grid=grid, space_order=8)
        u = TimeFunction(name='u', grid=grid, space_order=8, save=nt,
                         layers=NoLayers)
        u1 = TimeFunction(name='u', grid=grid, space_order=8, save=nt,
                          layers=NoLayers)

        f.data[:] = .01

        eqn = Eq(u.forward, eval(expr) + 1.)

        op0 = Operator(eqn, opt=opt0({'gpu-fit': [u]}))
        op1 = Operator(eqn, opt=opt1({'deriv-schedule': 'smart',
                                      'deriv-unroll': unroll,
                                      'gpu-fit': [u]}))

        # Check generated code
        weights = get_weights(op1)
        assert len(weights) == exp_n
        assert all(len(i.weights) == exp_len for i in weights)

        init_chessboard(u)
        init_chessboard(u1)

        op0.apply(time_M=nt-2)
        op1.apply(time_M=nt-2, u=u1)

        assert np.allclose(u.data, u1.data, rtol=1e-2)

    @pytest.mark.parametrize('unroll,expected', [(False, 2), ('full', 0)])
    def test_shifted_derivs(self, unroll, expected):
        grid = Grid(shape=(31, 17, 25))
        x, y, z = grid.dimensions

        p0 = TimeFunction(name='p', grid=grid, time_order=2, space_order=4)
        p1 = TimeFunction(name='p', grid=grid, time_order=2, space_order=4)

        eqn = Eq(p0.forward, p0.dx(x0=x+x.spacing/2) + p0.dy(x0=y+y.spacing/2) + 1.)

        op0 = Operator(eqn, subs=grid.spacing_map, opt=opt0())
        op1 = Operator(eqn, subs=grid.spacing_map,
                       opt=opt1({'deriv-schedule': 'smart',
                                 'deriv-unroll': unroll}))

        # Check generated code
        weights = get_weights(op1)
        assert len(weights) == expected
        assert all(len(i.weights) == 2 for i in weights)

        init_chessboard(p0)
        init_chessboard(p1)

        op0.apply(time_M=5)
        op1.apply(time_M=5, p=p1)

        assert np.allclose(p0.data, p1.data, rtol=1e-4)

    @pytest.mark.parametrize('unroll,expected', [
        (False, 1), ('inner', 0), ('full', 0)
    ])
    def test_shifted_derivs_mixed_v0(self, unroll, expected):
        grid = Grid(shape=(11, 11))
        x, y = grid.dimensions

        p0 = TimeFunction(name='p', grid=grid, time_order=2, space_order=4)
        p1 = TimeFunction(name='p', grid=grid, time_order=2, space_order=4)

        eqn = Eq(p0.forward, p0.dx(x0=x+x.spacing/2).dx(x0=x-x.spacing/2) + 1.)

        op0 = Operator(eqn, subs=grid.spacing_map, opt=opt0())
        op1 = Operator(eqn, subs=grid.spacing_map,
                       opt=opt1({'deriv-schedule': 'smart',
                                 'deriv-unroll': unroll}))

        # Check generated code
        weights = get_weights(op1)
        assert len(weights) == expected
        assert all(len(i.weights) == 4 for i in weights)

        init_chessboard(p0)
        init_chessboard(p1)

        op0.apply(time_M=5)
        op1.apply(time_M=5, p=p1)

        assert np.allclose(p0.data, p1.data, rtol=1e-4)

    @skipif('nocudalike')
    @pytest.mark.parametrize('unroll,expected', [
        (False, 2), ('inner', 1), ('full', 0)
    ])
    def test_shifted_derivs_mixed_v1(self, unroll, expected):
        grid = Grid(shape=(31, 17, 25))
        x, y, z = grid.dimensions

        p0 = TimeFunction(name='p', grid=grid, time_order=2, space_order=4)
        p1 = TimeFunction(name='p', grid=grid, time_order=2, space_order=4)

        pfirst = (p0.dx(x0=x+x.spacing/2) + p0.dy(x0=y+y.spacing/2))
        psecond = pfirst.dx(x0=x-x.spacing/2) + pfirst.dy(x0=y-y.spacing/2)

        eqn = Eq(p0.forward, psecond + 1.)

        op0 = Operator(eqn, subs=grid.spacing_map, opt=opt0())
        op1 = Operator(eqn, subs=grid.spacing_map,
                       opt=opt1({'deriv-schedule': 'smart',
                                 'deriv-unroll': unroll,
                                 'cire-mode': '1pass'}))

        # Check generated code
        weights = get_weights(op1)
        assert len(weights) == expected
        assert all(len(i.weights) == 2 for i in weights)

        init_chessboard(p0)
        init_chessboard(p1)

        op0.apply(time_M=5)
        op1.apply(time_M=5, p=p1)

        assert np.allclose(p0.data, p1.data, rtol=1e-4)

    def test_custom_coeffs(self):
        p = SpaceDimension('p')
        nstc = 8

        grid = Grid(shape=(10, 10, 10))
        x, _, _ = grid.dimensions

        f = Function(name='f', grid=grid, space_order=(2*nstc, 0, 0),
                     coefficients='symbolic')
        g = Function(name='g', grid=grid, space_order=(2*nstc, 0, 0))

        ax = Function(name='DD2x', space_order=0, shape=(2*nstc + 1,),
                      dimensions=(p,))
        stencil_coeffs_x_p1 = Coefficient(1, f, x, ax)
        stencil_coeffs = Substitutions(stencil_coeffs_x_p1)

        eqn = Eq(g, f.dx, coefficients=stencil_coeffs)

        op = Operator(eqn, opt=('advanced', {'deriv-unroll': 'full'}))

        # The coefficient array would be made of Indexeds `DD2x[0], DD2x[1], ...`,
        # so we expect them to be rather inlined into the generated code to avoid
        # useless memory traffic
        assert 'init_global0' not in op._func_table


def tti_sa_eqns(grid):
    # TODO:
    # This is currently a copy-paste from submodules/devito/tests/test_unexpansion
    # We'll be able to get rid of it once conftest.py will have been split into
    # a true conftest.py file (only py.test related things) and a testing utilities
    # file
    t = grid.stepping_dim
    x, y, z = grid.dimensions

    so = 4

    a = Function(name='a', grid=grid, space_order=so)
    f = Function(name='f', grid=grid, space_order=so)
    e = Function(name='e', grid=grid, space_order=so)
    r = Function(name='r', grid=grid, space_order=so)
    p0 = TimeFunction(name='p0', grid=grid, time_order=2, space_order=so)
    m0 = TimeFunction(name='m0', grid=grid, time_order=2, space_order=so)

    def g1(field, r, e):
        return (cos(e) * cos(r) * field.dx(x0=x+x.spacing/2) +
                cos(e) * sin(r) * field.dy(x0=y+y.spacing/2) -
                sin(e) * field.dz(x0=z+z.spacing/2))

    def g2(field, r, e):
        return - (sin(r) * field.dx(x0=x+x.spacing/2) -
                  cos(r) * field.dy(x0=y+y.spacing/2))

    def g3(field, r, e):
        return (sin(e) * cos(r) * field.dx(x0=x+x.spacing/2) +
                sin(e) * sin(r) * field.dy(x0=y+y.spacing/2) +
                cos(e) * field.dz(x0=z+z.spacing/2))

    def g1_tilde(field, r, e):
        return ((cos(e) * cos(r) * field).dx(x0=x-x.spacing/2) +
                (cos(e) * sin(r) * field).dy(x0=y-y.spacing/2) -
                (sin(e) * field).dz(x0=z-z.spacing/2))

    def g2_tilde(field, r, e):
        return - ((sin(r) * field).dx(x0=x-x.spacing/2) -
                  (cos(r) * field).dy(x0=y-y.spacing/2))

    def g3_tilde(field, r, e):
        return ((sin(e) * cos(r) * field).dx(x0=x-x.spacing/2) +
                (sin(e) * sin(r) * field).dy(x0=y-y.spacing/2) +
                (cos(e) * field).dz(x0=z-z.spacing/2))

    update_p = t.spacing**2 * a**2 / f * \
        (g1_tilde(f * g1(p0, r, e), r, e) +
         g2_tilde(f * g2(p0, r, e), r, e) +
         g3_tilde(f * g3(p0, r, e) + f * g3(m0, r, e), r, e)) + \
        (2 - t.spacing * a)

    update_m = t.spacing**2 * a**2 / f * \
        (g1_tilde(f * g1(m0, r, e), r, e) +
         g2_tilde(f * g2(m0, r, e), r, e) +
         g3_tilde(f * g3(m0, r, e) + f * g3(p0, r, e), r, e)) + \
        (2 - t.spacing * a)

    eqns = [Eq(p0.forward, update_p),
            Eq(m0.forward, update_m)]

    return eqns


class TestMisc:

    @skipif('nocudalike')
    def test_pickling_preserves_weights(self):
        grid = Grid(shape=(10, 10, 10))
        x, _, _ = grid.dimensions

        u = TimeFunction(name='u', grid=grid, space_order=2)

        eqn = Eq(u.forward, u.dx2 + 1.)

        op = Operator(eqn, opt=('advanced', {'expand': False}))

        op_pkl = pickle.dumps(op)
        op_new = pickle.loads(op_pkl)

        # Check entries are 1/Muls, not Pows!
        init_global = op_new._func_table['init_global0'].root
        init_value = init_global.body.body[0].function.initvalue
        assert str(list(init_value)) == '[1/(h_x*h_x), -2.0/(h_x*h_x), 1/(h_x*h_x)]'