"""

Ravaflow3G Mixture Model
========================
"""

# %% md
# 
# This example shows how to simulate mass flow on a topography using
# :class:`.Ravaflow3GMixture`.
#
#

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

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 `temp_Ravaflow3GMixture_example` to save output
# files generated by r.avaflow
cwd = os.getcwd()
if not os.path.exists('temp_Ravaflow3GMixture_example'):
    os.mkdir('temp_Ravaflow3GMixture_example')
dir_sim = os.path.join(cwd, 'temp_Ravaflow3GMixture_example')

# %% md
# 
# Now, we can create an instance of :class:`.Ravaflow3GMixture` by given
# ``dir_sim`` and leaving other parameters to their default values
# (Other parameters include ``time_step``, ``time_end``, ``curvature_control``,
# ``entrainment_control``, ``stopping_control``, etc.)

voellmy_model = Ravaflow3GMixture(dir_sim)

# %% md
#
# To run a simulation using above `voellmy_model`, one needs to specify 
#     1. ``elevation`` – Name of elevation raster file (including its path).
#     2. ``hrelease`` – Name of release height raster file (including its path).
#     3. ``prefix`` – Prefix required by r.avaflow to name output files.
#     4. If ``elevation`` is not a georeferenced file which can be used to create
#        a `GRASS Location`, ``EPSG`` must be provided.
#
# Other optional parameters include ``internal_friction``, ``basal_friction``,
# ``turbulent_friction``, ``entrainment_coef`` etc.
#
# The synthetic topography ``synthetic_topo.tif`` and release mass ``synthetic_rel.tif``
# are used here for illustration. They are located at the `/tests/data/` folder.

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')
prefix = 'synthetic'

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

# %% md
# 
# Given above inputs, we can create a `GRASS Location` and a shell file by calling
# :py:meth:`.Ravaflow3GMixture.preprocess` method, and run the simulation using
# :py:meth:`.Ravaflow3GMixture.run` method

grass_location, sh_file = voellmy_model.preprocess(
    prefix=prefix, elevation=elevation, hrelease=hrelease, EPSG='2326')
voellmy_model.run(grass_location, sh_file)

# %% md
# 
# Once the simulation is finished, the results are located in the folder
# `/dir_sim/your_prefix_results`. In this case, they are located in
# `/temp_Ravaflow3GMixture_example/synthetic_results`. We can extract desired
# outputs using respective method and visualize them.
# For example, we can extract the overall impact area, maximum flow velocity, or
# maximum flow velocity at specific locations:

import numpy as np

# overall impact area
impact_area = voellmy_model.extract_impact_area(prefix)
print(f"Impact_area is {impact_area} m^2")

# overall maximum flow velocity
v_mmax = voellmy_model.extract_qoi_max(prefix, 'v', aggregate=True)
print(f"Maximum flow velocity is {v_mmax} m/s^2")

# maximum flow velocity at specific locations
loc = np.array([[500, 2000], [1500, 2000]])
v_max_loc = voellmy_model.extract_qoi_max_loc(prefix, loc, 'v')
print(f"Maximum flow velocity at location {loc[0]} is {v_max_loc[0]} m/s^2")
print(f"Maximum flow velocity at location {loc[1]} is {v_max_loc[1]} m/s^2")

# %% md
# 
# We can visulize point-wise maximum flow height and velocity using heatmaps:

import matplotlib.pyplot as plt
import linecache

# extract head information of result ascii files and compute x and y coordinates
hmax_asc = os.path.join(dir_sim, f'{prefix}_results', f'{prefix}_ascii', f'{prefix}_hflow_max.asc')
        
header = [linecache.getline(hmax_asc, i) for i in range(1,6)]
header_values = [float(h.split()[-1].strip()) for h in header]
ncols, nrows, xll, yll, cellsize = header_values
ncols = int(ncols)
nrows = int(nrows)

x = np.arange(xll, xll+(cellsize*ncols), cellsize)
y = np.arange(yll, yll+(cellsize*nrows), cellsize)

# point-wise maximum flow height
h_max = voellmy_model.extract_qoi_max(prefix, 'h', aggregate=False)
# point-wise maximum flow velocity
v_max = voellmy_model.extract_qoi_max(prefix, 'v', aggregate=False)

fig, ax = plt.subplots(2, 1, figsize=(10,6))

fig0 = ax[0].contourf(x, y, h_max, levels=20)
ax[0].set_xlabel('x')
ax[0].set_ylabel('y')
ax[0].set_title('maximum flow height' + r' $(m/s)$')
cbar0 = plt.colorbar(fig0, ax=ax[0], format='%.1f', orientation='vertical', fraction=0.1)

fig1 = ax[1].contourf(x, y, v_max, levels=20)
ax[1].set_xlabel('x')
ax[1].set_ylabel('y')
ax[1].set_title('maximum flow velocity' + r' ($m/s^2$)')
cbar1 = plt.colorbar(fig1, ax=ax[1], format='%.1f', orientation='vertical', fraction=0.1)

plt.tight_layout()

# %% md
# 
# Here we delete the folder `temp_Ravaflow3GMixture_example` and all files therein.

import shutil

shutil.rmtree(dir_sim)