"""
Plotting different process variables on already existing Matplotlib figures / axes
==================================================================================

This tutorial covers plotting meshseries data using user defined matplotlib
objects for figure and / or axes. This is useful if different plotting functions
ogstools are to be used on different subplots within the same figure.
"""

# %%
import matplotlib.pyplot as plt

import ogstools as ot
from ogstools import examples

meshseries = examples.load_meshseries_THM_2D_PVD(spatial_unit="km")

ot.plot.setup.combined_colorbar = False

# %% [markdown]
# Compare different variables
# ----------------------------

# %%
# It is possible to plot various variables in different subplots of the
# same figure:

fig, ax = plt.subplots(2, 1, figsize=(15, 15))
ot.plot.contourf(
    meshseries.mesh(0), ot.variables.temperature, fig=fig, ax=ax[0]
)
ot.plot.contourf(
    meshseries.mesh(1), ot.variables.displacement, fig=fig, ax=ax[1]
)
fig.tight_layout()

# %% [markdown]
# Plot two time steps and their difference
# ----------------------------------------

# %%
# We can use the same approach to plot the difference between different time
# steps can be plotted. Color bars can be drawn automatically, if user provides
# both Figure and Axes objects:

fig, ax = plt.subplots(3, 1, figsize=(20, 30))

ot.plot.contourf(
    meshseries.mesh(0), ot.variables.temperature, fig=fig, ax=ax[0]
)
ot.plot.contourf(
    meshseries.mesh(1), ot.variables.temperature, fig=fig, ax=ax[1]
)
diff_mesh = ot.mesh.difference(
    meshseries.mesh(1), meshseries.mesh(0), ot.variables.temperature
)


ot.plot.contourf(diff_mesh, ot.variables.temperature, fig=fig, ax=ax[2])
ax[0].set_title(r"$T(\mathrm{t}_{0})$")
ax[1].set_title(r"$T(\mathrm{t}_{end})$")
ax[2].set_title(r"$T(\mathrm{t}_{end})$-$T(\mathrm{t}_{0})$")
fig.tight_layout()