%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from ttim import *
import pandas as pd

"""
Explanation: Test for anisotropic water-table aquifer
This test is taken from examples presented in MLU tutorial.
End of explanation
"""


b = 10 #aquifer thickness in m
Q = 172.8 #constant discharge rate in m^3/d
rw = 0.1 #well radius in m
rc = 0.1 #casing radius in m

"""
Explanation: Set basic parameters:
End of explanation
"""


r1 = 3.16 
r2 = 31.6
data0 = np.loadtxt('data/moench_pumped.txt', skiprows=1)
t0 = data0[:, 0] / 60 / 60 / 24 #convert time from seconds to days
h0 = -data0[:, 1]
data1 = np.loadtxt('data/moench_ps1.txt', skiprows=1)
t1 = data1[:, 0] / 60 / 60 / 24 #convert time from seconds to days
h1 = -data1[:, 1]
data2 = np.loadtxt('data/moench_pd1.txt', skiprows=1)
t2 = data2[:, 0] / 60 / 60 / 24 #convert time from seconds to days
h2 = -data2[:, 1]
data3 = np.loadtxt('data/moench_ps2.txt', skiprows=1)
t3 = data3[:, 0] / 60 / 60 / 24 #convert time from seconds to days
h3 = -data3[:, 1]
data4 = np.loadtxt('data/moench_pd2.txt', skiprows=1)
t4 = data4[:, 0] / 60 / 60 / 24 #convert time from seconds to days
h4 = -data4[:, 1]

"""
Explanation: Load datasets of observation wells:
End of explanation
"""


#Set kaq, Saq, Sy and kzoverkh as given in Moench (1997)
kaq = 1e-4 * 60 * 60 * 24 #convert from m/s to m/d
Sy = 0.2
Saq = 2e-5
zh = 0.5 #kzoverkh

ml1 = Model3D(kaq=kaq, z=[0, -0.1, -2.1, -5.1, -10.1], Saq=[Sy, Saq, Saq, Saq], \
              kzoverkh=zh, tmin=1e-5, tmax=3)
w1 = Well(ml1, xw=0, yw=0, rw=rw, rc=rc, tsandQ=[(0, Q)], layers=3)
ml1.solve()

hm1 = ml1.head(r1, 0, t1, layers=1)[0]
hm2 = ml1.head(r1, 0, t2, layers=3)[0]
hm3 = ml1.head(r2, 0, t3, layers=1)[0]
hm4 = ml1.head(r2, 0, t4, layers=3)[0]
hm0 = ml1.head(0, 0, t0, layers=3)[0]
plt.figure(figsize=(8, 5))
plt.loglog(t0, -h0, '.', label='pumped well')
plt.loglog(t0, -hm0, label='ttim pumped well')
plt.loglog(t1, -h1, '.', label='PS1')
plt.loglog(t1, -hm1, label='ttim PS1')
plt.loglog(t2, -h2, '.', label='PD1')
plt.loglog(t2, -hm2, label='ttim PD1')
plt.loglog(t3, -h3, '.', label='PS2')
plt.loglog(t3, -hm3, label='ttim PS2')
plt.loglog(t4, -h4, '.', label='PD2')
plt.loglog(t4, -hm4, label='ttim PD2')
plt.legend();

res1 = 0
res2 = 0
res3 = 0
res4 = 0
res0 = 0
for i in range(len(h1)):
    r = (h1[i] - hm1[i]) ** 2
    res1 = res1 + r
for i in range(len(h2)):
    r = (h2[i] - hm2[i]) ** 2
    res2 = res2 + r
for i in range(len(h3)):
    r = (h3[i] - hm3[i]) ** 2
    res3 = res3 + r
for i in range(len(h4)):
    r = (h4[i] - hm4[i]) ** 2
    res4 = res4 + r
for i in range(len(h0)):
    r = (h0[i] - hm0[i]) ** 2
    res0 = res0 + r
    
n = len(h1) + len(h2) + len(h3) + len(h4) + len(h0)
residuals = res1 + res2 + res3 + res4 + res0
rmse = np.sqrt(residuals/n)
print('RMSE:', rmse)

"""
Explanation: Check how well TTim can simulate drawdowns in a vertically anisotropic water-table aquifer:
End of explanation
"""


ml2 = Model3D(kaq=1, z=[0, -0.1, -2.1, -5.1, -10.1], Saq=[0.1, 1e-4, 1e-4, 1e-4], \
              kzoverkh=1, tmin=1e-5, tmax=3)
w2 = Well(ml2, xw=0, yw=0, rw=rw, rc=rc, tsandQ=[(0, Q)], layers=3)
ml2.solve()

ca2 = Calibrate(ml2)
ca2.set_parameter(name='kaq0_3', initial=1)
ca2.set_parameter(name='Saq0', initial=0.2)
ca2.set_parameter(name='Saq1_3', initial=1e-4, pmin=0)
ca2.set_parameter_by_reference(name='kzoverkh', parameter=ml2.aq.kzoverkh, \
                               initial=0.1, pmin=0)
