from landlab import RasterModelGrid
import numpy as np
from landlab.plot.imshow import imshow_grid
%matplotlib inline

"""
Explanation: <a href="http://landlab.github.io"><img style="float: left" src="../../landlab_header.png"></a>
Setting watershed boundary conditions on a raster grid

This tutorial ilustrates how to set watershed boundary conditions on a raster grid.
Note that a watershed is assumed to have a ring of nodes around the core nodes that are closed boundaries (i.e. no flux can cross these nodes, or more correctly, no flux can cross the faces around the nodes).
This means that automatically the nodes on the outer perimeter of the grid will be set to be closed boundary.
By definitation a watershed also has one outlet through which fluxes can pass.  Here the outlet is set as the node that has the lowest value, is not a nodata_value node, and is adjacent to at least one closed boundary node.  
This means that an outlet can be on the outer perimeter of the raster.  However, the outlet does not need to be on the outer perimeter of the raster.

The first example uses set_watershed_boundary_condition, which finds the outlet for the user.

First import what we need.
End of explanation
"""


mg1 = RasterModelGrid((5, 5), 1.)
z1 = mg1.add_ones('topographic__elevation', at='node')
mg1.at_node['topographic__elevation'][2] = 0.
mg1.at_node['topographic__elevation']

"""
Explanation: Now we create a 5 by 5 grid with a spacing (dx and dy) of 1.
We also create an elevation field with value of 1. everywhere, except at the outlet, where the elevation is 0.  In this case the outlet is in the middle of the bottom row, at location (0,2), and has a node id of 2.
End of explanation
"""


mg1.set_watershed_boundary_condition(z1)

"""
Explanation: The set_watershed_boundary_condition in RasterModelGrid will find the outlet of the watershed.
This method takes the node data, in this case z, and, optionally the no_data value.
This method sets all nodes that have no_data values to closed boundaries.
This example does not have any no_data values, which is fine.  
In this case, the code will set all of the perimeter nodes as BC_NODE_IS_CLOSED(boundary status 4) in order to create this boundary around the core nodes. 
The exception on the perimeter is node 2 (with elevation of 0).  Although it is on the perimeter, it has a value and it has the lowest value.  So in this case node 2 will be set as BC_NODE_IS_FIXED_VALUE (boundary status 1).
The rest of the nodes are set as a CORE_NODE (boundary status 0)
End of explanation
"""


imshow_grid(mg1, mg1.status_at_node, color_for_closed='blue')

"""
Explanation: Check to see that node status were set correctly.
imshow_grid will default to not plot the value of BC_NODE_IS_CLOSED nodes, which is why we override this below with the option color_for_closed
End of explanation
"""


mg2 = RasterModelGrid((5, 5), 10.)
z2 = mg2.add_ones('topographic__elevation', at='node')
mg2.at_node['topographic__elevation'][1] = 0.
mg2.at_node['topographic__elevation']

"""
Explanation: The second example uses set_watershed_boundary_condition_outlet_coords

In this case the user knows the coordinates of the outlet node.
First instantiate a new grid, with new data values.
End of explanation
"""


mg2.set_watershed_boundary_condition_outlet_coords((0, 1), z2)

"""
Explanation: Note that the node with zero elevation, which will be the outlet, is now at location (0,1).
Note that even though this grid has a dx & dy of 10., the outlet coords are still (0,1).
Set the boundary conditions.
End of explanation
"""


imshow_grid(mg2, mg2.status_at_node, color_for_closed='blue')

"""
Explanation: Plot grid of boundary status information
End of explanation
"""


mg3 = RasterModelGrid((5, 5), 5.)
z3 = mg3.add_ones('topographic__elevation', at='node')
mg3.at_node['topographic__elevation'][5] = 0.
mg3.at_node['topographic__elevation']

"""
Explanation: The third example uses set_watershed_boundary_condition_outlet_id

In this case the user knows the node id value of the outlet node.
First instantiate a new grid, with new data values.
End of explanation
"""


mg3.set_watershed_boundary_condition_outlet_id(5, z3)

"""
Explanation: Set boundary conditions with the outlet id.
Note that here we know the id of the node that has a value of zero and choose this as the outlet.  But the code will not complain if you give it an id value of a node that does not have the smallest data value.
End of explanation
"""


imshow_grid(mg3, mg3.status_at_node, color_for_closed='blue')

"""
Explanation: Another plot to illustrate the results.
End of explanation
"""


from landlab.io import read_esri_ascii
(grid_bijou, z_bijou) = read_esri_ascii('west_bijou_gully.asc', halo=1)

"""
Explanation: The final example uses set_watershed_boundary_condition on a watershed that was exported from Arc.

First import read_esri_ascii and then import the DEM data.
An optional value of halo=1 is used with read_esri_ascii.  This puts a perimeter of nodata values around the DEM.
This is done just in case there are data values on the edge of the raster.  These would have to become closed to set watershed boundary conditions, but in order to avoid that, we add a perimeter to the data.
End of explanation
"""


imshow_grid(grid_bijou, z_bijou)

"""
Explanation: Let's plot the data to see what the topography looks like.
End of explanation
"""


grid_bijou.set_watershed_boundary_condition(z_bijou, 0)

"""
Explanation: In this case the nodata value is zero.  This skews the colorbar, but we can at least see the shape of the watershed.
Let's set the boundary condition.  Remember we don't know the outlet id.
End of explanation
"""


imshow_grid(grid_bijou, grid_bijou.status_at_node, color_for_closed='blue')

"""
Explanation: Now we can look at the boundary status of the nodes to see where the found outlet was.
End of explanation
"""


imshow_grid(grid_bijou, z_bijou)

"""
Explanation: This looks sensible.
Now that the boundary conditions ae set, we can also look at the topography. 
imshow_grid will default to show boundaries as black, as illustrated below.  But that can be overwridden as we have been doing all along.
End of explanation
"""


