'''my_second_flopy_example.py: well modeling added
'''

import numpy as np
import flopy

# Assign name and create modflow model object
modelname = 'ex2'
mf = flopy.modflow.Modflow(modelname, exe_name='mf2005')

# Model domain and grid definition
Lx = 1000.
Ly = 1000.
ztop = 0.
zbot = -50.
nlay = 1
nrow = 10
ncol = 10
delr = Lx/ncol # spacings along a row, can be an array
delc = Ly/nrow # spacings along a column, can be an array
delv = (ztop - zbot) / nlay
botm = np.linspace(ztop, zbot, nlay + 1)

# Create the discretization object
dis = flopy.modflow.ModflowDis(mf, nlay, nrow, ncol, delr=delr, delc=delc,
                               top=ztop, botm=botm[1:])

# Variables for the BAS package
# active > 0, inactive = 0, or constant head < 0
ibound = np.ones((nlay, nrow, ncol), dtype=np.int32)
ibound[:, :, 0] = -1
ibound[:, :, -1] = -1
# intial head value also serves as boundary conditions
strt = np.ones((nlay, nrow, ncol), dtype=np.float32)
strt[:, :, 0] =  10.
strt[:, :, -1] = 0.
bas = flopy.modflow.ModflowBas(mf, ibound=ibound, strt=strt)

# Add LPF package to the MODFLOW model
# hk array of horizontal hydraulic conductivity. vka array of vertical hydraulic conductivity, ipakcb file number writing for cell-by-cell budget(need to be defined for the current version)
lpf = flopy.modflow.ModflowLpf(mf, hk=10., vka=10., ipakcb=53)

# Add the well package
# Remember to use zero-based layer, row, column indices!
pumping_rate = -50.
wcol = nrow/2 - 1 # x index for the well
wrow = ncol/2 - 1 # y index for the well
wel_sp = [[0, wrow, wcol, pumping_rate]] # lay, row, col index, pumping rate
stress_period_data = {0: wel_sp} # define well stress period {period, well info dictionary}
wel = flopy.modflow.ModflowWel(mf, stress_period_data=stress_period_data)

# Add OC package to the MODFLOW model
spd = {(0, 0): ['print head', 'print budget', 'save head', 'save budget']}
oc = flopy.modflow.ModflowOc(mf, stress_period_data=spd, compact=True)

# Add PCG package to the MODFLOW model
pcg = flopy.modflow.ModflowPcg(mf)

# Write the MODFLOW model input files
mf.write_input()

# Run the MODFLOW model
success, buff = mf.run_model()

# Post process the results
import matplotlib.pyplot as plt
import flopy.utils.binaryfile as bf
fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(1, 1, 1, aspect='equal')

# FV grid extent: (delr/2., Lx - delr/2., delc/2., Ly - delc/2.)
wpt = ((wcol+0.5)*delr, Lx - ((wrow + 0.5)*delc)) # origing at low upper..

hds = bf.HeadFile(modelname+'.hds')
times = hds.get_times() # simulation time, steady state
head = hds.get_data(totim=times[-1])

cbb = bf.CellBudgetFile(modelname+'.cbc') # read budget file
#kstpkper_list = cbb.get_kstpkper()
# cbb.textlist to get a list of data texts
frf = cbb.get_data(text='FLOW RIGHT FACE', totim=times[-1])[0]
fff = cbb.get_data(text='FLOW FRONT FACE', totim=times[-1])[0]

# flopy plot object
modelmap = flopy.plot.ModelMap(model=mf, layer=0)
# plot grid
lc = modelmap.plot_grid() # grid
# plot contour
cs = modelmap.contour_array(head, levels=np.linspace(0, 10, 21)) # head contour
plt.clabel(cs, fontsize=20, fmt='%1.1f', zorder=1) # contour label
# plot discharge quiver
quiver = modelmap.plot_discharge(frf, fff, head=head)
# plot well location
plt.plot(wpt[0],wpt[1],'ro')
plt.show()