from landlab.components import FlowAccumulator, ErosionDeposition
from landlab.plot import imshow_grid
from landlab import HexModelGrid
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import copy

"""
Explanation: <a href="http://landlab.github.io"><img style="float: left" src="../../../landlab_header.png"></a>
Introduction to the ErosionDeposition component
<hr>
<small>For more Landlab tutorials, click here: <a href="https://landlab.readthedocs.io/en/latest/user_guide/tutorials.html">https://landlab.readthedocs.io/en/latest/user_guide/tutorials.html</a></small>
<hr>

This tutorial introduces the ErosionDeposition component, which simulates erosion and deposition in a river network over long time scales.
Theory and Overview
ErosionDeposition models fluvial processes using the approach described by Davy and Lague (2009). The basic goal is to calculate rates of erosion or deposition in the various branches of a river network that is embedded in a gridded landscape (see, e.g., Tucker and Hancock, 2010). The literature has a number of different approaches to this problem. Some models assume transport-limited conditions, such that erosion or deposition result from local imbalances in sediment transport capacity (see, e.g., Willgoose, 2018). Some assume detachment-limited conditions, such that any eroded sediment is entirely removed  (e.g., Howard, 1994; Whipple and Tucker, 1999; Tucker and Whipple, 2002). Still others allow for both erosion of detachment-resistant material (bedrock or cohesive sediments), and re-deposition of that material. The Davy-Lague approach falls in the latter category, and is similar to the approach used in some shorter-term morphodynamic models as well as agricultural soil erosion models. The basic idea involves conservation of sediment mass in the water column of a river. One calculates the rate of entrainment of sediment from the bed into the water as a function of discharge and local slope. The bed also accumulates sediment that settles out from the water column, at a rate that depends on sediment concentration (treated as the ratio of sediment flux to water discharge) and a settling-velocity parameter.
The theory behind using this kind of dual erosion-deposition approach in the context of fluvial landscape evolution is described by Davy and Lague (2009). Tucker and Hancock (2010) provide a review of landscape evolution modeling that compares this with other approaches to fluvial erosion/deposition theory. The equations used in the Landlab ErosionDeposition component are described by Barnhart et al. (2019) as part of the terrainBento collection of Landlab-based landscape evolution models, which includes models that use ErosionDeposition.
End of explanation
"""


print(ErosionDeposition.__doc__)

"""
Explanation: Look at the top-level internal documentation for the ErosionDeposition component:
End of explanation
"""


print(ErosionDeposition.__init__.__doc__)

"""
Explanation: The __init__ docstring lists the parameters for this component:
End of explanation
"""


# Parameters
nrows = 41
ncols = 41
dx = 100.0
K = 0.0001  # erodibility coefficient, 1/yr
m_sp = 0.5  # exponent on drainage area or discharge, -
n_sp = 1.0  # exponent on slope, -
sp_crit = 0.0  # erosion threshold
v_s = 100.0  # settling velocity parameter (dimensionless if drainage area is used instead of discharge)
F_f = 0.5  # fraction of fines generated during bed erosion
initial_elevation = (
    200.0  # starting elevation of an "uplifted block" (rapid baselevel drop), m
)

run_duration = 120000.0  # duration of run, yr
dt = 100.0  # time-step duration, yr
plot_every = 40000.0  # time interval for plotting, yr

# Derived parameters
nsteps = int(run_duration / dt)
next_plot = plot_every

# set up colormap
cmap = copy.copy(mpl.cm.get_cmap("pink"))

"""
Explanation: Set some parameters:
End of explanation
"""


mg = HexModelGrid(
    (nrows, ncols), spacing=dx, node_layout="rect", orientation="vertical"
)
z = mg.add_zeros("topographic__elevation", at="node")
# add some roughness, as this lets "natural" channel planforms arise
np.random.seed(0)
initial_roughness = np.random.rand(z.size)
z[:] += initial_roughness
z[mg.core_nodes] += initial_elevation
z[mg.boundary_nodes] = 0.0

# close off boundaries on 3 sides
is_closed_boundary = np.logical_and(
    mg.status_at_node != mg.BC_NODE_IS_CORE,
    mg.x_of_node < (np.amax(mg.x_of_node) - 0.5 * dx),
)
mg.status_at_node[is_closed_boundary] = mg.BC_NODE_IS_CLOSED

"""
Explanation: Create a grid with one side open:
End of explanation
"""


fr = FlowAccumulator(mg, depression_finder="DepressionFinderAndRouter")
ed = ErosionDeposition(
    mg,
    K=K,
    m_sp=m_sp,
    n_sp=n_sp,
    sp_crit=sp_crit,
    v_s=v_s,
    F_f=F_f,
    solver="adaptive",  # use the adaptive time stepper, which is slightly faster
)

"""
Explanation: Instantiate components:
End of explanation
"""


for i in range(1, nsteps + 1):

    # route flow
    fr.run_one_step()  # run_one_step isn't time sensitive, so it doesn't take dt as input

    # do some erosion/deposition
    ed.run_one_step(dt)

    if i * dt >= next_plot:
        plt.figure()
        imshow_grid(
            mg,
            "topographic__elevation",
            grid_units=["m", "m"],
            var_name="Elevation (m)",
            cmap=cmap,
        )
        next_plot += plot_every

"""
Explanation: Run the model in a loop to evolve topography on the uplifted block:
End of explanation
"""


