from landlab import RasterModelGrid
from landlab.components import FastscapeEroder, FlowAccumulator
from landlab.plot import channel_profile, imshow_grid
from IPython.display import HTML, Image
import matplotlib.animation as animation
import matplotlib.pylab as plt
import numpy as np

"""
Explanation: <a href="http://landlab.github.io"><img style="float: left"
src="https://raw.githubusercontent.com/landlab/tutorials/release/landlab_header.png"></a>
Animate Landlab output
<hr>

<p><small>More Landlab tutorials:
<a href="https://landlab.readthedocs.io/en/latest/user_guide/tutorials.html">https://landlab.readthedocs.io/en/latest/user_guide/tutorials.html</a></small></p>

<hr>

WARNING: This tutorial is not tested. This is because it relys on ffmpeg and imagemagik. 
It has not been updated to work with Landlab 2.0
Introduction
This tutorial presents a workflow to animate model output. The workflow is presented in two phases of a Landlab model.
This workflow requires software that might not be installed on your computer. The software is open-source and freely available for Linux, MacOS, and Windows. The following software is required for the model phases of this tutorial:
* Phase 1: A mp4 video format is produced using FFMpeg (use conda install ffmpeg, or the pip equivalent, or visit https://www.ffmpeg.org/download.html).
* Phase 2: A gif video format is produced using ImageMagick (use conda install imagemagick, or the pip equivalent, or visit https://www.imagemagick.org/script/download.php).
Prepare the model for both phases
Import the modules called by this tutorial.
End of explanation
"""


mg = RasterModelGrid((40, 40), 100)
z = mg.add_zeros('topographic__elevation', at='node')
z += np.random.rand(z.size)

outlet_id = int(mg.number_of_node_columns * 0.5)
mg.set_watershed_boundary_condition_outlet_id(outlet_id, z)
mg.at_node['topographic__elevation'][outlet_id] = 0

fr = FlowAccumulator(mg)
sp = FastscapeEroder(mg, K_sp=3e-5, m_sp=0.5, n_sp=1)

"""
Explanation: Create a grid with random elevation, set boundary conditions, and initialize components.
End of explanation
"""


simulation_duration = 1e6
dt = 1000
n_timesteps = int(simulation_duration // dt) + 1
timesteps = np.linspace(0, simulation_duration, n_timesteps)

uplift_rate = 0.001
uplift_per_timestep = uplift_rate * dt

"""
Explanation: Set model time and uplift parameters.
End of explanation
"""


# Create a matplotlib figure for the animation.
fig, ax = plt.subplots(1, 1)

# Initiate an animation writer using the matplotlib module, `animation`.
# Set up to animate 6 frames per second (fps)
writer = animation.FFMpegWriter(fps=6)

# Setup the movie file.
writer.setup(fig, 'first_phase.mp4')

for t in timesteps:
    # Uplift and erode.
    z[mg.core_nodes] += uplift_per_timestep
    fr.run_one_step()
    sp.run_one_step(dt)

    # Update the figure every 50,000 years.
    if t % 5e4 == 0:
        imshow_grid(mg, z, colorbar_label='elevation (m)')
        plt.title('{:.0f} kiloyears'.format(t * 1e-3))

        # Capture the state of `fig`.
        writer.grab_frame()

        # Remove the colorbar and clear the axis to reset the
        # figure for the next animation timestep.
        plt.gci().colorbar.remove()
        ax.cla()

plt.close()

"""
Explanation: Phase 1: Animate elevation change using imshow_grid
We first prepare the animation movie file. The model is run and the animation frames are captured together.
End of explanation
"""


writer.finish()

"""
Explanation: Finish the animation
The method, writer.finish completes the processing of the movie and saves then it.
End of explanation
"""


HTML("""<div align="middle"> <video width="80%" controls loop>
     <source src="first_phase.mp4" type="video/mp4"> </video></div>""")

"""
Explanation: This code loads the saved mp4 and presents it in a Jupyter Notebook.
End of explanation
"""


increased_uplift_per_timestep = 10 * uplift_per_timestep

"""
Explanation: Phase 2: Animate multiple visualizations of elevation change over time
In the second model phase, we will create an animation similar to the one above, although with the following differences:
* The uplift rate is greater.
* The animation file format is gif.
* The figure has two subplots.
* The data of one of the subplots is updated rather than recreating the plot from scratch for each frame.
* The animation frame rate (fps) is lower.
Increase uplift rate prior to running the second phase of the model.
End of explanation
"""


# Create a matplotlib figure for the animation.
fig2, axes = plt.subplots(1, 2, figsize=(9, 3))
fig2.subplots_adjust(top=0.85, bottom=0.25, wspace=0.4)

# Layout right subplot.

time = 0

line, = axes[1].plot(time, z.mean(), 'k')

axes[1].set_title('mean elevation over time')
axes[1].set_xlim([0, 1000])
axes[1].set_ylim([0, 1000])
axes[1].set_xlabel('time (kyr)')
axes[1].set_ylabel('elevation (m)')

# Initiate a writer and set up a movie file.
writer = animation.ImageMagickWriter(fps=2)
writer.setup(fig2, 'second_phase.gif')

for t in timesteps:
    # Uplift and erode.
    z[mg.core_nodes] += increased_uplift_per_timestep
    fr.run_one_step()
    sp.run_one_step(dt)

    # Update the figure every 50,000 years.
    if t % 5e4 == 0:
        fig2.sca(axes[0])
        fig2.suptitle('{:.0f} kiloyears'.format(t * 1e-3))

        # Plot the left subplot.
        axes[0].set_title('topography')
        imshow_grid(mg, z, colorbar_label='elevation (m)')
        colorbar = plt.gci().colorbar

        # Update the right subplot.
        line.set_xdata(np.append(line.get_xdata(), t * 1e-3))
        line.set_ydata(np.append(line.get_ydata(), z.mean()))

        # Capture the state of `fig2`.
        writer.grab_frame()

        # Reset the figure for the next animation time step.
        plt.cla()
        colorbar.remove()

writer.finish()

plt.close()

"""
Explanation: Run the second phase of the model
Here we layout the figure with a left and right subplot.
* The left subplot will be an animation of the grid similar to phase 1. We will recreate the image of this subplot for each animation frame.
* The right subplot will be a line plot of the mean elevation over time. We will layout the subplot elements (labels, limits) before running the model, and then extend the plot line at each animation frame.
axes[0] and axes[1] refer to the left and right subplot, respectively.
A gif formatted movie is created in this model phase using the software, ImageMagick.
End of explanation
"""


Image(filename='second_phase.gif')

"""
Explanation: This code loads the saved mp4 and presents it in a Jupyter Notebook.
End of explanation
"""


