from landlab.components import LinearDiffuser
from landlab.plot import imshow_grid
from landlab import RasterModelGrid
import matplotlib as mpl
import matplotlib.cm as cm
from matplotlib.pyplot import figure, show, plot, xlabel, ylabel, title
import numpy as np

"""
Explanation: <a href="http://landlab.github.io"><img style="float: left" src="../../landlab_header.png"></a>
Getting to know the Landlab component library
<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 notebook walks you through the stages of creating and running a Landlab model using the Landlab component library.
We are going to create three models: firstly, a single-component driver implementing just linear diffusion; then a three-component driver implementing linear diffusion, flow routing, and stream power incision; and finally a similar model, but implementing a storm-interstorm sequence.
The basics: one component
Let's begin with the one-component diffusion model.
Firstly, import the library elements we'll need. The component classes can all be imported from the landlab.components library. They're all formatted in CamelCaseLikeThis. Anything you see in that folder that isn't formatted like this isn't a component!
End of explanation
"""


mg = RasterModelGrid((80, 80), xy_spacing=5.)
z = mg.add_zeros('topographic__elevation', at="node")

"""
Explanation: Let's start by creating the grid that we'll do the first part of this exercise with, and putting some data into its fields. Note that you need to create the fields that a component takes as inputs before instantiating a component - though you can put values into the arrays later if you need to (as illustrated below). For more info on working with fields, see the fields tutorial.
End of explanation
"""


LinearDiffuser.input_var_names

"""
Explanation: How did we know this was a field we needed as an input? Well, firstly because we read the component documentation (always do this!), but secondly we can get a reminder using the Landlab Component Standard Interface:
End of explanation
"""


LinearDiffuser.var_help('topographic__elevation')

"""
Explanation: Note we didn't have to instantiate the component to be able to do this! Other standard properties are output_var_names and optional_var_names; pass an input or output name to var_loc, var_type, var_units, and var_definition to get the centering ('node', 'link', etc.), array dtype (float, int), units (meters, etc.), and a descriptive string, respectively. var_help will give you a lot of this information at once:
End of explanation
"""


for edge in (mg.nodes_at_left_edge, mg.nodes_at_right_edge):
    mg.status_at_node[edge] = mg.BC_NODE_IS_CLOSED 
for edge in (mg.nodes_at_top_edge, mg.nodes_at_bottom_edge):
    mg.status_at_node[edge] = mg.BC_NODE_IS_FIXED_VALUE

"""
Explanation: It's also a good idea to set the grid boundary conditions before component instantiation. Let's have fixed value top and bottom and closed left and right (see the boundary conditions tutorial):
End of explanation
"""


lin_diffuse = LinearDiffuser(mg, linear_diffusivity=0.2)

"""
Explanation: You will find that all components within landlab share a similar interface. We'll examine how it looks first on the diffusion component.
Landlab components have a standardised instantiation signature. Inputs to the component can be fed in as arguments to the constructor (i.e., the function that gets called when you create a new instances of a component), rather than being fed in as strings from a text input file (though note, you an still do this, see below). This has two major advantages: firstly, components now have plainly declared default values, which are visible just as they would be in, say, a numpy function; secondly, because the inputs are now Python objects, it's a lot easier to work with spatially variable inputs that need to be passed in as arrays, and also to feed dynamically changing inputs into a component.
The standard signature to instantiate a component looks like this:
python
MyComponent(grid, input1=default1, input2=default2, input3=default3, ...)
Because defaults are provided, you can instantiate a component with default values very simply. The diffuser, for example, requires only that a linear_diffusity be supplied:
End of explanation
"""


