%matplotlib inline

# import plotting tools
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib as mpl

# import numpy
import numpy as np

# import necessary landlab components
from landlab import RasterModelGrid, HexModelGrid
from landlab.components import (FlowDirectorD8, FlowDirectorDINF,
                                FlowDirectorMFD, FlowDirectorSteepest)

# import landlab plotting functionality
from landlab.plot.drainage_plot import drainage_plot


# create a plotting routine to make a 3d plot of our surface.
def surf_plot(mg,
              surface='topographic__elevation',
              title='Surface plot of topography'):

    fig = plt.figure()
    ax = fig.gca(projection='3d')

    # Plot the surface.
    Z = mg.at_node[surface].reshape(mg.shape)
    color = cm.gray((Z - Z.min()) / (Z.max() - Z.min()))
    surf = ax.plot_surface(mg.x_of_node.reshape(mg.shape),
                           mg.y_of_node.reshape(mg.shape),
                           Z,
                           rstride=1,
                           cstride=1,
                           facecolors=color,
                           linewidth=0.,
                           antialiased=False)
    ax.view_init(elev=35, azim=-120)
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Elevation')
    plt.title(title)
    plt.show()

"""
Explanation: <a href="http://landlab.github.io"><img style="float: left" src="../../landlab_header.png"></a>
Introduction to FlowDirectors
Landlab directs flow and accumulates it using two types of components:
FlowDirectors use the topography to determine how flow moves between adjacent nodes. For every node in the grid it determines the nodes to receive flow and the proportion of flow to send from one node to its receiver. 
The FlowAccumulator uses the direction and proportion of flow moving between each node and (optionally) water runoff to calculate drainage area and discharge. 
The FlowDirectors are method-specific, and presently landlab supports four different methods for determining flow direction. 


FlowDirectorSteepest Flow is routed to only one node. The algorithm considers the link slopes leaving from each node and chooses the steepest downhill link to route flow along. In the case of a raster grid, only the links are considered (Landlab differentiates between links, which never cross and are located at North, South, East, and West on a raster grid, and diagonals, which cross and are located at North East, North West, South East, and South West). For raster grids, this method is also known as D4 flow routing. In the case of irregular grids, all links originating from a node are consideded.


FlowDirectorD8 (raster only) Flow is only routed to one node but diagonals are also considered. 


FlowDirectorMFD Flow is directed to all nodes that are located downhill of the source node. In the case of raster grid diagonals can be included using the keyword diagonals=True. Flow is partitioned between receiver nodes based on the relative slope along the links leading to the receiver nodes. The default method for partitioning is based on the sum of receiver slopes (partition_method='slope'). Partitioning can also be done on the basis of the square root of slope which gives the result of a steady kinematic wave (partition_method='square_root_of_slope').


FlowDirectorDINF (raster only) Flow is directed to two cells based on the slope of the triangular facets that can be defined between a node and its neighbors. The steepest downhill facet is chosen and then flow is partitioned between the receiver nodes at the bottom of that facet based on the relative slopes along the facet-bounding links. (The method, known as "D-infinity", is described by Tarboton (1997, Water Resources Research, 33(2), 309-319)).


In this tutorial we will go over how to initialize and run a FlowDirector, and briefly compare the different flow direction algorithms. For information on how to initialize and run the FlowAccumulator or more detailed examples that contrast the differences between each flow direction algorithm, refer to the other tutorials in this section. 
First, we import the necessary Python modules and make a small plotting routine.
End of explanation
"""


mg1 = RasterModelGrid((10, 10))
_ = mg1.add_field('topographic__elevation', mg1.y_of_node, at='node')
surf_plot(mg1, title='Grid 1: A basic ramp')

"""
Explanation: Topographic grid
For this tutorial we will consider one topographic surface. A separate tutorial exists that compares multiple grids and highlights the difference in flow direction and accumulation between the different methods.
End of explanation
"""


fd = FlowDirectorSteepest(mg1, 'topographic__elevation')