ca2.series(name='pumped', x=0, y=0, t=t0, h=h0, layer=3)
ca2.series(name='PS1', x=r1, y=0, t=t1, h=h1, layer=1)
ca2.series(name='PD1', x=r1, y=0, t=t2, h=h2, layer=3)
ca2.series(name='PS2', x=r2, y=0, t=t3, h=h3, layer=1)
ca2.series(name='PD2', x=r2, y=0, t=t4, h=h4, layer=3)
ca2.fit()

display(ca2.parameters)
print('RMSE:', ca2.rmse())

hm0_2 = ml2.head(0, 0, t0, layers=3)[0]
hm1_2 = ml2.head(r1, 0, t1, layers=1)[0]
hm2_2 = ml2.head(r1, 0, t2, layers=3)[0]
hm3_2 = ml2.head(r2, 0, t3, layers=1)[0]
hm4_2 = ml2.head(r2, 0, t4, layers=3)[0]
plt.figure(figsize=(8, 5))
plt.semilogx(t0, h0, '.', label='pumped')
plt.semilogx(t0, hm0_2, label='ttim pumped')
plt.semilogx(t1, h1, '.', label='PS1')
plt.semilogx(t1, hm1_2, label='ttim PS1')
plt.semilogx(t2, h2, '.', label='PD1')
plt.semilogx(t2, hm2_2, label='ttim PD1')
plt.semilogx(t3, h3, ',', label='PS2')
plt.semilogx(t3, hm3_2, label='ttim PS2')
plt.semilogx(t4, h4, '.', label='PD2')
plt.semilogx(t4, hm4_2, label='ttim PD2')
plt.legend();

"""
Explanation: Try calibrating model to find the parameters:
End of explanation
"""


ml3 = Model3D(kaq=1, z=[0, -0.1, -2.1, -5.1, -10.1], Saq=[0.1, 1e-4, 1e-4, 1e-4], \
              kzoverkh=1, tmin=1e-5, tmax=3)
w3 = Well(ml3, xw=0, yw=0, rw=rw, rc=rc, tsandQ=[(0, Q)], layers=3)
ml3.solve()

ca3 = Calibrate(ml3)
ca3.set_parameter(name='kaq0', initial=1, pmin=0)
ca3.set_parameter(name='kaq1_3', initial=1)
ca3.set_parameter(name='Saq0', initial=0.2, pmin=0)
ca3.set_parameter(name='Saq1_3', initial=1e-4, pmin=0)
ca3.set_parameter_by_reference(name='kzoverkh', parameter=ml3.aq.kzoverkh, \
                               initial=0.1, pmin=0)
ca3.series(name='pumped', x=0, y=0, t=t0, h=h0, layer=3)
ca3.series(name='PS1', x=r1, y=0, t=t1, h=h1, layer=1)
ca3.series(name='PD1', x=r1, y=0, t=t2, h=h2, layer=3)
ca3.series(name='PS2', x=r2, y=0, t=t3, h=h3, layer=1)
ca3.series(name='PD2', x=r2, y=0, t=t4, h=h4, layer=3)
ca3.fit()

display(ca3.parameters)
print('RMSE:', ca3.rmse())

hm0_3 = ml3.head(0, 0, t0, layers=3)[0]
hm1_3 = ml3.head(r1, 0, t1, layers=1)[0]
hm2_3 = ml3.head(r1, 0, t2, layers=3)[0]
hm3_3 = ml3.head(r2, 0, t3, layers=1)[0]
hm4_3 = ml3.head(r2, 0, t4, layers=3)[0]
plt.figure(figsize=(8, 5))
plt.semilogx(t0, h0, '.', label='pumped')
plt.semilogx(t0, hm0_3, label='ttim pumped')
plt.semilogx(t1, h1, '.', label='PS1')
plt.semilogx(t1, hm1_3, label='ttim PS1')
plt.semilogx(t2, h2, '.', label='PD1')
plt.semilogx(t2, hm2_3, label='ttim PD1')
plt.semilogx(t3, h3, ',', label='PS2')
plt.semilogx(t3, hm3_3, label='ttim PS2')
plt.semilogx(t4, h4, '.', label='PD2')
plt.semilogx(t4, hm4_3, label='ttim PD2');

"""
Explanation: Try calibrating model with stratified kaq:
End of explanation
"""


ca3.parameters['optimal'].values

ta = pd.DataFrame(columns=['Moench', 'TTim', 'TTim-stratified'],\
                 index=['k0[m/d]', 'k[m/d]', 'Sy[-]', 'Ss[1/m]', 'kz/kh'])
ta.loc[:, 'TTim-stratified'] = ca3.parameters['optimal'].values
ta.loc[1:, 'TTim'] = ca2.parameters['optimal'].values
ta.loc[1:, 'Moench'] = [8.640, 0.2, 2e-5, 0.5]
ta.loc['RMSE'] = [0.061318, ca2.rmse(), ca3.rmse()]
ta

"""
Explanation: Summary of calibrated values
End of explanation
"""


