import flopy
import numpy as np

modelname = 'test'
mf = flopy.modflow.Modflow(modelname = modelname, exe_name = 'mf2005')

q = 50
Lx = 500
Ly = 500
ncol = 50
nrow = 50
nlay = 1
ztop = 0.0
zbot = -50
botm = np.linspace(ztop, zbot, nlay + 1)
delr = Ly / nrow
delc = Lx / ncol
dis = flopy.modflow.ModflowDis(mf,
                               nlay = nlay,
                               nrow = nrow,
                               ncol = ncol,
                               delc = delc,
                               delr = delr,
                               ztop = ztop,
                               botm = botm[1:])

ibound = np.ones((nlay, nrow, ncol))
ibound[:,0,:] = 0
ibound[:,-1,:] = 0
ibound[:, :, 0] = -1

strt = 10 * np.ones((nlay, nrow, ncol))
strt[:, :, 0] = 0.0
strt[:, :, -1] = 10.0
bas = flopy.modflow.ModflowBas(mf, ibound = ibound, strt=strt)

lpf = flopy.modflow.ModflowLpf(mf,
                               hk = 50,
                               vka = 1,
                               sy = 0.1,
                               ss = 1.0e-4,
                               laytyp = 1)

spd = []
for il in range(nlay):
    for ir in range(nrow):
        spd.append([il, ir, 0, 0, q])

fhb = flopy.modflow.ModflowFhb(mf, ds5 = spd)

pcg = flopy.modflow.ModflowPcg(mf)

spd2 = {(0, 0): ["print head", "print budget", "save head", "save budget"]}
oc = flopy.modflow.ModflowOc(mf, stress_period_data=spd2, compact=True)

# Write the model input files
mf.write_input()

# Run the model
success, mfoutput = mf.run_model()
if not success:
    raise Exception("MODFLOW did not terminate normally.")

# Imports
import matplotlib.pyplot as plt
import flopy.utils.binaryfile as bf


hds = bf.HeadFile(modelname + ".hds")
head = hds.get_data(totim=1.0)

extent = (delr / 2.0, Lx - delr / 2.0, Ly - delc / 2.0, delc / 2.0)
fig = plt.figure(figsize=(6, 6))
ax = fig.add_subplot(1, 1, 1, aspect="equal")
ax.contour(head[0, :, :], levels=np.arange(1, 10, 1), extent=extent)