%matplotlib inline
import os
import numpy as np
import matplotlib.pyplot as plt
import flopy
from flopy.utils.gridgen import Gridgen

"""
Explanation: FloPy Creating Layered Quadtree Grids with GRIDGEN
FloPy has a module that can be used to drive the GRIDGEN program.  This notebook shows how it works.
The Flopy GRIDGEN module requires that the gridgen executable can be called using subprocess.
End of explanation
"""


Lx = 100.
Ly = 100.
nlay = 2
nrow = 51
ncol = 51
delr = Lx / ncol
delc = Ly / nrow
h0 = 10
h1 = 5
top = h0
botm = np.zeros((nlay, nrow, ncol), dtype=np.float32)
botm[1, :, :] = -10.

ms = flopy.modflow.Modflow(rotation=20.)
dis = flopy.modflow.ModflowDis(ms, nlay=nlay, nrow=nrow, ncol=ncol, delr=delr,
                               delc=delc, top=top, botm=botm)

"""
Explanation: Setup Base MODFLOW Grid
GRIDGEN works off of a base MODFLOW grid.  The following information defines the basegrid.
End of explanation
"""


model_ws = os.path.join('.', 'data')
g = Gridgen(dis, model_ws=model_ws)

"""
Explanation: Create the Gridgen Object
End of explanation
"""


# setup the active domain
adshp = os.path.join(model_ws, 'ad0')
adpoly = [[[(0, 0), (0, 60), (40, 80), (60, 0), (0, 0)]]]
# g.add_active_domain(adpoly, range(nlay))

"""
Explanation: Add an Optional Active Domain
Cells outside of the active domain will be clipped and not numbered as part of the final grid.  If this step is not performed, then all cells will be included in the final grid.
End of explanation
"""


x = Lx * np.random.random(10)
y = Ly * np.random.random(10)
wells = zip(x, y)
g.add_refinement_features(wells, 'point', 3, range(nlay))

river = [[[(-20, 10), (60, 60)]]]
g.add_refinement_features(river, 'line', 3, range(nlay))

g.add_refinement_features(adpoly, 'polygon', 1, range(nlay))

"""
Explanation: Refine the Grid
End of explanation
"""


g.build()

"""
Explanation: Build the Grid
End of explanation
"""


fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
g.plot(ax, linewidth=0.5)
flopy.plot.plot_shapefile(adshp, ax=ax, facecolor='yellow', alpha=0.1)

"""
Explanation: Plot the Grid
End of explanation
"""


mu = flopy.modflow.Modflow(model_ws=model_ws, modelname='mfusg')
disu = g.get_disu(mu)
# disu.write_file()
# print(disu)

"""
Explanation: Create a Flopy ModflowDisu Object
End of explanation
"""


adpoly_intersect = g.intersect(adpoly, 'polygon', 0)
print(adpoly_intersect.dtype.names)
print(adpoly_intersect)
print(adpoly_intersect.nodenumber)

well_intersect = g.intersect(wells, 'point', 0)
print(well_intersect.dtype.names)
print(well_intersect)
print(well_intersect.nodenumber)

river_intersect = g.intersect(river, 'line', 0)
print(river_intersect.dtype.names)
# print(river_intersect)
# print(river_intersect.nodenumber)

"""
Explanation: Intersect Features with the Grid
End of explanation
"""


a = np.zeros((g.nodes), dtype=np.int)
a[adpoly_intersect.nodenumber] = 1
a[well_intersect.nodenumber] = 2
a[river_intersect.nodenumber] = 3
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
g.plot(ax, a=a, masked_values=[0], edgecolor='none', cmap='jet')

"""
Explanation: Plot Intersected Features
End of explanation
"""


