import pandas as pd
import pastas as ps
import matplotlib.pyplot as plt

ps.show_versions(numba=True)
ps.set_log_level("INFO")

"""
Explanation: Non-linear recharge models
R.A. Collenteur, University of Graz
This notebook explains the use of the RechargeModel stress model to simulate the combined effect of precipitation and potential evaporation on the groundwater levels. For the computation of the groundwater recharge, four recharge models are currently available:

Linear (Berendrecht et al., 2003; von Asmuth et al., 2008)
Berendrecht (Berendrecht et al., 2006)
FlexModel (Collenteur et al., 2021)
Peterson (Peterson & Western, 2014)

The first model is a simple linear function of precipitation and potential evaporation while the latter two are simulate a non-linear response of recharge to precipitation using a soil-water balance concepts. Detailed descriptions of these models can be found in articles listed in the References at the end of this notebook.
<div class="alert alert-info">

<b>Tip</b> 

To run this notebook and the related non-linear recharge models, it is strongly recommended to install Numba (http://numba.pydata.org). This Just-In-Time (JIT) compiler compiles the computationally intensive part of the recharge calculation, making the non-linear model as fast as the Linear recharge model.

</div>
End of explanation
"""


head = pd.read_csv("../data/B32C0639001.csv",  parse_dates=['date'], 
                   index_col='date', squeeze=True) 

# Make this millimeters per day
evap = ps.read_knmi("../data/etmgeg_260.txt", variables="EV24").series * 1e3
rain = ps.read_knmi("../data/etmgeg_260.txt", variables="RH").series * 1e3

ps.plots.series(head, [evap, rain], figsize=(10,6), 
                labels=["Head [m]", "Evap [mm/d]", "Rain [mm/d]"]);

"""
Explanation: Read Input data
Input data handling is similar to other stressmodels. The only thing that is necessary to check is that the precipitation and evaporation are provided in mm/day. This is necessary because the parameters for the non-linear recharge models are defined in mm for the length unit and days for the time unit. It is possible to use other units, but this would require manually setting the initial values and parameter boundaries for the recharge models.
End of explanation
"""


ml = ps.Model(head)

# Select a recharge model
rch = ps.rch.FlexModel(gw_uptake=True)
#rch = ps.rch.Berendrecht()
#rch = ps.rch.Linear()
#rch = ps.rch.Peterson()

rm = ps.RechargeModel(rain, evap, recharge=rch, rfunc=ps.Gamma, name="rch")
ml.add_stressmodel(rm)

ml.solve(noise=True, tmin="1990", report="basic")
ml.plots.results(figsize=(10,6));

"""
Explanation: Make a basic model
The normal workflow may be used to create and calibrate the model.
1. Create a Pastas Model instance
2. Choose a recharge model. All recharge models can be accessed through the recharge subpackage (ps.rch).
3. Create a RechargeModel object and add it to the model
4. Solve and visualize the model
End of explanation
"""


recharge = ml.get_stress("rch").resample("A").sum()
ax = recharge.plot.bar(figsize=(10,3))
ax.set_xticklabels(recharge.index.year)
plt.ylabel("Recharge [mm/year]");

"""
Explanation: Analyze the estimated recharge flux
After the parameter estimation we can take a look at the recharge flux computed by the model. The flux is easy to obtain using the get_stress method of the model object, which automatically provides the optimal parameter values that were just estimated. After this, we can for example look at the yearly recharge flux estimated by the Pastas model.
End of explanation
"""


