#!/usr/bin/env python
# coding: utf-8
Copyright 2020 Marco Dal Molin et al.

This file is part of SuperflexPy.

SuperflexPy is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

SuperflexPy is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with SuperflexPy. If not, see <https://www.gnu.org/licenses/>.

This file is part of the SuperflexPy modelling framework. For details about it,
visit the page https://superflexpy.readthedocs.io

CODED BY: Marco Dal Molin
DESIGNED BY: Marco Dal Molin, Fabrizio Fenicia, Dmitri Kavetski
# # SuperflexPy example 02: Calibrate a model
# 
# Author: Marco Dal Molin
# 
# Collaborators: Fabrizio Fenicia, Dmitri Kavetski

# ## What's in this example
# 
# This example will cover the following aspects:
# - Calibration of a simple model
# - [Benefits of using the numba implementation](https://superflexpy.readthedocs.io/en/latest/numerical_solver.html#computational-efficiency-with-numpy-and-numba)
# 
# By clicking on the items you will be redirected to the documentation page (when available) that explains the arguments in a more detailed way. 

# ## What's not in this example
# 
# The following aspects are already covered in other examples:
# 
# - [Initialize already implemented elements](./03_init_single_element_model.ipynb)
# - [Put elements together to initialize a Unit](./04_init_single_unit_model.ipynb)
# - [Run a model](./01_run_simple_model.ipynb)
# - [Change states and parameters](./01_run_simple_model.ipynb)
# 
# For this reason we will put only the code need, without further explanations. You can check the other examples by clicking on the items above.

# ## Do you want to use this example as a script?
# 
# Examples of SuperflexPy are created and distributed using Jupyter notebooks because they enable to conjugate runnable code with text explanations. We have decided to not provide the content of the notebooks as script because this would mean duplicating the same content in multiple places, bringing to maintainability problems.
# 
# If the user wants to download the content of this notebook as a python script, it is possible following the steps:
# 1. If not already done, open this notebook using [Binder](https://mybinder.org/v2/gh/dalmo1991/superflexPy/master?filepath=examples%2F02_calibrate_a_model.ipynb)
# 2. Go on File -> Download as -> Python (.py)
# 3. Select the saving path
# 
# The result is a Python script that contains all the markdown text as comment and the code contained inside the cells.

# ## STEP 0: Import of external libraries
# 
# The execution of the code in this example relies on the following external libraries:
# - **numpy**: arrays handling
# - **datetime**: measure the time
# 
# We assume that those libraries are already installed together with the latest version of SuperflexPy. Keep in mind that not all the libraries listed above are strictly needed to execute SuperflexPy and, therefore, they are not listed as requirements of SuperflexPy.

# In[1]:


import numpy as np
from datetime import datetime


# ## STEP 1: Initialize the model
# 
# In this example we will use the Unit as maximum complexity level of the model. The same operations illustrated in this example can be done similarly with any SuperflexPy component (i.e. Element, Unit, Node, Network).
# 
# The model is imported and initialized in the following cell.

# In[2]:


from superflexpy.implementation.root_finders.pegasus import PegasusPython
from superflexpy.implementation.numerical_approximators.implicit_euler import ImplicitEulerPython
from superflexpy.implementation.elements.hbv import PowerReservoir
from superflexpy.framework.unit import Unit

# Initialize numercal routines
root_finder = PegasusPython()
numeric_approximator = ImplicitEulerPython(root_finder=root_finder)

# Initialize the elements
pr_1 = PowerReservoir(
    parameters={'k': 0.1, 'alpha': 1.0},
    states={'S0': 10.0},
    approximation=numeric_approximator,
    id='PR-1'
)

pr_2 = PowerReservoir(
    parameters={'k': 0.01, 'alpha': 2.0},
    states={'S0': 1.0},
    approximation=numeric_approximator,
    id='PR-2'
)

# Initialize the Unit
model = Unit(
    layers=[
        [pr_1],
        [pr_2]
    ],
    id='model'
)


# ## STEP 2: Prepare the model to run
# 
# This involves:
# - Read the input arrays (in this example they are generated randomly)
# - Assign inputs to the model
# - Set the time step
# 
# This is done in the following cell.

# In[3]:


# Fix the seed
SEED = 2
rng = np.random.RandomState(seed=SEED)

# Create the precipitation input
P = np.zeros(100)
P[:10] = rng.randint(10, size=10)
P[25:30] = rng.randint(20, size=5)
P[40:60] = rng.randint(5, size=20)
P[80:83] = rng.randint(30, 50, size=3)

# Assign the input
model.set_input([P])