"""
Explanation: Initalizing and running a FlowDirector
To instantiate a flow director, you must pass it the minimum of a model grid that has a field called 'topographic__elevation'. Alternatively, you can pass it the name of another field name at node, or an array with length number of nodes.
End of explanation
"""


fd.run_one_step()

receivers = fd.direct_flow()

print(receivers)

"""
Explanation: The FlowDirectors are method specific, so if you want to do D8 flow directing, you must use FlowDirectorD8.
All of the FlowDirectors have two public methods: run_one_step() and direct_flow().
Both use the values of the surface provided to identify flow directions (and in the case of directing to more than one receiver, proportions). Both store the same information about receivers, proportions, and other calculated values to the model grid as fields. The difference is that run_one_step() does not return any values while direct_flow() returns the receiver array in FlowDirectorSteepest and FlowDirectorD8, and the receiver and proportion arrays in FlowDirectorMFD and FlowDirectorDINF.
End of explanation
"""


plt.figure()
drainage_plot(mg1, title='Basic Ramp using FlowDirectorSteepest')

"""
Explanation: We can illustrate the receiver node FlowDirectionSteepest has assigned to each donor node using a plotting function in Landlab called drainage_plot. We will see many of these plots in this tutorial so let's take a moment to walk through the plot and what it contains. 


The background image (white to black) shows the values of topographic elevation of the underlying surface or any other at_node field we choose to plot.


The colors of the dots inside of each pixel show the locations of the nodes and the type of node. 


The arrows show the direction of flow, and the color shows the proportion of flow that travels along that link. 


The X on top of a node shows indicates that the node is a local sink and flows to itself. 


Note that in Landlab Boundary Nodes, or nodes that are on the edge of a grid, do not have area and do not contribute flow to nodes. These nodes can either be Fixed Gradient Nodes, Fixed Value Nodes, or Closed Nodes. With the exception of Closed Nodes the boundary nodes can receive flow. 
An important step in all flow direction and accumulation is setting the proper boundary condition. Refer to the boundary condition tutorials for more information.
End of explanation
"""


mg1 = RasterModelGrid((10, 10))
_ = mg1.add_field('topographic__elevation', mg1.y_of_node, at='node')
fd = FlowDirectorD8(mg1)
fd.run_one_step()
plt.figure()
drainage_plot(mg1, title='Basic Ramp using FlowDirectorD8')

"""
Explanation: Reassuringly we can see that the flow is being sent from high elevations at the top of the grid to low elevations at the bottom of the grid. We can also see that all of the arrows are yellow, and thus all of the flow is traveling on these links. 
Brief Comparison between FlowDirector methods
Now let's see how the other FlowDirectors direct the flow on this simple grid. We don't need to specify the surface so long as it is the field 'topographic__elevation'.
End of explanation
"""


mg1 = RasterModelGrid((10, 10))
_ = mg1.add_field('topographic__elevation', mg1.y_of_node, at='node')
fd = FlowDirectorMFD(mg1,
                     diagonals=False)  # diagonals=False is the default option
fd.run_one_step()
plt.figure()
drainage_plot(mg1, title='Basic Ramp using FlowDirectorMFD without diagonals')

"""
Explanation: For this ramp, the steepest slope is down a link, and not a diagonal, so FlowDirectorD8 gives the same result as FlowDirectorSteepest.
End of explanation
"""


mg1 = RasterModelGrid((10, 10))
_ = mg1.add_field('topographic__elevation', mg1.y_of_node, at='node')
fd = FlowDirectorMFD(mg1, diagonals=True)
fd.run_one_step()
plt.figure()
drainage_plot(mg1, title='Basic Ramp using FlowDirectorMFD with diagonals')

"""
Explanation: Similarly, while there is more than one node below each core node, there is only one node that is connected by a link and not a diagonal. Thus FlowDirectorMFD with the keyword diagonals set to True provides the same results as FlowDirectorSteepest and FlowDirectorD8
End of explanation
"""


