"""

RunSimulator: Ravaflow3G Mixture Model
======================================
"""

# %% md
# 
# This example shows how to run multiple Ravaflow3GMixture simulations in serial
# and parallelly using :class:`.RunSimulator`.
#

# %% md
# 
# First, we import the class :class:`.Ravaflow3GMixture`.

from psimpy.simulator import Ravaflow3GMixture

# %% md
# 
# To create an instance of this class, we must specify the parameter ``dir_sim``.
# It represents the directory in which output files generated by r.avaflow will be
# saved.

import os

# Here we create a folder called `temp1_run_Ravaflow3GMixture_example` to save
# output files generated by r.avaflow
cwd = os.getcwd()
if not os.path.exists('temp1_run_Ravaflow3GMixture_example'):
    os.mkdir('temp1_run_Ravaflow3GMixture_example')
dir_sim = os.path.join(cwd, 'temp1_run_Ravaflow3GMixture_example')

# %% md
# 
# Given ``dir_sim``,  we can create an instance of :class:`.Ravaflow3GMixture`.
# To reduce simulation time, we set ``time_end`` to :math:`50`. Other parameters
# are set to their default values.

voellmy_model = Ravaflow3GMixture(dir_sim, time_end=50)

# %% md
# 
# The `simulator` in this example is defined based on ``voellmy_model`` as follows.
# It takes required inputs and returns the overall impact area.

import numpy as np

def simulator(prefix, elevation, hrelease, basal_friction, turbulent_friction, EPSG):
    """Preprocess required inputs, run simulation, and return output as a numpy array."""
    grass_location, sh_file = voellmy_model.preprocess(prefix=prefix, elevation=elevation,
        hrelease=hrelease, basal_friction=basal_friction, turbulent_friction=turbulent_friction, EPSG=EPSG)
    voellmy_model.run(grass_location, sh_file) # run this line, r.avaflow will write outputs to dir_sim
    impact_area = voellmy_model.extract_impact_area(prefix)
    
    return np.array([impact_area]) # define dir_out and set save_out to True, returned numpy array will be saved to dir_out


# %% md
# 
# To demonstrate how to run multiple simulations using :class:`.RunSimulator`, 
# we choose ``basal_friction`` and ``turbulent_friction`` as variable input
# parameters. Each has two different values, leading to four simulations.

import itertools

var_inp_parameter = ['basal_friction', 'turbulent_friction']
basal_friction = [20, 30] 
turbulent_friction = [3, 4]
var_samples = np.array(
    [x for x in itertools.product(basal_friction, turbulent_friction)])


# %% md
#
# Other parameters of the `simulator`, including ``elevation``, ``hrelease``, and
# ``EPSG`` are treated as fixed input. It means that their values are the same in
# all simulations.

dir_data = os.path.abspath('../../../tests/data/')
elevation = os.path.join(dir_data, 'synthetic_topo.tif')
hrelease = os.path.join(dir_data, 'synthetic_rel.tif')
fix_inp = {'elevation': elevation, 'hrelease': hrelease, 'EPSG': '2326'}

# %% md
# 
# The parameter ``prefix`` of the `simulator` is special. It is not involved in
# the computational model, but only used to name output files generated by r.avaflow.
# Such parameter is not defined in ``var_inp_parameter`` or ``fix_inp`` of
# :class:`.RunSimulator`. Instead, we use a seperate parameter, called ``o_parameter``
# for this purpose.

o_parameter = 'prefix'


# %% md
# 
# We may want to save outputs returned by the `simulator` at each simulation for
# later inspection or processing. In that case, we need to define ``dir_out`` and
# set ``save_out`` as `True`.

import os

cwd = os.getcwd()
if not os.path.exists('temp2_run_Ravaflow3GMixture_example'):
    os.mkdir('temp2_run_Ravaflow3GMixture_example')
dir_out = os.path.join(cwd, 'temp2_run_Ravaflow3GMixture_example')

# %% md
#
# .. note:: Please note that ``dir_out`` and ``dir_sim`` are two different
#    parameters for different purposes.
#  
#     - ``dir_out`` is a parameter of :class:`.RunSimulator`. The `simulator`
#       returns output of interest as a numpy array for each simulation. If
#       we want to save this returned output of interest (here impact area)
#       to our local machine, we need to specify the value of ``dir_out`` which
#       represents the directory in which output of interest will be saved and
#       set ``save_out`` to `True`. Otherwise, we do not need ``dir_out`` and
#       leave ``save_out`` to `False`.
#     - ``dir_sim`` is a parameter of :class:`.Ravaflow3GMixture`. :class:`.Ravaflow3GMixture`
#       relies on the third party software r.avaflow 2.4. When we call
#       :py:meth:`.Ravaflow3GMixture.run` in the function body of above `simulator`,
#       r.avaflow 2.4 will be run and it generates output files. The value of
#       ``dir_sim`` specifies the directory in which output files generated by
#       r.avaflow 2.4 are going to be saved.
#
#    The value of ``dir_out`` and ``dir_sim`` can be the same if file names 
#    have no conflict. We recommend to keep them seperate. In addition, if the function
#    body of the `simulator` does not write files to disk, ``dir_sim`` is not required. Our
#    :class:`.MassPointModel` is an example. Similarly, if we do not want to save
#    returned numpy array of the `simulator`, ``dir_out`` is not needed.

# %% md
#
# Now we can define an object of :class:`.RunSimulator` by

from psimpy.simulator import RunSimulator

run_simulator = RunSimulator(simulator, var_inp_parameter, fix_inp, o_parameter,
    dir_out, save_out=True)


# %% md
# 
# Before running simulations, we need to specify values of ``prefixes`` which
# will be used to name files generated by each r.avaflow simulation and
# each returned numpy array of `simulator`.

serial_prefixes = ["serial"+str(i) for i in range(len(var_samples))]
parallel_prefixes = ["parallel"+str(i) for i in range(len(var_samples))]

# %% md
#
# Now we can use :py:meth:`.RunSimulator.serial_run` method or
# :py:meth:`.RunSimulator.parallel_run` method to run the simulations
# in serial or parallelly.

import time

start = time.time()
run_simulator.serial_run(var_samples=var_samples, prefixes=serial_prefixes)
serial_time = time.time() - start
serial_output = run_simulator.outputs
print(f"serial_output: {serial_output}")

start = time.time()
# append setting to True means that simulation outputs of parallel run will be
# appended to above serial run outputs
run_simulator.parallel_run(var_samples, prefixes=parallel_prefixes,
    max_workers=2, append=True)
parallel_time = time.time() - start
    
parallel_output = run_simulator.outputs[len(var_samples):]
print(f"parallel_output: {parallel_output}")

    
print("Serial run time: ", serial_time)
print("Parallel run time: ", parallel_time)

# %% md
# 
# All output files returned by the `simulator` are

os.listdir(dir_out)

# %% md
# 
# All output files/directories generated by r.avaflow simulations 
# (including files/directories generated during preprocessing) are

os.listdir(dir_sim)


# %% md
# .. warning:: If one simulation failed due to whatever reason, the error massage
#    will be printed to the screen but other simulations will continue. In that 
#    case, the output file of failed simulation will not be writted to ``dir_out``.
#    Also, the element of :py:attr:`.RunSimulator.outputs` corresponding to that
#    simulation will be a string representing the error message, instead of a
#    numpy array.


# %% md
# 
# In the end, we delete the folder `temp1_run_Ravaflow3GMixture_example` and
# `temp2_run_Ravaflow3GMixture_example` and all files therein.

import shutil

shutil.rmtree(dir_sim)
shutil.rmtree(dir_out)