import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import pastas as ps

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

"""
Explanation: Testing an AR(1) Noise Model for Pastas
R.A. Collenteur, University of Graz, May 2020
In this notebook the classical Autoregressive AR1 noise model is tested for Pastas models. This noise model is tested against synthetic data generated with Numpy or Statsmodels' ARMA model. This noise model is tested on head time series with a regular time step.
End of explanation
"""


# Read in some data
rain = ps.read.read_knmi('../examples/data/etmgeg_260.txt', variables='RH').series
evap = ps.read.read_knmi('../examples/data/etmgeg_260.txt', variables='EV24').series

# Set the True parameters
Atrue = 800
ntrue = 1.4
atrue = 200
dtrue = 20

# Generate the head
step = ps.Gamma().block([Atrue, ntrue, atrue], cutoff=0.9999)
h = dtrue * np.ones(len(rain) + step.size)
for i in range(len(rain)):
    h[i:i + step.size] += rain[i] * step
head = pd.DataFrame(index=rain.index, data=h[:len(rain)],)
head = head['1990':'2015']

# Plot the head without noise
plt.figure(figsize=(10,2))
plt.plot(head,'k.', label='head')
plt.legend(loc=0)
plt.ylabel('Head (m)')
plt.xlabel('Time (years)');

"""
Explanation: 1. Develop the AR(1) Noise Model for Pastas
The following formula is used to calculate the noise according to the ARMA(1,1) process:
$$\upsilon(t_i) = r(t_i) - r(t_{i-1}) \text{e}^{-\Delta t_i / \alpha}$$
where $\upsilon$ is the noise, $\Delta t_i$ the time step between the residuals ($r$), and $\alpha$ [days] the AR parameter of the model. The model is named NoiseModel and can be found in noisemodel.py. This is the default noise model that is added to a Pastas model, but it can also be added manually as follows: ml.add_noisemodel(ps.NoiseModel())
2. Generate synthetic head time series
End of explanation
"""


# reproduction of random numbers
np.random.seed(1234)
alpha= 0.8

# generate samples using Numpy
random_seed = np.random.RandomState(1234)
noise = random_seed.normal(0,1,len(head)) * np.std(head.values) * 0.2
a = np.zeros_like(head[0])

for i in range(1, noise.size):
    a[i] = noise[i] + a[i - 1] * alpha

head_noise = head[0] + a

"""
Explanation: 3. Generate AR(1) noise and add it to the synthetic heads
In the following code-block, noise is generated using an AR(1) process using Numpy.
End of explanation
"""


ml = ps.Model(head_noise)
sm = ps.StressModel(rain, ps.Gamma, name='recharge', settings='prec')
ml.add_stressmodel(sm)
ml.add_noisemodel(ps.NoiseModel())

ml.solve(tmin="1991", tmax='2015-06-29', noise=True, report=False)

# Plot the results
axes = ml.plots.results(figsize=(10,5));
axes[-1].plot(ps.Gamma().step([Atrue, ntrue, atrue], cutoff=0.999))

"""
Explanation: 4. Create and solve a Pastas Model
End of explanation
"""


print(np.exp(-1./ml.parameters.loc["noise_alpha", "optimal"]).round(2), "vs", alpha)

"""
Explanation: 5. Did we find back the original AR parameter?
End of explanation
"""


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

"""
Explanation: The estimated parameters for the noise model are almost the same as the true parameters, showing that the model works for regular time steps.
6. So is the autocorrelation removed correctly?
End of explanation
"""


index = pd.read_csv("../examples/data/test_index.csv", parse_dates=True, 
                    index_col=0).index.round("D").drop_duplicates()
head_irregular = head_noise.reindex(index)

ml = ps.Model(head_irregular)
sm = ps.StressModel(rain, ps.Gamma, name='recharge', settings='prec', cutoff=0.99)
ml.add_stressmodel(sm)
ml.add_noisemodel(ps.NoiseModel())

# ml.set_parameter("noise_alpha", 9.5, vary=False)
# ml.set_parameter("recharge_A", 800, vary=False)
# ml.set_parameter("recharge_n", 1.2, vary=False)
# ml.set_parameter("recharge_a", 200, vary=False)

ml.solve(tmin="1991", tmax='2015-06-29', noise=True, report=False)
axes = ml.plots.results(figsize=(10,5));
axes[-1].plot(ps.Gamma().step([Atrue, ntrue, atrue]))

print(np.exp(-1./ml.parameters.loc["noise_alpha", "optimal"]).round(2), "vs", alpha)

ml.plots.diagnostics(figsize=(10,4), acf_options={"bin_width": 0.5});

"""
Explanation: That seems okay. It is important to understand that this noisemodel will only help in removing autocorrelations at the first time lag, but not at larger time lags, compared to its AR(1) counterpart.
7. Test the NoiseModel for irregular time steps
In this final step the ArmaModel is tested for irregular timesteps, using the indices from a real groundwater level time series.
End of explanation
"""


v = ml.noise()
print("", v.loc[:"2010-01-01"].std())
print("", v.loc["2010-01-01":].std())

"""
Explanation: The above autocorrelation plot shows that there is some autocorrelation in the first time steps. However, this is probably more a numerical artifict due to estimating the ACF for irregular time steps than "real" autocorrelation. This can be shown when estimating the ACF for the period with more frequent observations.
End of explanation
"""


