%matplotlib inline
import os
import sys
import glob
import platform
import shutil
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import flopy

print(sys.version)
print('numpy version: {}'.format(np.__version__))
print('matplotlib version: {}'.format(mpl.__version__))
print('pandas version: {}'.format(pd.__version__))
print('flopy version: {}'.format(flopy.__version__))

#Set name of MODFLOW exe
#  assumes executable is in users path statement
exe_name = 'mf2005'
if platform.system() == 'Windows':
    exe_name += '.exe'

path = 'temp'
if not os.path.isdir(path):
    os.mkdirs(path)
gpth = os.path.join('..', 'data', 'mf2005_test', 'UZFtest2.*')
for f in glob.glob(gpth):
    shutil.copy(f, path)

"""
Explanation: FloPy
UZF example
Demonstrates functionality of the flopy UZF module using the example from Niswonger and others (2006). This is the same as the SFR example problem from Prudic and others (2004;
p. 13–19), except the UZF package replaces the ET and RCH packages.
Problem description:

Grid dimensions: 1 Layer, 15 Rows, 10 Columns  
Stress periods: 12  
Units are in seconds and days
Flow package: LPF  
Stress packages: SFR, GHB, UZF 
Solver: SIP  

<img src="./img/Niswonger2006_fig13.png" width="400" height="500"/>
End of explanation
"""


m = flopy.modflow.Modflow.load('UZFtest2.nam', version='mf2005', exe_name=exe_name, 
                               model_ws=path, load_only=['ghb', 'dis', 'bas6', 'oc', 'sip', 'lpf', 'sfr'])

"""
Explanation: Load example dataset, skipping the UZF package
End of explanation
"""


m.external_fnames

rm = [True if '.uz' in f else False for f in m.external_fnames]

m.external_fnames = [f for i, f in enumerate(m.external_fnames) if not rm[i]]
m.external_fnames

m.external_binflag = [f for i, f in enumerate(m.external_binflag) if not rm[i]]
m.external_output = [f for i, f in enumerate(m.external_output) if not rm[i]]
m.external_units = [f for i, f in enumerate(m.external_output) if not rm[i]]

"""
Explanation: remove previous uzf external file references
(so they don't conflict with the ones made by flopy)
End of explanation
"""


fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
modelmap = flopy.plot.ModelMap(model=m)
quadmesh = modelmap.plot_ibound()
linecollection = modelmap.plot_grid()

"""
Explanation: izufbnd array

in the example, the UZF package izufbnd array is the same as the ibound
End of explanation
"""


irnbndpth = os.path.join('..', 'data', 'uzf_examples', 'irunbnd.dat')
irunbnd = np.loadtxt(irnbndpth)

fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
modelmap = flopy.plot.ModelMap(model=m)
irunbndplt = modelmap.plot_array(irunbnd)
plt.colorbar(irunbndplt, ax=ax, label='SFR segment')
linecollection = modelmap.plot_grid()

"""
Explanation: Set up the irunbnd array

read this in from an external file
End of explanation
"""


vksbndpth = os.path.join('..', 'data', 'uzf_examples', 'vks.dat')
vks = np.loadtxt(vksbndpth)

fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
modelmap = flopy.plot.ModelMap(model=m)
vksplt = modelmap.plot_array(vks)
plt.colorbar(vksplt, ax=ax, label='Layer 1 Kv')
linecollection = modelmap.plot_grid()

m2 = flopy.modflow.Modflow.load('UZFtest2.nam', version='mf2005', exe_name='mf2005', 
                               model_ws=path)

"""
Explanation: vks (unsaturated zone vertical hydraulic conductivity) array
End of explanation
"""


finf = np.loadtxt(os.path.join('..', 'data', 'uzf_examples', 'finf.dat'))
finf = np.reshape(finf, (m2.nper, m2.nrow, m2.ncol))
plt.imshow(finf[0], interpolation='none')
plt.colorbar()

plt.plot(m2.dis.perlen.array.cumsum()/864600, 
         [a.mean() * 86400 * 365 * 12 for a in finf[:, :, :]], marker='o')
plt.xlabel('Time, in days')
plt.ylabel('Average infiltration rate, inches per year')

"""
Explanation: finf array

load infiltration rates from a file into a 3D array
finf can be submitted to FloPy as a 3D array, list of 2D arrays, list of numeric values, or single numeric value
End of explanation
"""


extwc = np.loadtxt(os.path.join('..', 'data', 'uzf_examples', 'extwc.dat'))
plt.imshow(extwc, interpolation='none')
plt.colorbar()

"""
Explanation: extwc (extinction water content) array
End of explanation
"""


uzgag = {-68: [-68],
         65: [2, 5, 65, 1], #Print time, head, uz thickness and cum. vols of infiltration, recharge, storage, change in storage and ground-water discharge to land surface.
         66: [5, 2, 66, 2], #Same as option 1 except rates of infiltration, recharge, change in storage, and ground-water discharge also are printed.
         67: [9, 4, 67, 3]} #Prints time, ground-water head, thickness of unsaturated zone, followed by a series of depths and water contents in the unsaturated zone. 