total_t = 200000.
dt = 1000.
uplift_rate = 0.001
nt = int(total_t // dt)
# ^note if we didn't know a priori that there are a round number of steps dt in the
# total time, we'd have to take care to account for the "extra" time (see example below)
for i in range(nt):
    lin_diffuse.run_one_step(dt)
    z[mg.core_nodes] += uplift_rate * dt  # add the uplift
    # add some output to let us see we aren't hanging:
    if i % 50 == 0:
        print(i*dt)

"""
Explanation: We'll see some other ways of initializing (e.g., from an input file) below.
Now we're ready to run the component! Run methods are also standardized. Most Landlab components have a standard run method named run_one_step, and it looks like this:
python
my_comp.run_one_step(dt)
If the component is time-dependent, dt, the timestep, will be the first argument. (In Landlab 1.x, some components have subsequent keywords, which will typically be flags that control the way the component runs, and usually can be left as their default values; these extra keywords are absent in Landlab 2.x). Note that nothing is returned from a run method like this, but that nonetheless the grid fields are updated.
This dt is properly thought of as the external model timestep; it controls essentially the frequency at which the various Landlab components you're implementing can exchange information with each other and with the driver (e.g., frequency at which uplift steps are added to the grid). If your model has a stability condition that demands a shorter timestep, the external timestep will be subdivided internally down to this shorter timescale.
So let's do it. It's up to you as the component designer to make sure your driver script accounts properly for the total time the model runs. Here, we want to run for 200000 years with a timestep of 1000 years, with an uplift rate of 0.001 m/y. So:
End of explanation
"""


# the following line makes figures show up correctly in this document (only needed for Jupyter notebook)
%matplotlib inline

# Create a figure and plot the elevations
figure(1)
im = imshow_grid(mg, 'topographic__elevation', grid_units = ['m','m'],
                 var_name='Elevation (m)')

figure(2)
elev_rast = mg.node_vector_to_raster(z)
ycoord_rast = mg.node_vector_to_raster(mg.node_y)
ncols = mg.number_of_node_columns
im = plot(ycoord_rast[:, int(ncols // 2)], elev_rast[:, int(ncols // 2)])
xlabel('horizontal distance (m)')
ylabel('vertical distance (m)')
title('topographic__elevation cross section')

"""
Explanation: Note that we're using z to input the uplift here, which we already bound to the Landlab field mg.at_node['topographic__elevation] when we instantiated that field. This works great, but always be careful to update the values inside the array, not to reset the variable as equal to something else, i.e., to put new values in the field do::
python
z[:] = new_values  # values copied into the existing field
not
python
z = new_values  # z is now "new_values", not the field!
Now plot the output!
End of explanation
"""


z[:] = 0.  # reset the elevations to zero
k_diff = mg.zeros('node', dtype=float)
k_diff.fill(1.)
k_diff *= (mg.node_x.max() - 0.9 * mg.x_of_node) / mg.x_of_node.max()
k_field = mg.add_field('linear_diffusivity', k_diff, at='node', clobber=True)
imshow_grid(mg, k_diff, var_name='k_diff', cmap='winter')  # check it looks good

"""
Explanation: Now, let's repeat this exercise, but illustrating the way we can input fields as some parameters for components. We're going to make the diffusivity spatially variable, falling by a factor of ten as we move across the grid.
End of explanation
"""


lin_diffuse = LinearDiffuser(mg, linear_diffusivity='linear_diffusivity')
# we could also have passed in `k_diff` in place of the string

"""
Explanation: Now we re-initialize the component instance to bind the k_diff field to the component:
End of explanation
"""


for i in range(nt):
    lin_diffuse.run_one_step(dt)
    z[mg.core_nodes] += uplift_rate * dt  # add the uplift
    # add some output to let us see we aren't hanging:
    if i % 50 == 0:
        print(i*dt)
figure(3)
im = imshow_grid(mg, 'topographic__elevation', grid_units = ['m','m'],
                 var_name='Elevation (m)')

"""
Explanation: ...and run just as before. Note this will be slower than before; the internal timestep is shorter because we've modified the diffusivities.
End of explanation
"""


from landlab.components import FlowAccumulator, FastscapeEroder
from landlab import load_params

"""
Explanation: Running two or more components
Now we're going to take a similar approach but this time combine the outputs of three distinct Landlab components: the diffuser, the monodirectional flow router, and the stream power incisor. For clarity, we're going to repeat the whole process from the start.
So first, let's import everything we don't already have:
End of explanation
"""


input_file = './coupled_params.txt'
inputs = load_params(input_file) # load the data into a dictionary

nrows = inputs['nrows']
ncols = inputs['ncols']
dx = inputs['dx']
uplift_rate = inputs['uplift_rate']
total_t = inputs['total_time']
dt = inputs['dt']

nt = int(total_t // dt) #this is how many loops we'll need
uplift_per_step = uplift_rate * dt

# illustrate what the MPD looks like:
print(inputs)

"""
Explanation: More components means more input parameters. So this time, we're going to make our lives easier by instantiating our components from an input file. Note also that we've now switched length units to km from m.
We're going to handle our input file using the very powerful load_params Landlab function. This function can read input text files formatted in a variety of different ways, including  the yaml standard. It automatically types the values it finds in the input file (i.e., makes them int, float, string, etc.), and returns them as a Python dictionary. This dictionary is the model parameter dictionary (MPD). However, feel free to use your own way of reading in a text file. The important thing is that you end up with a dictionary that contains 'input_parameter_name': parameter_value pairs. Note that the file format has subsets of parameters grouped, using indentation:
yaml
stream_power:
  K_sp: 0.3
  m_sp: 0.5
linear_diffuser:
  linear_diffusivity: 0.0001
When read into a dictionary, this forms two sub-dictionaries, with the keys stream_power and linear_diffuser. We will pass these two sub-dictionaries as **kwargs arguments to the FastscapeEroder and LinearDiffuser components, respectively.
End of explanation
"""


mg = RasterModelGrid((nrows, ncols), dx)
z = mg.add_zeros('topographic__elevation', at='node')
# add some roughness, as this lets "natural" channel planforms arise
initial_roughness = np.random.rand(z.size) / 100000.
z += initial_roughness
for edge in (mg.nodes_at_left_edge, mg.nodes_at_right_edge):
    mg.status_at_node[edge] = mg.BC_NODE_IS_CLOSED
for edge in (mg.nodes_at_top_edge, mg.nodes_at_bottom_edge):
    mg.status_at_node[edge] = mg.BC_NODE_IS_FIXED_VALUE

"""
Explanation: Now instantiate the grid, set the initial conditions, and set the boundary conditions:
End of explanation
"""


fr = FlowAccumulator(mg)
sp = FastscapeEroder(mg, **inputs['stream_power'])
lin_diffuse = LinearDiffuser(mg, **inputs['linear_diffuser'])

"""
Explanation: So far, so familiar.
Now we're going to instantiate all our components, using the MPD. We can do this using a bit of Python magic that lets you pass dictionaries into functions as sets of keywords. We do this by passing the dictionary as the final input, with to asterisks - ** in front of it:
End of explanation
"""


for i in range(nt):
    # lin_diffuse.run_one_step(dt) no diffusion this time
    fr.run_one_step() # run_one_step isn't time sensitive, so it doesn't take dt as input
    sp.run_one_step(dt)
    mg.at_node['topographic__elevation'][mg.core_nodes] += uplift_per_step # add the uplift
    if i % 20 == 0:
        print ('Completed loop %d' % i)

"""
Explanation: What's happening here is that the component is looking inside the dictionary for any keys that match its keywords, and using them. Values in the dictionary will override component defaults, but note that you cannot provide a keyword manually that is also defined in a supplied dictionary, i.e., this would result in a TypeError:
```python



lin_diffuse = LinearDiffuser(mg, linear_diffusivity=1.,
                                 **{'linear_diffusivity': 1.})
TypeError
```



A note on the FlowAccumulator. This component provides a variety of options for the flow direction method used (e.g., D4/SteepestDescent, D8, MFD etc.). By default it uses D4 flow routing and does not deal with depression finding and routing.
In order to use the DepressionFinderAndRouter inside the FlowAccumulator specify depression_finder = 'DepressionFinderAndRouter'.
If you are using the FlowAccumulator in additional projects or using this notebook as a starting place for additional work, work through the three tutorials on the FlowDirectors and the FlowAccumulator first.
And now we run! We're going to run once with the diffusion and once without.
End of explanation
"""


figure('topo without diffusion')
imshow_grid(mg, 'topographic__elevation', grid_units=['km','km'], var_name='Elevation (km)')

"""
Explanation: You'll need to give the above code a few seconds to run.
End of explanation
"""


z[:] = initial_roughness
for i in range(nt):
    lin_diffuse.run_one_step(dt)  # no diffusion this time
    fr.run_one_step() # run_one_step isn't time sensitive, so it doesn't take dt as input
    sp.run_one_step(dt)
    mg.at_node['topographic__elevation'][mg.core_nodes] += uplift_per_step # add the uplift
    if i % 20 == 0:
        print ('Completed loop %d' % i)

figure('topo with diffusion')
imshow_grid(mg, 'topographic__elevation', grid_units=['km','km'],
            var_name='Elevation (km)')

"""
Explanation: And now let's reset the grid elevations and do everything again, but this time, with the diffusion turned on:
End of explanation
"""


from landlab.components import ChannelProfiler, PrecipitationDistribution
from matplotlib.pyplot import loglog

z[:] = initial_roughness

"""
Explanation: Beautiful! We've smoothed away the fine-scale channel roughness, as expected, and produced some lovely convex-up hillslopes in its place. Note that even though the initial conditions were identical in both cases, including the roughness, the channel positions have been moved significantly by the hillslope diffusion into the channel.
As a final step, we're going to show off some of Landlab's fancier functionality. We're going to repeat the above coupled model run, but this time we're going to plot some evolving channel profiles, and we're going to drive the simulation with a sequence of storms, not just a fixed timestep. We'll also produce a slope-area plot for the final conditions.
Working with timesteps of varying length requires a bit more bookkeeping, but the principle is the same as what we've seen before.
So, load the new landlab objects we'll need, then reset the initial conditions:
End of explanation
"""


dt = 0.1
total_t = 250.

storm_inputs = load_params('./coupled_params_storms.txt')
precip = PrecipitationDistribution(total_t=total_t, delta_t=dt, **storm_inputs)
print(storm_inputs)

# make a color mapping appropriate for our time duration
norm = mpl.colors.Normalize(vmin=0, vmax=total_t)
map_color = cm.ScalarMappable(norm=norm, cmap="viridis")

"""
Explanation: Instantiate the storm generator. This time, we're going to mix an input file for some components with manual definition of others (that we already defined above).
End of explanation
"""


out_interval = 20.
last_trunc = total_t # we use this to trigger taking an output plot
for (interval_duration, rainfall_rate) in precip.yield_storm_interstorm_duration_intensity():
    if rainfall_rate > 0.:
        # note diffusion also only happens when it's raining...
        fr.run_one_step()
        sp.run_one_step(interval_duration)
        lin_diffuse.run_one_step(interval_duration)
    z[mg.core_nodes] += uplift_rate * interval_duration
    this_trunc = precip.elapsed_time // out_interval
    if this_trunc != last_trunc:  # time to plot a new profile!
        print('made it to time %d' % (out_interval * this_trunc))
        last_trunc = this_trunc
        figure("long_profiles")
        # get and plot the longest profile
        cp = ChannelProfiler(mg)
        cp.run_one_step()
        cp.plot_profiles(color = map_color.to_rgba(precip.elapsed_time))
        
    # no need to track elapsed time, as the generator will stop automatically
# make the figure look nicer:
figure("long_profiles")
xlabel('Distance upstream (km)')
ylabel('Elevation (km)')
title('Long profiles evolving through time')
mpl.pyplot.colorbar(map_color)

"""
Explanation: Now run:
End of explanation
"""


figure('topo with diffusion and storms')
imshow_grid(mg, 'topographic__elevation', grid_units=['km','km'], var_name='Elevation (km)')

"""
Explanation: Note that the "wobbles" in the long profile here are being created by the stochastic storm sequence. We could reduce their impact by reducing the storm-interstorm timescales, or allowing diffusion while it's not raining, but we've chosen not to here to show that the storms are having an effect.
End of explanation
"""


cp = ChannelProfiler(mg, number_of_watersheds=7, minimum_channel_threshold=0.01, main_channel_only=False)
cp.run_one_step()
cp.plot_profiles_in_map_view()

"""
Explanation: We can also plot the location of the channels in map view. 
Here we plot all channel segments with drainage area greater than 0.01 square kilometers in the seven biggest drainage basins.
End of explanation
"""


figure('final slope-area plot')
loglog(mg.at_node['drainage_area'], mg.at_node['topographic__steepest_slope'],'.')
xlabel('Drainage area (km**2)')
ylabel('Local slope')
title('Slope-Area plot for whole landscape')

"""
Explanation: Next we make a slope area plot.
End of explanation
"""


