import numpy as np
import matplotlib.pyplot as plt
import timeit
import pstats, cProfile
from GWTSA import *
print 'packages succesfully imported!'
%matplotlib inline

"""
Explanation: Unsaturated Zone Module Testing
In this notebook we perform some tests to validate the unsaturated zone module applied in GWTSA for non-linear time serie models. Three models are available in GWTSA, however this notebook only considers the percolation models to serve as an example:
(1)\begin{equation}
    \frac{dS_r}{dt} = Pe(t)- K_{p}\left( \frac{S_r(t)}{S_{rmax}}\right) ^\gamma - E_a
\end{equation}
This non-linear differential equation cannot be solved analytically and is therefore solved numerically. Details on the mathematics can be found in the Appendix of the Documentation of GWTSA GWTSA_Manual.pdf. Testing of the numerical model is important to validate that the code does what it should do.
GWTSA has two numerical methods to solve the equation, activated by the solver keyword-argument: explicit (solver=0) and implicit Euler (solver=1). Since these calculations are computationally intensive, the scripts are written in Cython and compiled to the C programming language. 
The following tests are used:

Analytical solution versus the numerical solution
Use of a numerical model to check emptying and filling of the unsaturated zone (check model logics)
Random forcings to check model under different conditions
Comparison of the python script and the cythonized script (C-language)
End of explanation
"""


# Provide the forcings precipitation and potential evapotransapiration
n = 50
P = E = np.zeros(n)

# Provide the model parameters for the soil module
S_rmax = 0.01 # is equal to ±100 milimeters
K_p = 0.001 # is equal to 0.01 m/d
Gamma = 1.0
Imax = 0.0015

# Provide some details for the timesteps to calculate the soil state
t = np.arange(0,n,1)
dt= 1

S0 = 0.5 * S_rmax # Calculate the initial soil state (assumed to be half full)
S = Unsat_Zone.perc(t, P, E, S_rmax, K_p, Gamma, Imax, dt, solver=0)[1]; #1 index return the soil state S
S1 = Unsat_Zone.perc(t, P, E, S_rmax, K_p, Gamma, Imax, dt, solver=1)[1]; #1 index return the soil state S
S2 = S0 * np.exp((-K_p)*t/(S_rmax)) #Plot the exact solution when P and E are zero and Beta=1

# Make a plot of the two solutions for visual comparison
plt.plot(t,S)
plt.plot(t,S1)
plt.plot(t,S2)

plt.legend(['explicit', 'implicit', 'analytical'], loc = 'best')
plt.ylabel('S [meter]')
plt.xlabel('Time [days]')
plt.title('exact solution vs. numerical solution')

"""
Explanation: 1. ANALYTICAL SOLUTION VS. NUMERICAL SOLUTION
As stated above, solving the unsaturated zone model is not trivial. However, under certain conditions equation $(1)$ can be simplyfied such that it can be solved analytically. This provides a way to compare the numerical solution to the analytical solution. In this example we test the unsaturated zone model for the case where the precipitation and the evapotranspiration are both equal to zero. Equation $(1)$ then simplifies to:
\begin{equation}
\frac{dS_r}{dt} = - K_{sat}\left( \frac{S_r(t)}{S_{rmax}}\right) ^\gamma
\end{equation}
Which, when we fix $\beta=1.0$, we can solve analytically:
\begin{equation}
S_r(t) = S_r(0) e^{\frac{-K_p}{S_{rmax}}}
\end{equation}
We now create a timeseries of evaporation and precipitation equal to zero and define the parameter $\beta=0$ as input data for the percolation model and compare the analytical solution to the numerical solution.
End of explanation
"""


n = 365
t = np.arange(0,n,1)
P = np.zeros(n)
P[:int(n/2.0)] = 2.0 / 1000 # All forcings must be provided in meters
E = 0.001 * np.ones(n)

# Provide the model parameters for the soil module
S_rmax = 0.01 # is equal to ±100 milimeters
K_p = 0.001 # is equal to 0.01 m/d
Gamma = 1.0
Imax = 0.0015

