"""

RunSimulator: Mass Point Model
==============================
"""

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


# %% md
# 
# First, we import the class :class:`.MassPointModel` and define a `simulator`
# based on :py:meth:`.MassPointModel.run` method.

from psimpy.simulator import MassPointModel

mpm = MassPointModel()
mpm_simulator = mpm.run


# %% md
# 
# This simulator takes required input data, solves the mass point model, and
# returns time history of the mass point's location and velocity
# (see :py:meth:`.MassPointModel.run`).
#
# Required inputs for `mpm_simulator` include:
#     1. topographic data: digital elevation model (in `ESRI ascii` format)
#     2. friction coefficients: coulomb friction and turbulent friction coefficients
#     3. initial state: initial location and initial velocity of the masspoint
#     4. computational parameters: such as time step, end time, etc.
#
# The synthetic topography ``synthetic_topo.asc`` is used here for illustration.
# It is located at the `/tests/data/` folder.

import os

dir_data = os.path.abspath('../../../tests/data/')
elevation = os.path.join(dir_data, 'synthetic_topo.asc')

# %% md
# .. note:: You may need to modify ``dir_data`` according to where you save
#    ``synthetic_topo.asc`` on your local machine.

# %% md
# 
# For this example, we are going to run multiple simulations at different values
# of coulomb friction coefficient (``coulomb_friction``) and turbulent friction
# coefficient (``turbulent_friction``) while keep other inputs fixed. Namely,
# ``coulomb_friction`` and ``turbulent_friction`` are variable input parameters.
# All other inputs are fixed input parameters.
#

import numpy as np
import itertools

# Variable iput parameters are defined as a list of strings. Their values will be
# passed as a numpy array (var_samples) when run simulations.
var_inp_parameter = ['coulomb_friction', 'turbulent_friction']
coulomb_friction = np.arange(0.1, 0.31, 0.1)
turbulent_friction = np.arange(500, 2001, 400)
var_samples = np.array(
    [x for x in itertools.product(coulomb_friction, turbulent_friction)])
print("Number of variable input samples are: ", len(var_samples))

# Fixed input parameters are defined as a dictionary. Their values are given.
fix_inp = {'elevation': elevation, 'x0': 200, 'y0': 2000, 'tend': 50}

# %% md
# .. note:: Parameters of `mpm_simulator` which are not included in
#    ``var_inp_parameter`` and ``fix_inp``, such as ``ux0``, ``uy0`` and
#    ``curvature``, will automatically take their default values.

# %% md
# 
# We may want to save outputs returned by `mpm_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

# Here we create a folder called `temp_run_MassPointModel_example` to save outputs
# returned at each simulation.
cwd = os.getcwd()
if not os.path.exists('temp_run_MassPointModel_example'):
    os.mkdir('temp_run_MassPointModel_example')
dir_out = os.path.join(cwd, 'temp_run_MassPointModel_example')


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

from psimpy.simulator import RunSimulator

run_mpm_simulator = RunSimulator(mpm_simulator, var_inp_parameter, fix_inp,
    dir_out=dir_out, save_out=True)

# %% md
# .. note:: Since outputs of each simulation are saved in the same folder ``dir_out``,
#    we need to give each file a unique name in order to avoid conflict. This is
#    realized by defining ``prefixes``.

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
# 
# Using :py:meth:`.RunSimulator.serial_run` method or
# :py:meth:`.RunSimulator.parallel_run` method, we can run the simulations
# in serial or parallelly.

import time

start = time.time()
run_mpm_simulator.serial_run(var_samples=var_samples, prefixes=serial_prefixes)
serial_time = time.time() - start
serial_output = run_mpm_simulator.outputs

start = time.time()
# max_workers controls maximum number of tasks running in parallel
run_mpm_simulator.parallel_run(var_samples, prefixes=parallel_prefixes, max_workers=4)
parallel_time = time.time() - start
parallel_output = run_mpm_simulator.outputs
    
print("Serial run time: ", serial_time)
print("Parallel run time: ", parallel_time)


# %% md
# 
# Once a simulation is done, its output is saved to ``dir_out``. All output files
# generated by above simulations are as follows: 

os.listdir(dir_out)


# %% md
# 
# Once all simulations are done, their outputs can also be accessed by
# :py:attr:`.RunSimulator.outputs` attribute, which is a list.

# Here we print the maximum velocity of each simulation in the parallel run.
max_v = [np.max(output[:,5]) for output in parallel_output]
print(
    f"Maximum velocity of {parallel_prefixes[0]} to {parallel_prefixes[-1]} are: ",
    '\n',
    max_v)


# %% 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
# 
# Here we delete the folder `temp_run_MassPointModel_example` and all files therein.

import shutil

shutil.rmtree(dir_out)