"""
Model time discretization
=========================

iMOD Python provides nice functionality to discretize your models into stress
periods, depending on the timesteps you assigned your boundary conditions. This
functionality is activated with the ``create_time_discretization()`` method.

"""

# %%
# Basics
# ------
#
# To demonstrate the ``create_time_discretization()`` method, we first have to
# create a Model object. In this case we'll use a MODFLOW 6 simulation, but the
# ``imod.wq.SeawatModel`` also supports this.
#
# .. note::
#
#    The simulation created in this example misses some mandatory packes, so is
#    not able to write a functional simulation. It is purely intended to describe
#    the workings of the ``time_discretization`` method. For a fully functional
#    MODFLOW 6 model, see the examples in :ref:`mf6-introduction`.
#
# Wel'll start off with the usual imports:

import numpy as np
import pandas as pd
import xarray as xr

import imod

# %%
#
# We can discretize a simulation as follows, this creates a
# :class:`imod.mf6.TimeDiscretization` object under the key
# ``"time_discretization"``.

simulation = imod.mf6.Modflow6Simulation("example")
simulation.create_time_discretization(
    additional_times=["2000-01-01", "2000-01-02", "2000-01-04"]
)

simulation["time_discretization"]

# %%
#
# To view the data inside :class:`imod.mf6.TimeDiscretization` object:

simulation["time_discretization"].dataset

# %%
# Notice that even though we specified three points in time, only two
# timesteps are included in the ``time`` coordinate, this is because Modflow
# requires a start time and a duration of each stress period. iMOD Python
# therefore uses three points in time to compute two stress periods with a
# duration.

simulation["time_discretization"].dataset["timestep_duration"]

# %%
# These two stress periods use their respective start time in their ``time``
# coordinate.
#
# Boundary Conditions
# -------------------
#
# The ``create_time_discretization`` method becomes especially useful if we add boundary
# conditions to our groundwater model. We'll first still have to initialize a
# groundwater flow model though:

gwf_model = imod.mf6.GroundwaterFlowModel()

# %%
# We'll create a dummy grid, which is used to provide an appropriate grid to
# the boundary condition packages.

template_data = np.ones((2, 1, 2, 2))

layer = [1]
y = [1.5, 0.5]
x = [0.5, 1.5]

dims = ("time", "layer", "y", "x")

# %%
# Next, we can assign a Constant Head boundary with two timesteps:

chd_times = pd.to_datetime(["2000-01-01", "2000-01-04"])
chd_data = xr.DataArray(
    data=template_data,
    dims=dims,
    coords={"time": chd_times, "layer": layer, "y": y, "x": x},
)

gwf_model["chd"] = imod.mf6.ConstantHead(head=chd_data)

gwf_model["chd"].dataset

# %%
# We'll also assign a Recharge boundary with two timesteps, which differ from
# the ConstantHead boundary:

rch_times = pd.to_datetime(["2000-01-01", "2000-01-02"])
rch_data = xr.DataArray(
    data=template_data,
    dims=dims,
    coords={"time": rch_times, "layer": layer, "y": y, "x": x},
)

gwf_model["rch"] = imod.mf6.Recharge(rate=rch_data)

gwf_model["rch"].dataset

# %%
# We can now let iMOD Python figure out how the simulation's time should be
# discretized. It is important that we provide an endtime, otherwise the
# duration of the last stress period cannot be determined:

endtime = pd.to_datetime(["2000-01-06"])

simulation_bc = imod.mf6.Modflow6Simulation("example_bc")
simulation_bc["gwf_1"] = gwf_model

simulation_bc.create_time_discretization(additional_times=endtime)

simulation_bc["time_discretization"].dataset

# %%
# Notice that iMOD Python figured out that the two boundary conditions, both
# with two timesteps, should lead to three stress periods!
#
# Extra settings
# --------------
#
# The :class:`imod.mf6.TimeDiscretization` package also supports other settings,
# like the amount of timesteps which Modflow should use within a stress period,
# as well as a timestep multiplier, to gradually increase the timesteps modflow
# uses within a stress period. This can be useful when boundary conditions
# change very abruptly between stress periods. These settings are set by
# accessing their respective variables in the ``dataset``.

times = simulation_bc["time_discretization"].dataset.coords["time"]

simulation_bc["time_discretization"].dataset["timestep_multiplier"] = 1.5
simulation_bc["time_discretization"].dataset["n_timesteps"] = xr.DataArray(
    data=[2, 4, 4], dims=("time",), coords={"time": times}
)

simulation_bc["time_discretization"].dataset

# %%
#
# Adaptive Time Stepping
# ----------------------
#
# The :class:`imod.mf6.AdaptiveTimeStepping` package can be used to let Modflow
# decide the length of each timestep within a stress period. This can be useful
# when boundary conditions change abruptly, or when a model is highly dynamic.
# The :class:`imod.mf6.AdaptiveTimeStepping` package can be added to a
# :class:`imod.mf6.Modflow6Simulation` as follows:

simulation_bc["ats"] = imod.mf6.AdaptiveTimeStepping(
    dt_init=1e-4, dt_min=1e-4, dt_max=2.0, dt_multiplier=2.0, dt_fail_divisor=5.0
)

simulation_bc["ats"]

# %%
#
# The ``dt_init`` parameter sets the initial timestep length, ``dt_min`` and
# ``dt_max`` set the minimum and maximum timestep length allowed,
# ``dt_multiplier`` multiplies the timestep length for each consecutive timestep
# if the model converges well, and ``dt_fail_divisor`` decreases the timestep
# length if the model has convergence problems.
#
# There is one further way to influence the timestep length, which is to set the
# ``ats_percel`` parameter in the :class:`imod.mf6.AdvectionTVD` package of a
# transport model. This parameter sets the maximum fraction of a cell that a
# parcel of solute is allowed to travel in a single timestep. This can be used
# to ensure that the timestep is small enough to get accurate transport results.
# Once this is set, MODFLOW 6 will choose the smallest necessary timestep based
# on the convergence criteria and the ``ats_percel`` criteria.

transport_model = imod.mf6.GroundwaterTransportModel()
transport_model["adv"] = imod.mf6.AdvectionTVD(ats_percel=0.95)
simulation_bc["gwt_1"] = transport_model

transport_model["adv"]

# %%
#
# Conclusion
# ----------
#
# The :meth:`imod.mf6.Modflow6Simulation.create_time_discretization` method can
# be used to easily generate stress periods for a simulation, based on the
# timesteps of the boundary conditions assigned to the model. The
# :class:`imod.mf6.AdaptiveTimeStepping` package can be used to let MODFLOW 6
# decide the length of each timestep within a stress period, based on
# convergence criteria and, in case of using the ``ats_percel`` options, the
# velocity field.