R, S = Unsat_Zone.perc(t, P, E, S_rmax, K_p, Gamma, Imax, dt, solver=0)
R1, S1 = Unsat_Zone.perc(t, P, E, S_rmax, K_p, Gamma, Imax, dt, solver=1)

plt.figure(figsize=(15,4))
plt.subplot(121)
plt.plot(t,S)
plt.plot(t,S1)
plt.ylabel('S [meter]')
plt.xlabel('Time [days]')
plt.axhline(S_rmax, color='r')
plt.legend(['explicit', 'implicit', r'$S_rmax$'], loc='best')
plt.title('Soil State over Time')
plt.ylim(0, S_rmax+0.001)

plt.subplot(122)
plt.plot(R)
plt.plot(R1)
plt.title('recharge over time')
plt.ylabel('Recharge [m/d]')
plt.xlabel('Time [days]')
plt.legend(['explicit', 'implicit', r'$S_rmax$'], loc='best')
plt.ylim(0, max(R)+0.0001)



"""
Explanation: 2. EMPTYING AND FILLING THE UNSATURATED ZONE
In this example we investigate the filling and emptying of the unsaturated zone. The unsaturated zone has certain physical boundaries, as it cannot have negative values for $S_r$ not values larger than $S_{rmax}$. 
$P = 0.002$ m/d for:  $0 < t < 183$
$P = 0.00$ m/d for:  $183 < t $
We model the situation where the unsaturated zone is full and where it is empty, to validate that the the amount of water cannot excess the soil capacity $S_{cap}$ and cannot go below zero. We therefore first consider half a year with very wet conditions, followed by half a year with no rain.
End of explanation
"""


# Provide the forcings precipitation and potential evapotransapiration
C = np.genfromtxt('./KNMI_Bilt.txt' , delimiter=',', skip_header=10000, usecols=[2, 3]);
P = C[:,0]/10000.
E = C[:,1]/10000.
n = len(P)
t = np.arange(0,n,1)
print E

# Provide the model parameters for the soil module
S_rmax = 0.28 # is equal to ±100 milimeters
K_p = 0.009 # is equal to 0.01 m/d
Gamma = 2.0
Imax = 0.0015

R, S = Unsat_Zone.perc(t, P, E, S_rmax, K_p, Gamma, Imax, dt, solver=0)
R1, S1 = Unsat_Zone.perc(t, P, E, S_rmax, K_p, Gamma, Imax, dt, solver=1)

# Bar plot of the precipitation excess ( Takes a while to plot so commented here)
plt.figure(figsize=(17,8))
#plt.subplot(311)
#plt.bar(t, (P-E), lw=0)
#plt.xlim(0, n); plt.ylim(0)
#plt.title('Precipitation excess and the resulting soil state')
#plt.ylabel('precipitation excess [m]') 

# Plot of the resulting Soil State
plt.subplot(312)
plt.plot(t, S)
plt.plot(t,S1)
plt.ylabel('S [m]')
plt.xlabel('Time [d]')
plt.xlim(0, n);
plt.legend(['explicit', 'implicit'])

plt.subplot(313)
plt.plot(t, R)
plt.plot(t,R1)
plt.ylabel('R [m]')
plt.xlabel('Time [d]')
plt.xlim(0, n);
plt.ylim(0,0.005)
plt.legend(['explicit', 'implicit'])

"""
Explanation: 3. RANDOM FORCINGS
End of explanation
"""


#Create an variable to store the percolation functions as the wrapper does not work with function?
X1 = Unsat_Zone.perc
X2 = Unsat_Zone_Python.perc

# Write a wrapper for the Cython file and time it
def wrapper(X1, t, P, E, S_cap, K_sat, Beta, D, dt, solver=1):
    def wrapped():
        return X1(t, P, E, S_cap, K_sat, Beta, D, dt, solver=1)
    return wrapped

wrapped = wrapper(X1,t, P, E, S_cap, K_sat, Beta, D, dt, solver=1)

