%matplotlib inline
import sys
import shutil
import os
import glob
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
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__))

"""
Explanation: FloPy
MODPATH example
This notebook demonstrates how to create and run forward and backward tracking with MODPATH. The notebooks also shows how to create subsets of pathline and endpoint information, plot MODPATH results on ModelMap objects, and export endpoints and pathlines as shapefiles.
End of explanation
"""


mffiles = glob.glob('../data/mp6/EXAMPLE.*')
for f in mffiles:
    shutil.copy(f, 'temp/' + os.path.split(f)[1])

"""
Explanation: Copy modflow datasets to scratch directory
End of explanation
"""


model_ws = 'temp/'
m = flopy.modflow.Modflow.load('EXAMPLE.nam', model_ws=model_ws)
m.get_package_list()

nrow, ncol, nlay, nper = m.nrow_ncol_nlay_nper
nrow, ncol, nlay, nper 

m.dis.steady.array

m.write_input()

hdsfile = flopy.utils.HeadFile(model_ws + 'EXAMPLE.HED')
hdsfile.get_kstpkper()

hds = hdsfile.get_data(kstpkper=(0, 2))

"""
Explanation: Load MODFLOW model
End of explanation
"""


plt.imshow(hds[4, :, :])
plt.colorbar();

fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
modelmap = flopy.plot.ModelMap(model=m, layer=4)
quadmesh = modelmap.plot_ibound()
linecollection = modelmap.plot_grid()
riv = modelmap.plot_bc('RIV', color='g', plotAll=True)
quadmesh = modelmap.plot_bc('WEL', kper=1, plotAll=True)
contour_set = modelmap.contour_array(hds, 
                                     levels=np.arange(np.min(hds),np.max(hds),0.5), colors='b')
plt.clabel(contour_set, inline=1, fontsize=14)

"""
Explanation: Plot RIV bc and head results
End of explanation
"""


mp = flopy.modpath.Modpath(modelname='ex6',
                           exe_name='mp6',
                           modflowmodel=m,
                           model_ws='temp',
                           dis_file=m.name+'.dis',
                           head_file=m.name+'.hed',
                           budget_file=m.name+'.bud')

mpb = flopy.modpath.ModpathBas(mp, hdry=m.lpf.hdry, laytyp=m.lpf.laytyp, ibound=1, prsity=0.1)

# start the particles at begining of per 3, step 1, as in example 3 in MODPATH6 manual
# (otherwise particles will all go to river)
sim = mp.create_mpsim(trackdir='forward', simtype='pathline', packages='RCH', start_time=(2, 0, 1.))
mp.write_input()

mp.run_model(silent=False)

"""
Explanation: Create forward particle tracking simulation where particles are released at the top of each cell in layer 1

specifying the recharge package in create_mpsim releases a single particle on iface=6 of each top cell  
start the particles at begining of per 3, step 1, as in example 3 in MODPATH6 manual
End of explanation
"""


epobj = flopy.utils.EndpointFile('temp/ex6.mpend')
well_epd = epobj.get_destination_endpoint_data(dest_cells=[(4, 12, 12)]) 
# returns record array of same form as epobj.get_all_data()

well_epd[0:2]

fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
modelmap = flopy.plot.ModelMap(model=m, layer=2)
quadmesh = modelmap.plot_ibound()
linecollection = modelmap.plot_grid()
riv = modelmap.plot_bc('RIV', color='g', plotAll=True)
quadmesh = modelmap.plot_bc('WEL', kper=1, plotAll=True)
contour_set = modelmap.contour_array(hds, 
                                     levels=np.arange(np.min(hds),np.max(hds),0.5), colors='b')
plt.clabel(contour_set, inline=1, fontsize=14)
modelmap.plot_endpoint(well_epd, direction='starting', colorbar=True)

"""
Explanation: Read in the endpoint file and plot particles that terminated in the well
End of explanation
"""


epobj.write_shapefile(well_epd, direction='starting', shpname='temp/starting_locs.shp', sr=m.sr)

"""
Explanation: Write starting locations to a shapefile
End of explanation
"""


# make a selection of cells that terminate in the well cell = (4, 12, 12)
pthobj = flopy.utils.PathlineFile('temp/ex6.mppth')
well_pathlines = pthobj.get_destination_pathline_data(dest_cells=[(4, 12, 12)])
well_pathlines # this is a record array of same form as pthobj._data

"""
Explanation: Read in the pathline file and subset to pathlines that terminated in the well
End of explanation
"""


fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
modelmap = flopy.plot.ModelMap(model=m, layer=2)
quadmesh = modelmap.plot_ibound()
linecollection = modelmap.plot_grid()
riv = modelmap.plot_bc('RIV', color='g', plotAll=True)
quadmesh = modelmap.plot_bc('WEL', kper=1, plotAll=True)
contour_set = modelmap.contour_array(hds, 
                                     levels=np.arange(np.min(hds),np.max(hds),0.5), colors='b')
plt.clabel(contour_set, inline=1, fontsize=14)

modelmap.plot_endpoint(well_epd, direction='starting', colorbar=True)
# for now, each particle must be plotted individually 
#(plot_pathline() will plot a single line for recarray with multiple particles)
for pid in np.unique(well_pathlines.particleid):
    modelmap.plot_pathline(pthobj.get_data(pid), layer='all', colors='red');

"""
Explanation: Plot the pathlines that terminate in the well and the starting locations of the particles
End of explanation
"""


# one line feature per particle
pthobj.write_shapefile(well_pathlines,
                       direction='starting', shpname='temp/pathlines.shp')

# one line feature for each row in pathline file 
# (can be used to color lines by time or layer in a GIS)
pthobj.write_shapefile(well_pathlines, one_per_particle=False, 
                       shpname='temp/pathlines_1per.shp')

"""
Explanation: Write pathlines to a shapefile
End of explanation
"""


model_ws = 'temp/'
m2 = flopy.modflow.Modflow.load('EXAMPLE.nam', model_ws=model_ws, exe_name='mf2005')
m2.get_package_list()

m2.nrow_ncol_nlay_nper

m2.wel.stress_period_data.data

node_data = np.array([(3, 12, 12, 'well1', 'skin', -1, 0, 0, 0, 1., 2., 5., 6.2),
                      (4, 12, 12, 'well1', 'skin', -1, 0, 0, 0, 0.5, 2., 5., 6.2)], 
                     dtype=[('k', np.int), ('i', np.int), ('j', np.int), 
                            ('wellid', np.object), ('losstype', np.object), 
                            ('pumploc', np.int), ('qlimit', np.int), 
                            ('ppflag', np.int), ('pumpcap', np.int), 
                            ('rw', np.float), ('rskin', np.float), 
                            ('kskin', np.float), ('zpump', np.float)]).view(np.recarray)

stress_period_data = {0: np.array([(0, 'well1', -150000.0)], dtype=[('per', np.int), ('wellid', np.object), 
                                                            ('qdes', np.float)])}

m2.name = 'Example_mnw'
m2.remove_package('WEL')
mnw2 = flopy.modflow.ModflowMnw2(model=m2, mnwmax=1,
                                 node_data=node_data, 
                                 stress_period_data=stress_period_data, 
                                 itmp=[1, -1, -1])
m2.get_package_list()

"""
Explanation: Replace WEL package with MNW2; create backward tracking simulation using particles released at MNW well
End of explanation
"""


m2.write_input()

m2.run_model(silent=False)

"""
Explanation: Write and run MODFLOW
End of explanation
"""


mp = flopy.modpath.Modpath(modelname='ex6mnw',
                           exe_name='mp6',
                           modflowmodel=m2,
                           model_ws='temp',
                           dis_file=m.name+'.dis',
                           head_file=m.name+'.hds',
                           budget_file=m.name+'.cbc')

mpb = flopy.modpath.ModpathBas(mp, hdry=m2.lpf.hdry, laytyp=m2.lpf.laytyp, ibound=1, prsity=0.1)
sim = mp.create_mpsim(trackdir='backward', simtype='pathline', packages='MNW2')

mp.write_input()

mp.run_model(silent=False)

"""
Explanation: Create new MODPATH object
End of explanation
"""


pthobj = flopy.utils.PathlineFile('temp/ex6mnw.mppth')
epdobj = flopy.utils.EndpointFile('temp/ex6mnw.mpend')
well_epd = epdobj.get_alldata()
well_pathlines = pthobj.get_alldata() # returns a list of recarrays; one per pathline

fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
modelmap = flopy.plot.ModelMap(model=m2, layer=2)
quadmesh = modelmap.plot_ibound()
linecollection = modelmap.plot_grid()
riv = modelmap.plot_bc('RIV', color='g', plotAll=True)
quadmesh = modelmap.plot_bc('MNW2', kper=1, plotAll=True)
contour_set = modelmap.contour_array(hds, 
                                     levels=np.arange(np.min(hds),np.max(hds),0.5), colors='b')
plt.clabel(contour_set, inline=1, fontsize=14)

modelmap.plot_pathline(well_pathlines, travel_time='<10000',
                       layer='all', colors='red');

"""
Explanation: Read in results and plot
End of explanation
"""


