import pandas as pd
import pastas as ps

ps.set_log_level("ERROR")
ps.show_versions()

"""
Explanation: Adding Trends
R.A. Collenteur (University of Graz), O.N. Ebbens (Artesia)
In this notebook it is explained how to use linear and step trend models to improve the simulation of groundwater levels.
End of explanation
"""


# Load the input data
evap = ps.read_knmi("../data/etmgeg_260.txt", variables="EV24").series * 1e3
rain = ps.read_knmi("../data/etmgeg_260.txt", variables="RH").series * 1e3
head = ps.read_dino("../data/B32C0609001_1.csv")

# Create a Pastas model
ml = ps.Model(head)

# Add a recharge model
rch = ps.rch.FlexModel()
rm = ps.RechargeModel(rain, evap, recharge=rch, rfunc=ps.Exponential, name="rch")
ml.add_stressmodel(rm)

# Solve and plot the model
ml.solve(noise=False, tmin="1990", report=False)  # Get better initial estimated first
ml.solve(noise=True, tmin="1990", initial=False, report=False)
ml.plots.results(figsize=(10, 6));

"""
Explanation: 1. Modeling a linear trend
In this first example we look at a model where a linear trend is used to improve the simulation of the groundwater levels. The linear trend is modeled using the LinearTrend stress model. We start with a model where we try to explain the groundwater level fluctuations using precipitation and evaporation. A simple non-linear recharge model is used to translate these fluxes into recharge and finally groundwater levels.
End of explanation
"""


# Add a linear trend
tm = ps.LinearTrend(start="1990-01-01", end="2020-01-01", name="trend")
ml.add_stressmodel(tm)

# Solve the model
ml.solve(noise=False, tmin="1990", report=False)  # Get better initial estimated first
ml.solve(noise=True, tmin="1990", initial=False, report=False)
ml.plots.results(figsize=(10, 6));

"""
Explanation: Add linear trend to the model
Clearly the model fit with the data in the above figure is not so good. Looking at the model residuals (simulation - observation) we can observe a steady upward trend in the residuals. Let's try and add a linear trend to the model to improve the groundwater level simulation.
End of explanation
"""


p = pd.read_csv("../data/nb18_rain.csv", index_col=0, parse_dates=True, squeeze=True)
e = pd.read_csv("../data/nb18_evap.csv", index_col=0, parse_dates=True, squeeze=True)
h = pd.read_csv("../data/nb18_head.csv", index_col=0, parse_dates=True, squeeze=True)

ml = ps.Model(h.iloc[::10])

sm = ps.RechargeModel(p, e, name="recharge", rfunc=ps.Exponential, recharge=ps.rch.Linear())
ml.add_stressmodel(sm)
ml.solve(report=False)

ml.plots.results(figsize=(10,5));

"""
Explanation: 2. Modeling a step trend
In this example the modeling of step trends in groundwater level time series is explored. Step trends can be used when a system change has taken place during the observation period, for example a lowering of the surrounding water levels. Here we model a groundwater level time series observed near the city of Eindhoven in the Netherlands that has undergone a structural change during the time of observation. The change has taken place in 2012, but unfortunately no observations are available for the period when the change was made.
model with precipitation and evaporation
First a model with only precipitation and potential evaporation as explanatory variables is created. It can be observed that the peak in the groundwater levels after 2012 lie about 0.5 meters lower that the peaks before 2012. This can also be observed by studying the model residuals, which show a different mean for the period before 2012 and after 2012.
End of explanation
"""


step = ps.StepModel(tstart=pd.Timestamp("2012-07-01"), name="step", up=None)
ml.add_stressmodel(step)
ml.solve(report=False)

ml.plots.results(figsize=(10,5));

"""
Explanation: Adding a step trend
Clearly, the model creating above has room for improvements. Since we know a system change has taken place around 2012, we can try and simulate this system change by adding a step trend. We do not know the absolute size of this step trend, only the time when it has taken place (1st of July 2012). The StepModel is available in Pastas to add a step trend to a model. This stress model does not require a independent time series, but a(indication of) the date of the system change instead. By deault the StepModel applies an instantaneous step in the groundwater level, using the One response function. it is however also possible to choose a more gradual change by applying another type of response function.
End of explanation
"""


ml.plots.diagnostics(figsize=(10, 4));

"""
Explanation: Diagnostic checking
In the following, the uncertainty of the estimated step trend is evaluated. Before we quantify the uncertainty, we perform some model diagnostic checks on the noise time series using ps.plots.diagnostics(). From this plot, we can say that the noise looks normally distributed and has no significant autocorrelation. Thus we may use the standard errors of the estimated parameters to quantify the uncertainty of the estimated step trend.
End of explanation
"""


ci = ml.fit.ci_contribution("step", alpha=0.05)
axes = ml.plots.results(adjust_height=False, figsize=(10,5))
axes[-2].fill_between(ci.index, ci.iloc[:, 0], ci.iloc[:, 1], zorder=-10, alpha=0.5);

"""
Explanation: 3. Uncertainty of step trend
Below we draw the 95% confidence interval of the step trend. We may use this confidence interval to interpret the results. For example, if we know or expect that the step trend is negative, we can reject the step trend from the model if the confidence interval also includes positive step trends. In this example this is not the case and the estimated step trend is always negative as expected.
End of explanation
"""


# method 1
ml_before = ps.Model(h.iloc[::10][:"2012-07-01"])

sm = ps.RechargeModel(p, e, name="recharge", rfunc=ps.Exponential, recharge=ps.rch.Linear())
ml_before.add_stressmodel(sm)
ml_before.solve(report=False)

ax = h.iloc[::10].plot(marker='.', color='k', ls='none', label='head', figsize=(10, 3))
ml_before.simulate().plot(ax=ax, label='model fit')
ml_before.simulate(tmin="2012-07-01", 
                   tmax=h.iloc[::10].index[-1]).plot(ax=ax, label='simulation')

ax.legend(ncol=3)
ax.grid()

# method 2
ml_after = ps.Model(h.iloc[::10][pd.Timestamp("2012-07-01"):])

sm = ps.RechargeModel(p, e, name="recharge", rfunc=ps.Exponential, 
                      recharge=ps.rch.Linear())
ml_after.add_stressmodel(sm)
ml_after.solve(report=False)


ax = h.iloc[::10].plot(marker='.', color='k', ls='none', label='head', figsize=(10, 3))
ml_after.simulate().plot(ax=ax, label='model fit')
ml_after.simulate(tmin=h.iloc[::10].index[0], 
                  tmax=pd.Timestamp("2012-07-01")).plot(ax=ax, label='simulation')

ax.legend(ncol=3)
ax.grid()

"""
Explanation: 4. Effect of system changes
Sometimes changes to a groundwater system cannot be easily modelled using a step trend. For example when the system change causes a different response of the groundwater to recharge. In that case there are other methods to visualise the effect of the system change on the groundwater head. Two of these methods are shown below:

Fit the model on the observations before the system change. Use this model to simulate the groundwater head after the system change. The differences between simulated groundwater heads an observations are an indication of the effect of the system change on the groundwater head.
similar to method 1. Now the model is fit on the period after the system change and groundwater heads are simulated for the period before the system change.
End of explanation
"""