Cython = timeit.timeit(wrapped, number=1000)

print 'Time taken is', Cython

# Write a wrapper for the Python file and time it
def wrapper(X2, t, P, E, S_cap, K_sat, Beta, D, dt):
    def wrapped():
        return X2(t, P, E, S_cap, K_sat, Beta, D, dt)
    return wrapped

wrapped = wrapper(X2,t, P, E, S_cap, K_sat, Beta, D, dt)

Python = timeit.timeit(wrapped, number=1000)

print 'using Cython File is', (Python/Cython), 'times faster than python to solve the soil module'

"""
Explanation: Compare performance of Cythonized and Python Module
(Non-functioning 29/11/2015)
The computations in this script are computationally demanding and slows down the calibration of a non-linear time series model significantly. Therefore, the python script for the unsaturated zone was ported to the programming C-language using Cython. The Python-script to be cythonized (or in other words: compiled) is Unsat_Zone.pyx using the 
python script setup_unsat_zone.py. Compiling the script on Mac goes as follows:

Open Terminal
Move to the directory where the setup_unsat_zone.py and the Unsat_Zone.pyx files are located
In your terminal, type: 'python setup_unsat_zone.py build_ext --inplace' and press enter

The file will now be compiled to an Unsat_Zone.c and an Unsat_Zone.so file (Shared Object) that can be imported as usual in python. To show the speed-up that is gained by Cythonizing our script, some performance tests are done below. It is shown that the Cythonized file of the python script is approximatly 20-25 times faster than the python script, a significant time saving.
End of explanation
"""


cProfile.runctx("Unsat_Zone.perc(t, P, E, S_cap, K_sat, Beta, D, dt, solver=1)", globals(), locals(), "Profile.prof", sort=-1)

s = pstats.Stats("Profile.prof")
s.strip_dirs().sort_stats("time").print_stats()

"""
Explanation: Speeding it up further
To identify bottlenecks in the computation of the unsaturated zone, profiling can be done. An example script is given below. 
Another method that is provided by Cython is highlighting how 'pythonic' each line of the script is. To do so, perform the following steps:
- Open Terminal
- Move to the directory where your .pyx file is stored
- In terminal, type the command: Cython -a filename.pyx
- In your finder, find the HTML file that is created and open it in your browser
- The darker yellow the lines, the more pythonic that line in your script is. The better yourt script the less yellow :)
End of explanation
"""


n = 30
t = np.arange(0,n,1)
P = np.zeros(n)
P[0] = 10.0 / 1000 # All forcings must be provided in meters
E = 0.00 * np.ones(n)
dt=1

# Provide the model parameters for the soil module
S_rmax = 0.01 # is equal to ±100 milimeters
K_p = 0.001 # is equal to 0.01 m/d
Gamma = 1.0
Imax = 0.0015

R, S = Unsat_Zone.perc(t, P, E, S_rmax, K_p, Gamma, Imax, dt, solver=0)
R1, S1 = Unsat_Zone.perc(t, P, E, S_rmax, K_p, Gamma, Imax, dt, solver=1)

plt.figure(figsize=(15,4))
plt.subplot(121)
plt.plot(t,S)
plt.plot(t,S1)
plt.ylabel('S [meter]')
plt.xlabel('Time [days]')
plt.axhline(S_rmax, color='r')
plt.legend(['explicit', 'implicit', r'$S_rmax$'], loc='best')
plt.title('Soil State over Time')
plt.ylim(0, S_rmax+0.001)

plt.subplot(122)
plt.plot(R)
plt.plot(R1)
plt.title('recharge over time')
plt.ylabel('Recharge [m/d]')
plt.xlabel('Time [days]')
plt.legend(['explicit', 'implicit', r'$S_rmax$'], loc='best')
plt.ylim(0, max(R)+0.0001)

"""
Explanation: Delay in the unsaturated zone?
In this section we will study if the unsaturated zone as modeled has a delaying effect of the arrival time of the peak
End of explanation
"""


