from __future__ import print_function

# Some magic to make plots appear within the notebook
%matplotlib inline

import numpy as np # In case we need to use numpy

from pymt import plugins

"""
Explanation: Sedflux3D

Link to this notebook: https://github.com/csdms/pymt/blob/master/docs/demos/sedflux3d.ipynb
Install command: $ conda install notebook pymt_sedflux

Import the Sedflux3D component from pymt. All of the components available to pymt are located in pymt.components. Here I've renamed the component to be Model to show that you could run these same commands with other models as well. For instance, you could instead import Child with from pymt.components import Child as Model and repeat this exercise with Child instead.
End of explanation
"""


model = plugins.Sedflux3D()

"""
Explanation: Instantiate the model. If you wanted, you could have multiple instances of the same model - each with their own state. A PyMT component is instantiate without any arguments.
End of explanation
"""


# help(model)

rm -rf _model # Clean up for the next step

"""
Explanation: PyMT adds a docstring to the model. It is important to us that credit it given where credit is due. We want to make it clear that when we include a component in our framework the component is not ours but belong to the model's author. To that end, the model documentation includes a list of model authors, DOIs, as well as full citations to use if someone uses the model. In addition, to keep models from being black boxes, and because they are all open source, we provide links to source code.
End of explanation
"""


config_file, initdir = model.setup('_model', run_duration=3650000.)

"""
Explanation: The following code is generic for all pymt components - not just Sedflux3D. First we instantiate the component and then call its setup method to create a model simulation with the necessary input files (note that this is not a BMI method but something added by PyMT). setup takes an optional argument that gives a path to a folder that will contain the necessary input files (if not provided, a temporary folder will be used). setup returns the name of the config file and the path to the folder containing it (we'll use this information later when we run initialize).
End of explanation
"""


ls _model

"""
Explanation: If we look inside the _model folder, we can see all of the necessary sedflux input files ready to run.
End of explanation
"""


model.initialize(config_file, dir=initdir)

model.set_value("channel_exit_water_flow__speed", 1.2)
model.set_value("channel_exit_x-section__mean_of_width", 400.)
model.set_value("channel_exit_x-section__mean_of_depth", 4.)
model.set_value("channel_exit_water_sediment~suspended__mass_concentration", .01)

"""
Explanation: The initialize method reads all of the input files and gets the model ready to update through time.
End of explanation
"""


for t in range(10):
    model.update()
    print(model.get_current_time())

"""
Explanation: Here we update the model for 10 time steps, printing the model time after each step.
End of explanation
"""


model.time_units

"""
Explanation: From the output it appears that sedflux is operating on a 1 year time step and is using units of days. We can double check that assumption.
End of explanation
"""


model.get_time_step()

"""
Explanation: Yep, days. And the time step...
End of explanation
"""


model.update_until(200., units='year')
model.get_current_time()

model.time_units = 'year'
print(model.get_current_time())

"""
Explanation: The update_until method gives us more control by allowing us to update to a particular time. Using the units keyword allows us to specify the "until" time. Here we update the model until 20 years.
End of explanation
"""


model.output_var_names

model.input_var_names

"""
Explanation: We've covered how to setup and run a model, now we'll have a look at how to get data out of the model. This is done with the get_value method but first we need to look at what variables the model provides. This is held as a tuple of variable names and accessed through the output_var_names attribute.
End of explanation
"""


model.get_value('sea_water__depth')

"""
Explanation: So to get the values of, say, water depth,
End of explanation
"""


model.get_value('sea_water__depth', units='parsec')

model.quick_plot('sea_water__depth', vmin=-200, vmax=200, cmap='BrBG_r')

model.set_value("channel_exit_water_sediment~suspended__mass_concentration", .5)
for t in range(10):
    model.update()
model.quick_plot('sea_water__depth', vmin=-200, vmax=200, cmap='BrBG_r')

"""
Explanation: As with previous methods, you can use the units keyword to specify units. Here we get water depths in units of parsecs.
End of explanation
"""


