from devito import Eq, Operator, Function, TimeFunction


def iso_stencil(field, model, **kwargs):
    """
    Stencil for the scalar isotropic visco- acoustic variable density
    self adjoint wave equation:

        b/v^2 [ P.dt2 + w/Q P.dt ] = (b P.dx).dx + (b P.dy).dy + (b P.dz).dz + s

    Note derivative shifts are omitted for simplicity above.
    See implementation notebook sa_01_iso_implementation1.ipynb for more details.

    Parameters
    ----------
    field : TimeFunction, required
        The pressure wavefield computed solution.
    model : Dictionary <string>:<Function>, contains:
        'b': Buoyancy = reciprocal density (units: m^3/kg)
        'v': Velocity (units: m/msec or km/sec)
        'wOverQ': The w/Q field for dissipation only attenuation.
    forward : bool, optional
        The propagation direction. Defaults to True.
    q : TimeFunction, Function or float, optional
        Full-space/time source of the wave-equation.

    Returns
    ----------
    The time update stencil.
    """
    # Get the Functions for buoyancy, velocity, and wOverQ
    vp, b, wOverQ = model.vp, model.b, model.damp

    # Define time step of pressure wavefield to be updated
    forward = kwargs.get('forward', True)

    if forward:
        field_next = field.forward
        field_prev = field.backward
    else:
        field_next = field.backward
        field_prev = field.forward