print(mg1.at_node['flow__receiver_proportions'][15])

"""
Explanation: When we permit flow along diagonal connections between nodes and flow to all downhill nodes, we see a difference in the directing pattern on this simple ramp. The flow is partitioned between the three downhill nodes, and there is more flow being sent to along the link as compared with the diagonals (the links are a lighter color blue than the diagonals). We can print out the proportion of flow, which is stored as a (number of nodes, number of possible receivers) array for FlowDirectorMFD. We will just print one node for clarity. The order of the receivers is in line with the standard Landlab ordering. The array will list the receiver proportions in the following order:
[link to East,
link to North,
link to West,
link to South,
diagonal to North East,
diagonal to North West,
diagonal to South West,
diagonal to South East]
End of explanation
"""


mg1 = RasterModelGrid((10, 10))
_ = mg1.add_field('topographic__elevation', mg1.y_of_node, at='node')

fd = FlowDirectorMFD(mg1,
                     diagonals=True,
                     partition_method='square_root_of_slope')
fd.run_one_step()
plt.figure()
drainage_plot(
    mg1,
    title=
    'Basic Ramp using FlowDirectorMFD with diagonals and steady kinematic wave partitioning'
)

"""
Explanation: As we can see, the flow is moving on the link to the South, the South West and South East. There is about a third of the flow moving on the links and diagonals, though more flow is partitioned to the link. This makes sense as the link is steeper than the diagonals.
End of explanation
"""


mg1 = RasterModelGrid((10, 10))
_ = mg1.add_field('topographic__elevation', mg1.y_of_node, at='node')

fd = FlowDirectorDINF(mg1)
fd.run_one_step()
plt.figure()
drainage_plot(mg1, title='Basic Ramp using FlowDirectorDINF')

"""
Explanation: In the FlowDirectorMFD we can partition flow based on slope (default) or based on the square root of the slope. This option gives the solution to a steady kinematic wave. We can see that the proportion of flow going on the links and diagonals is more similar.
End of explanation
"""


hmg = HexModelGrid((9, 5))
_ = hmg.add_field('topographic__elevation', hmg.y_of_node, at='node')
fd = FlowDirectorSteepest(hmg)
fd.run_one_step()
plt.figure()
drainage_plot(hmg)

"""
Explanation: In FlowDirectorDINF flow is partitioned to two nodes based on steepness of the eight triangular facets surrounding each node. The partitioning is based on the relation between the link and diagonal slope that form the edge of the facet and the slope of the facet itself. When one of the facet edges has the same slope as the facet, as is the case in this ramp example, all of the flow is partitioned along that edge. 
FlowDirectors on irregular grids
FlowDirectorSteepest and FlowDirectorMFD work on raster, hex, and irregular grids. FlowDirectorD8 and FlowDirectorDINF don't work on hex or irregular grids because the directing and partitioning algorithm that underlies each of them is specific to raster grids. 
Before concluding this tutorial, we will briefly consider a hexagonal example.
End of explanation
"""


hmg = HexModelGrid((9, 5))
_ = hmg.add_field('topographic__elevation', hmg.y_of_node, at='node')

fd = FlowDirectorMFD(hmg)
fd.run_one_step()
plt.figure()
drainage_plot(hmg)

hmg.at_node["topographic__elevation"][hmg.core_nodes].max()

"""
Explanation: Just as in the raster example, FlowDirectorSteepest sends flow from donor nodes to only one receiver node. You may ask why the flow is going to the southwest while the links are just as steep to the southeast. When two links have exactly the same slope, in route-to-one methods like FlowDirectorSteepest and FlowDirectorD8 landlab will choose the first steepest link it encounters while searching through the links. When diagonal links are also considered in FlowDirectorD8 the diagonals are examined after the links. Thus if a link and a diagonal leading from one node downhill had exactly the same slope, the flow would be sent down the link. 
For a more thorough discussion and comparison of each of the FlowDirector algorithms, please see the tutorial "Comparison of FlowDirectors".
End of explanation
"""