"""
Explanation: Set up the gages (observation points)

supplied as a dictionary keyed by IFTUNIT
A positive value [of IFTUNIT] is for output of individual cells whereas a negative value is for output that is summed over all model cells. 
values are a list of [IUZROW, IUZCOL, IFTUNIT, IUZOPT]
IUZROW and IUZCOL are zero based
End of explanation
"""


uzf = flopy.modflow.ModflowUzf1(m,
                                nuztop=1, iuzfopt=1, irunflg=1, ietflg=1,
                                ipakcb=0, 
                                iuzfcb2=61,# binary output of recharge and groundwater discharge
                                ntrail2=25, nsets=20, nuzgag=4,
                                surfdep=1.0, uzgag=uzgag,
                                iuzfbnd=m2.bas6.ibound.array, 
                                irunbnd=irunbnd, 
                                vks=vks, # saturated vertical hydraulic conductivity of the uz
                                finf=finf, #infiltration rates
                                eps=3.5, # Brooks-Corey relation of water content to hydraulic conductivity (epsilon)
                                thts = 0.35, # saturated water content of the uz in units of volume of water to total volume
                                pet=5.000000E-08, # potential ET
                                extdp=15., # ET extinction depth(s)
                                extwc=extwc, #extinction water content below which ET cannot be removed from the unsaturated zone
                                unitnumber=19)

m.write_input()

"""
Explanation: Make the UZF package
End of explanation
"""


m.run_model()

"""
Explanation: run the model
End of explanation
"""


uzfbdobjct = flopy.utils.CellBudgetFile(os.path.join(path,'UZFtest2.uzfcb2.bin'))
uzfbdobjct.list_records()

r = uzfbdobjct.get_data(text='UZF RECHARGE')
et = uzfbdobjct.get_data(text='GW ET')
plt.imshow(r[6], interpolation='None')

rtot = [rp.sum() for rp in r]
ettot = [etp.sum() for etp in et]
sltot = [sl.sum() for sl in uzfbdobjct.get_data(text='SURFACE LEAKAGE')]

plt.plot(rtot, label='simulated recharge')
plt.plot(np.abs(ettot), label='simulated actual et')
plt.plot(np.abs(sltot), label='simulated surface leakage')
plt.xlabel('Timestep')
plt.ylabel('Volume, in cubic feet')
plt.legend()

"""
Explanation: Results
Look at the budget output
End of explanation
"""


fpth = os.path.join(path, 'UZFtest2.uzf68.out')
dtype = [('TIME', np.float), ('APPLIED-INFIL', np.float), ('RUNOFF', np.float), 
         ('ACTUAL-INFIL', np.float), ('SURFACE-LEAK', np.float), 
         ('UZ-ET', np.float), ('GW-ET', np.float), ('UZSTOR-CHANGE', np.float), 
         ('RECHARGE', np.float)]
# read data from file
df = np.genfromtxt(fpth, skip_header=3, dtype=dtype)
# convert numpy recarray to pandas dataframe
df = pd.DataFrame(data=df)
# set index to the time column
df.set_index(['TIME'], inplace=True)
# plot the data
ax = df.plot(legend=False, figsize=(15, 10))
patches, labels = ax.get_legend_handles_labels()
ax.legend(patches, labels, loc=1)
ax.set_ylabel('Volume for whole model, in cubic feet')

"""
Explanation: Look at the gages
End of explanation
"""


fpth = os.path.join(path, 'UZFtest2.uzf67.out')
data = []
with open(fpth) as input:
    for i in range(3):
        next(input)
    for line in input:
        line = line.strip().split()
        if len(line) == 6:
            layer = int(line.pop(0))
            time = float(line.pop(0))
            head = float(line.pop(0))
            uzthick = float(line.pop(0))
        depth = float(line.pop(0))
        watercontent = float(line.pop(0))
        data.append([layer, time, head, uzthick, depth, watercontent])

df3 = pd.DataFrame(data, columns=['layer', 'time', 'head', 'uzthick', 'depth', 'watercontent'])
df3.head(41)

wc = df3.watercontent.reshape(len(df3.time.unique()), 40).T
wc = pd.DataFrame(wc, columns=df3.time.unique(), index=df3.depth[0:40])
wc.head()

fig, ax = plt.subplots(figsize=(15, 10))
plt.imshow(wc, interpolation='None')
ax.set_aspect(3)
r, c = wc.shape
xcol_locs = np.linspace(0, c-1, 8, dtype=int)
ycol_locs = np.linspace(0, r-1, 5, dtype=int)
ax.set_xticks(xcol_locs)

xlabels = wc.columns
ax.set_xticklabels(xlabels[xcol_locs])
ax.set_ylabel('Depth, in feet')
ax.set_yticks(ycol_locs)
ax.set_yticklabels(wc.index[ycol_locs])
ax.set_xlabel('Time, in seconds')
plt.colorbar(label='Water content')

"""
Explanation: Plot water content profile through time at row 10, column 5
End of explanation
"""


