#! /usr/bin/env python
import numpy as np

from landlab.components import LinearDiffuser

_TINY_DIFFUSIVITY = 1.0e-20


class SimpleSubmarineDiffuser(LinearDiffuser):
    r"""
    Transport marine sediment using a water-depth-dependent diffusion model.

    This component models sediment transport as a diffusion process with a
    coefficient that depends on water depth :math:`h` as follows:

    .. math::
        D(h) = D_0 f_1(h) f_2(h)

    Here :math:`D_0` is the maximum value, corresponding to the input
    parameter :code:`shallow_water_diffusivity`.

    The function :math:`f_1(h)` describes the decrease in transport efficiency
    below the wave base depth :math:`h_w`. It is defined as unity for depth
    above the wave base, and as

    .. math::
        f_1(h) = \exp( -(h - h_w) / h_w)

    for :math:`h > h_w`.

    The function :math:`f_2(h)` handles the transition in transport efficiency
    around the shoreline. If :code:`tidal_range`, :math:`R_t`, is zero, then
    :math:`f_2` is set to unity underwater (:math:`h \ge 0`), and a tiny value
    above water (not zero, because that would cause a divide-by-zero error in
    the base class).
    If :math:`R_t > 0`, then a :math:`tanh` function is used to model
    a smooth decrease in :math:`D` from the low to high tide level:

    .. math::
        f_2(h) = (\tanh ( -h / R_t) + 1) / 2

    with an addition tiny value added to locations above water to avoid
    division by zero.

    Examples
    --------
    >>> from landlab import RasterModelGrid
    >>> from landlab.components import SimpleSubmarineDiffuser
    >>> grid = RasterModelGrid((3, 7), xy_spacing=100.0)
    >>> grid.set_closed_boundaries_at_grid_edges(False, True, False, True)
    >>> topo = grid.add_zeros('topographic__elevation', at='node')
    >>> topo[:] = -10.0
    >>> topo[9:14] = [0., 10., 10., 5., 5.]
    >>> ssd = SimpleSubmarineDiffuser(grid, tidal_range=0.0)
    >>> ssd.run_one_step(dt=5.0)
    >>> topo[8:13]
    array([ -9.5,   0. ,   9.5,  10. ,   5. ])
    >>> grid.at_node["sediment_deposit__thickness"][8:13]
    array([ 0.5,  0. , -0.5,  0. ,  0. ])
    """

    _name = "SimpleSubmarineDiffuser"

    _time_units = "y"

    _info = {
        "sea_level__elevation": {
            "dtype": "float",
            "intent": "in",
            "optional": False,
            "units": "m",
            "mapping": "grid",
            "doc": "Sea level elevation",
        },
        "topographic__elevation": {
            "dtype": "float",
            "intent": "inout",
            "optional": False,
            "units": "m",
            "mapping": "node",
            "doc": "Land surface topographic elevation",  # and seafloor
        },
        "water__depth": {
            "dtype": "float",
            "intent": "out",
            "optional": False,
            "units": "m",
            "mapping": "node",
            "doc": "depth of water under current sea level",
        },
        "sediment_deposit__thickness": {
            "dtype": "float",
            "intent": "out",