"""Landlab component that simulates overland flow.

This component simulates overland flow using the 2-D numerical model of
shallow-water flow over topography using the de Almeida et al., 2012
algorithm for storage-cell inundation modeling.

.. codeauthor:: Jordan Adams

Examples
--------
>>> import numpy as np
>>> from landlab import RasterModelGrid
>>> from landlab.components.overland_flow import OverlandFlow

Create a grid on which to calculate overland flow.

>>> grid = RasterModelGrid((4, 5))

The grid will need some data to provide the overland flow component. To
check the names of the fields that provide input to the overland flow
component use the *input_var_names* class property.

>>> OverlandFlow.input_var_names
('surface_water__depth', 'topographic__elevation')

Create fields of data for each of these input variables.

>>> grid.at_node['topographic__elevation'] = np.array([
...     0., 0., 0., 0., 0.,
...     1., 1., 1., 1., 1.,
...     2., 2., 2., 2., 2.,
...     3., 3., 3., 3., 3.])
>>> grid.at_node['surface_water__depth'] = np.array([
...     0. , 0. , 0. , 0. , 0. ,
...     0. , 0. , 0. , 0. , 0. ,
...     0. , 0. , 0. , 0. , 0. ,
...     0.1, 0.1, 0.1, 0.1, 0.1])

Instantiate the `OverlandFlow` component to work on this grid, and run it.

>>> of = OverlandFlow(grid, steep_slopes=True)
>>> of.run_one_step()

After calculating the overland flow, new fields have been added to the
grid. Use the *output_var_names* property to see the names of the fields that
have been changed.

>>> of.output_var_names
('surface_water__depth', 'surface_water__discharge', 'water_surface__gradient')

The `surface_water__depth` field is defined at nodes.

>>> of.var_loc('surface_water__depth')
'node'
>>> grid.at_node['surface_water__depth'] # doctest: +NORMALIZE_WHITESPACE
array([  1.00000000e-05,   1.00000000e-05,   1.00000000e-05,
         1.00000000e-05,   1.00000000e-05,   1.00000000e-05,
         1.00000000e-05,   1.00000000e-05,   1.00000000e-05,
         1.00000000e-05,   1.00000000e-05,   2.00100000e-02,
         2.00100000e-02,   2.00100000e-02,   1.00000000e-05,
         1.00010000e-01,   1.00010000e-01,   1.00010000e-01,
         1.00010000e-01,   1.00010000e-01])

The `surface_water__discharge` field is defined at links. Because our initial
topography was a dipping plane, there is no water discharge in the horizontal
direction, only toward the bottom of the grid.

>>> of.var_loc('surface_water__discharge')
'link'
>>> q = grid.at_link['surface_water__discharge'] # doctest: +NORMALIZE_WHITESPACE
>>> np.all(q[grid.horizontal_links] == 0.)
True
>>> np.all(q[grid.vertical_links] <= 0.)
True

The *water_surface__gradient* is also defined at links.

>>> of.var_loc('water_surface__gradient')
'link'
>>> grid.at_link['water_surface__gradient'] # doctest: +NORMALIZE_WHITESPACE
array([ 0. ,  0. ,  0. ,  0. ,
        0. ,  1. ,  1. ,  1. ,  0. ,
        0. ,  0. ,  0. ,  0. ,
        0. ,  1. ,  1. ,  1. ,  0. ,
        0. ,  0. ,  0. ,  0. ,
        0. ,  1.1,  1.1,  1.1,  0. ,
        0. ,  0. ,  0. ,  0. ])
"""
from landlab import Component, FieldError
import numpy as np
from landlab.grid.structured_quad import links
from landlab.utils.decorators import use_file_name_or_kwds


_SEVEN_OVER_THREE = 7.0 / 3.0