# Set the time step
model.set_timestep(1.0)


# ## STEP 3: Calibrate the model
# 
# This example will, actually, not show how to calibrate the model to data but it will illustrate how to perform all the steps that are usually needed to calibrate a model. The steps are:
# 
# 1. Generate a set of candidate parameters
# 2. Assign the parameters to the model
# 3. Reset the states before running
# 4. Run the model and return the output time series
# 
# Steps 1 to 4 are usually repeated several times until reaching satisfactory values for the parameters.

# ### 01. Generate a set of candidate parameters
# 
# We want to calibrate only 2 out of 4 parameters: `model_PR-1_k` and `model_PR-2_k`. Normally, the calibration library will provide the value of these parameters. In this case, we will generate them randomly since the goal is to show the process and not to actually calibrate the model.

# In[4]:


# Generate the parameters randomly
parameters = rng.rand(2)

# Construct the dictionary for the set_parameters method
parameters_to_set = {
    'model_PR-1_k': parameters[0],
    'model_PR-2_k': parameters[1],   
}


# ### 02. Assign the parameters to the model
# 
# This can be done using the `set_parameters` method

# In[5]:


model.set_parameters(parameters=parameters_to_set)


# ### 03. Reset the states before running
# 
# To avoid errors related to previous runs of the model, it is good practice to reset the states to their initialization value.

# In[6]:


model.reset_states()


# ### 04. Run the model and return the output time series
# 
# No the model can be run using the `get_output` method.

# In[7]:


output = model.get_output()


# ## STEP 4: Potential of using the numba implementation
# 
# During calibration (or uncertainty quantification) the steps illustrated above can be repeated several hundreds (if not thousands) of times. For this reason, using the pure Python implementation of the numerical solvers may not be the most efficient thing to do.
# 
# We will here show the advantage of using the numba implementation. Further details can be found [here](TODO).

# ### 00. Create a function to calibrate the model
# 
# This function will performs the steps seen in "STEP 3".

# In[8]:


def calibrate(model, rng, num_iterations):
    for _ in range(num_iterations):
        parameters = rng.rand(2)
        parameters_to_set = {
            'model_PR-1_k': parameters[0],
            'model_PR-2_k': parameters[1],   
        }
        model.set_parameters(parameters=parameters_to_set)
        model.reset_states()
        output = model.get_output()


# ### 01. Prepare a model that uses numba
# 
# The only thing that need to be change in order to use the model with numba is the choice of the root finder and of the numerical approximator to use.

# In[9]:


from superflexpy.implementation.root_finders.pegasus import PegasusNumba
from superflexpy.implementation.numerical_approximators.implicit_euler import ImplicitEulerNumba

# Initialize numercal routines
root_finder_numba = PegasusNumba()
numeric_approximator_numba = ImplicitEulerNumba(root_finder=root_finder_numba)

# Initialize the elements
pr_1_numba = PowerReservoir(
    parameters={'k': 0.1, 'alpha': 1.0},
    states={'S0': 10.0},
    approximation=numeric_approximator_numba,
    id='PR-1'
)

pr_2_numba = PowerReservoir(
    parameters={'k': 0.01, 'alpha': 2.0},
    states={'S0': 1.0},
    approximation=numeric_approximator_numba,
    id='PR-2'
)

# Initialize the Unit
model_numba = Unit(
    layers=[
        [pr_1_numba],
        [pr_2_numba]
    ],
    id='model'
)

# Assign the input
model_numba.set_input([P])

# Set the time step
model_numba.set_timestep(1.0)


# ### 02. Test the Python model
# 
# We will simulate a calibration that requires 1000 model runs to calibrate.

# In[10]:


start = datetime.now()
calibrate(model=model, rng=rng, num_iterations=1000)
end = datetime.now()

print('Calibration time: {}'.format(end - start))


# ### 03. Test the numba model
# 
# We will simulate a calibration that requires 1000 model runs to calibrate.

# In[11]:


start = datetime.now()
calibrate(model=model_numba, rng=rng, num_iterations=1000)
end = datetime.now()

print('Calibration time: {}'.format(end - start))


# ### 04. Conclusions
# 
# Although running times depends on the hardware used, numba models can be up to 30 times faster than pure Python solutions.
# 
# Note that, when using numba for the first time, most of the time is used to compile the methods to machine code (in the machine used when writing the example, compilation takes roughly 1 second while running  the model once is in the order of $10^{-3}$ seconds). Compilation needs to be done only once as long as data types do not change and, therefore, it is more efficient to use numba when the model needs to be run several times, like for calibration.