% matplotlib inline
import sys
sys.path.append('/Users/aleaf/Documents/GitHub/flopy3/')
import platform
import os
import numpy as np
import glob
import shutil
import matplotlib as mpl
import matplotlib.pyplot as plt
import flopy
import flopy.utils.binaryfile as bf
from flopy.utils.sfroutputfile import SfrFile

mpl.rcParams['figure.figsize'] = (11, 8.5)

print(sys.version)
print('numpy version: {}'.format(np.__version__))
print('matplotlib version: {}'.format(mpl.__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'

"""
Explanation: SFR package example
Demonstrates functionality of Flopy SFR module using the example documented by Prudic and others (2004):  
Problem description:

Grid dimensions: 1 Layer, 15 Rows, 10 Columns  
Stress periods: 1 steady  
Flow package: LPF  
Stress packages: SFR, GHB, EVT, RCH  
Solver: SIP  

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


path = 'data'
if os.path.isfile(path):
    os.remove(path)
elif os.path.isdir(path):
    shutil.rmtree(path)
os.mkdir(path)
gpth = os.path.join('..', 'data', 'mf2005_test', 'test1ss.*')
for f in glob.glob(gpth):
    shutil.copy(f, path)
gpth = os.path.join('..', 'data', 'mf2005_test', 'test1tr.*')
for f in glob.glob(gpth):
    shutil.copy(f, path)

"""
Explanation: copy over the example files to the working directory
End of explanation
"""


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

oc = m.oc
oc.stress_period_data

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


rpth = os.path.join('..', 'data', 'sfr_examples', 'test1ss_reach_data.csv')
reach_data = np.genfromtxt(rpth, delimiter=',', names=True)
reach_data

"""
Explanation: Read pre-prepared reach and segment data into numpy recarrays using numpy.genfromtxt()
Reach data (Item 2 in the SFR input instructions), are input and stored in a numpy record array
http://docs.scipy.org/doc/numpy/reference/generated/numpy.recarray.html
This allows for reach data to be indexed by their variable names, as described in the SFR input instructions.
For more information on Item 2, see the Online Guide to MODFLOW:
http://water.usgs.gov/nrp/gwsoftware/modflow2000/MFDOC/index.html?sfr.htm
End of explanation
"""


spth = os.path.join('..', 'data', 'sfr_examples', 'test1ss_segment_data.csv')
ss_segment_data = np.genfromtxt(spth, delimiter=',', names=True)
segment_data = {0: ss_segment_data}
segment_data[0][0:1]['width1']

"""
Explanation: Segment Data structure
Segment data are input and stored in a dictionary of record arrays, which
End of explanation
"""


channel_flow_data = {0: {1: [[0.5, 1.0, 2.0, 4.0, 7.0, 10.0, 20.0, 30.0, 50.0, 75.0, 100.0],
                             [0.25, 0.4, 0.55, 0.7, 0.8, 0.9, 1.1, 1.25, 1.4, 1.7, 2.6],
                             [3.0, 3.5, 4.2, 5.3, 7.0, 8.5, 12.0, 14.0, 17.0, 20.0, 22.0]]}}

"""
Explanation: define dataset 6e (channel flow data) for segment 1
dataset 6e is stored in a nested dictionary keyed by stress period and segment,
with a list of the following lists defined for each segment with icalc == 4
FLOWTAB(1) FLOWTAB(2) ... FLOWTAB(NSTRPTS)
DPTHTAB(1) DPTHTAB(2) ... DPTHTAB(NSTRPTS)
WDTHTAB(1) WDTHTAB(2) ... WDTHTAB(NSTRPTS)
End of explanation
"""


channel_geometry_data = {0: {7: [[0.0, 10.0, 80.0, 100.0, 150.0, 170.0, 240.0, 250.0],
                                 [20.0, 13.0, 10.0, 2.0, 0.0, 10.0, 13.0, 20.0]],
                             8: [[0.0, 10.0, 80.0, 100.0, 150.0, 170.0, 240.0, 250.0],
                                 [25.0, 17.0, 13.0, 4.0, 0.0, 10.0, 16.0, 20.0]]}}

"""
Explanation: define dataset 6d (channel geometry data) for segments 7 and 8
dataset 6d is stored in a nested dictionary keyed by stress period and segment,
with a list of the following lists defined for each segment with icalc == 4
FLOWTAB(1) FLOWTAB(2) ... FLOWTAB(NSTRPTS)
DPTHTAB(1) DPTHTAB(2) ... DPTHTAB(NSTRPTS)
WDTHTAB(1) WDTHTAB(2) ... WDTHTAB(NSTRPTS)
End of explanation
"""


nstrm = len(reach_data) # number of reaches
nss = len(segment_data[0]) # number of segments
nsfrpar = 0 # number of parameters (not supported)
nparseg = 0
const = 1.486 # constant for manning's equation, units of cfs
dleak = 0.0001 # closure tolerance for stream stage computation
ipakcb = 53 # flag for writing SFR output to cell-by-cell budget (on unit 53)
istcb2 = 81 # flag for writing SFR output to text file
dataset_5 = {0: [nss, 0, 0]} # dataset 5 (see online guide)

"""
Explanation: Define SFR package variables
End of explanation
"""


sfr = flopy.modflow.ModflowSfr2(m, nstrm=nstrm, nss=nss, const=const, dleak=dleak, ipakcb=ipakcb, istcb2=istcb2, 
                                reach_data=reach_data,
                                segment_data=segment_data,
                                channel_geometry_data=channel_geometry_data,
                                channel_flow_data=channel_flow_data,
                                dataset_5=dataset_5, unit_number=15)

sfr.reach_data[0:1]

"""
Explanation: Instantiate SFR package
Input arguments generally follow the variable names defined in the Online Guide to MODFLOW
End of explanation
"""


sfr.plot(key='iseg');

"""
Explanation: Plot the SFR segments
any column in the reach_data array can be plotted using the key argument
End of explanation
"""


chk = sfr.check()

m.external_fnames = [os.path.split(f)[1] for f in m.external_fnames]
m.external_fnames

m.write_input()

#m.run_model()

"""
Explanation: Check the SFR dataset for errors
End of explanation
"""


sfr_outfile = os.path.join('..', 'data', 'sfr_examples', 'test1ss.flw')
sfrout = SfrFile(sfr_outfile)
df = sfrout.get_dataframe()
df.head()

"""
Explanation: Load SFR formated water balance output into pandas dataframe using the SfrFile class

requires the pandas library
End of explanation
"""


inds = df.segment == 3
ax = df.ix[inds, ['Qin', 'Qaquifer', 'Qout']].plot(x=df.reach[inds])
ax.set_ylabel('Flow, in cubic feet per second')
ax.set_xlabel('SFR reach')

"""
Explanation: Plot streamflow and stream/aquifer interactions for a segment
End of explanation
"""


streambed_top = m.sfr.segment_data[0][m.sfr.segment_data[0].nseg == 3][['elevup', 'elevdn']][0]
streambed_top

df['model_top'] = m.dis.top.array[df.row.values - 1, df.column.values -1]
fig, ax = plt.subplots()
plt.plot([1, 6], list(streambed_top), label='streambed top')
ax = df.ix[inds, ['stage', 'model_top']].plot(ax=ax, x=df.reach[inds])
ax.set_ylabel('Elevation, in feet')
plt.legend()

"""
Explanation: Look at stage, model top, and streambed top
End of explanation
"""


bpth = os.path.join('data', 'test1ss.cbc')
cbbobj = bf.CellBudgetFile(bpth)
cbbobj.list_records()

sfrleak = cbbobj.get_data(text='  STREAM LEAKAGE')[0]
sfrleak[sfrleak == 0] = np.nan # remove zero values

"""
Explanation: Get SFR leakage results from cell budget file
End of explanation
"""


im = plt.imshow(sfrleak[0], interpolation='none', cmap='coolwarm', vmin = -3, vmax=3)
cb = plt.colorbar(im, label='SFR Leakage, in cubic feet per second');

"""
Explanation: Plot leakage in plan view
End of explanation
"""


sfrQ = sfrleak[0].copy()
sfrQ[sfrQ == 0] = np.nan
sfrQ[df.row.values-1, df.column.values-1] = df[['Qin', 'Qout']].mean(axis=1).values
im = plt.imshow(sfrQ, interpolation='none')
plt.colorbar(im, label='Streamflow, in cubic feet per second');

"""
Explanation: Plot total streamflow
End of explanation
"""


sfrout_tr = SfrFile('data/test1tr.flw')
dftr = sfrout_tr.get_dataframe()
dftr.head()

"""
Explanation: Reading transient SFR formatted output
the SfrFile class handles this the same way
files for the transient version of the above example were already copied to the data folder in the third cell above
first run the transient model to get the output:
```

mf2005 test1tr.nam
```
End of explanation
"""


fig, axes = plt.subplots(2, 1, sharex=True)
dftr8 = dftr.loc[(dftr.segment == 8) & (dftr.reach == 5)]
dftr8.Qout.plot(ax=axes[0])
axes[0].set_ylabel('Simulated streamflow, cfs')
dftr8.Qaquifer.plot(ax=axes[1])
axes[1].set_ylabel('Leakage to aquifer, cfs')

"""
Explanation: plot a hydrograph
plot Qout (simulated streamflow) and Qaquifer (simulated stream leakage) through time
End of explanation
"""