class OverlandFlow(Component):

    """Simulate overland flow using de Almeida approximations.

    Landlab component that simulates overland flow using the de Almeida
    et al., 2012 approximations of the 1D shallow water equations to be used
    for 2D flood inundation modeling.

    This component calculates discharge, depth and shear stress after some
    precipitation event across any raster grid. Default input file is named
    "overland_flow_input.txt' and is contained in the
    landlab.components.overland_flow folder.

    Parameters
    ----------
    grid : RasterModelGrid
        A landlab grid.
    h_init : float, optional
        Thicknes of initial thin layer of water to prevent divide by zero
        errors (m).
    alpha : float, optional
        Time step coeffcient, described in Bates et al., 2010 and
        de Almeida et al., 2012.
    mannings_n : float, optional
        Manning's roughness coefficient.
    g : float, optional
        Acceleration due to gravity (m/s^2).
    theta : float, optional
        Weighting factor from de Almeida et al., 2012.
    rainfall_intensity : float, optional
        Rainfall intensity.



    The primary method of this class is :func:`run_one_step`.


    Construction::

        OverlandFlow(grid, default_fixed_links=False, h_init=0.00001,
                 alpha=0.7, mannings_n=0.03, g=9.81, theta=0.8,
                 rainfall_intensity=0.0, steep_slopes=False, **kwds)

"""
    _name = 'OverlandFlow'

    _input_var_names = (
        'surface_water__depth',
        'topographic__elevation',
    )

    _output_var_names = (
        'surface_water__depth',
        'surface_water__discharge',
        'water_surface__gradient',
    )

    _var_units = {
        'surface_water__depth': 'm',
        'surface_water__discharge': 'm3/s',
        'topographic__elevation': 'm',
        'water_surface__gradient': '-',
    }

    _var_mapping = {
        'surface_water__depth': 'node',
        'topographic__elevtation': 'node',
        'surface_water__discharge': 'link',
        'water_surface__gradient': 'link',
    }

    _var_doc = {
        'surface_water__depth': 'The depth of water at each node.',
        'topographic__elevtation': 'The land surface elevation.',
        'surface_water__discharge': 'The discharge of water on active links.',
        'water_surface__gradient': 'Downstream gradient of the water surface.',
    }

    @use_file_name_or_kwds
    def __init__(self, grid, default_fixed_links=False, h_init=0.00001,
                 alpha=0.7, mannings_n=0.03, g=9.81, theta=0.8,
                 rainfall_intensity=0.0, steep_slopes=False, **kwds):
        """Create a overland flow component.

        Parameters
        ----------
        grid : RasterModelGrid
            A landlab grid.
        h_init : float, optional
            Thicknes of initial thin layer of water to prevent divide by zero
            errors (m).
        alpha : float, optional
            Time step coeffcient, described in Bates et al., 2010 and
            de Almeida et al., 2012.
        mannings_n : float, optional
            Manning's roughness coefficient.
        g : float, optional
            Acceleration due to gravity (m/s^2).
        theta : float, optional
            Weighting factor from de Almeida et al., 2012.
        rainfall_intensity : float, optional
            Rainfall intensity.
        """
        super(OverlandFlow, self).__init__(grid, **kwds)

        # First we copy our grid
        self._grid = grid

        self.h_init = h_init
        self.alpha = alpha
        
        if type(mannings_n) is str:
            self.mannings_n = self.grid.at_link[mannings_n]
        else:
            self.mannings_n = mannings_n

        self.g = g
        self.theta = theta
        self.rainfall_intensity = rainfall_intensity
        self.steep_slopes = steep_slopes


        # Now setting up fields at the links...
        # For water discharge
        try:
            self.q = grid.add_zeros('surface_water__discharge', at='link',
                units=self._var_units['surface_water__discharge'])

        except FieldError:
            # Field was already set; still, fill it with zeros
            self.q = grid.at_link['surface_water__discharge']
            self.q.fill(0.)

        # For water depths calculated at links
        try:
            self.h_links = grid.add_zeros('surface_water__depth', at='link',
                                          units=self._var_units[
                                              'surface_water__depth'])
        except FieldError:
            self.h_links = grid.at_link['surface_water__depth']
            self.h_links.fill(0.)

        self.h_links += self.h_init

        try:
            self.h = grid.add_zeros('surface_water__depth', at='node',
                units=self._var_units['surface_water__depth'])

        except FieldError:
            # Field was already set
            self.h = grid.at_node['surface_water__depth']

        self.h += self.h_init

        # For water surface slopes at links
        try:
            self.water_surface_slope = grid.add_zeros('water_surface__gradient', at='link')
        except FieldError:
            self.water_surface_slope = grid.at_link['water_surface__gradient']
            self.water_surface_slope.fill(0.)

        # Start time of simulation is at 1.0 s
        self.elapsed_time = 1.0

        self.dt = None
        self.dhdt = grid.zeros()

        # When we instantiate the class we recognize that neighbors have not
        # been found. After the user either calls self.set_up_neighbor_array
        # or self.overland_flow this will be set to True. This is done so
        # that every iteration of self.overland_flow does NOT need to
        # reinitalize the neighbors and saves computation time.
        self.neighbor_flag = False

        # When looking for neighbors, we automatically ignore inactive links
        # by default. However, what about when we want to look at fixed links
        # too? By default, we ignore these, but if they are important to your
        # model and will be updated in your driver loop, they can be used by
        # setting the flag in the initialization of  the class to 'True'
        self.default_fixed_links = default_fixed_links

        # Assiging a class variable to the elevation field.
        self.z = self._grid.at_node['topographic__elevation']