# import our functionality
import numpy as np
import matplotlib.pyplot as plt
from scipy.special import erfc

kD = 900. # m2/d
S = 0.1
A = 4.0
times = np.linspace(0.1, 3, 10) # d
L = 800. #m
x = np.linspace(0, 800, 33)

for t in times:
    u = np.sqrt(x**2 * S/(4 * kD * t))
    s = A * erfc(u)
    plt.plot(x, s)

plt.xlabel('times [d]')
plt.ylabel('head [m]')
plt.title('Change of groundwater head as func of x, for different times')
plt.show()

"""
Explanation: Exercises chapter 5.3
Pre-exercise
How does the head change at 400 m from the river due to a sudden change of 4 m of the river stage which lasts 2 days.

kD = 900 m2/d
S  = 0.1
A  = 4 m
L = 400 m

Our solution is
$$ s(x,t)  = A \, erfc(u), \,\,\, u = \sqrt{ \frac {x^2 S} {4 kD t} } $$
End of explanation
"""


times = np.linspace(0.1, 3, 100) # d
Q0  = A * np.sqrt( kD * S / (np.pi * times))
plt.xlabel('times [d]')
plt.ylabel('Qlow [m2/d]')
plt.title('Flow at x = 0')
plt.plot(times, Q0)
plt.show()

"""
Explanation: The flux is
$$ Q = -kD \frac {\partial s} {\partial x} = A \sqrt{\frac {kD S} {\pi t} } e^{-u^2} $$
End of explanation
"""


kD = 900. # m2/d
S = 0.1
A = 4.0
times = np.linspace(0.1, 3, 10) # d
L = 800. #m
x = np.linspace(0, 800, 200)
dt = 2 # delay, sudden fall back of river stage

for t in times:
    u1 = np.sqrt(x**2 * S/(4 * kD * t))    
    u2 = np.sqrt(x**2 * S/(4 * kD * (t - dt)))
    s = A * erfc(u1)
    if t> dt:
        s = s -A * erfc(u2)
    plt.plot(x, s, label='t = {:.1f} d'.format(t))

plt.xlabel('distance x [m]')
plt.ylabel('head [m]')
plt.title('Change of groundwater head as func of x, for different times')
plt.legend(fontsize='small')
plt.show()

kD = 900. # m2/d
S = 0.1
A = 4.0
times = np.linspace(0.1, 30, 1000) # d
L = 800. #m
distances = [0.1, 200, 400, 800]
dt = 2 # delay, sudden fall back of river stage

period2 = times>dt

for x in distances:
    u1 = np.sqrt(x**2 * S/(4 * kD * times))    
    u2 = np.sqrt(x**2 * S/(4 * kD * (times[period2] - dt)))
    s = A * erfc(u1)
    s[period2] = s[period2] -A * erfc(u2)
    plt.plot(times, s, label='x = {:.1f} m'.format(x))

plt.ylim(0, 5)
plt.xlabel('timed [d]')
plt.ylabel('head [m]')
plt.title('Change of groundwater head as func of times, for different x')
plt.legend(fontsize='small')
plt.show()

"""
Explanation: How about the situation that after 2 days the river stage falls suddenly back to zero??
End of explanation
"""


