"""
2D steady-state modeling by Freyberg (Groundwater, 1988)
You have to download or git-clone the Freyberg example from flopy github repository
"""
import flopy
import os
import matplotlib.pyplot as plt

model_ws = "data"

# load freyberg example (Groundwater, 1988)
mf = flopy.modflow.Modflow.load("freyberg.nam", model_ws=model_ws, verbose=True,check=False, exe_name="mf2005")

# check the input files and print out loaded packages
mf.check()

# spatial coordinate
print(mf.dis.sr)


# export shapefile (need to install pyshp package)
#mf.dis.export("data/freyberg_dis.shp")

# you can read shapefile
#import shapefile
#sf = shapefile.Reader("data/freyberg_dis.shp")

mf.write_input()
mf.run_model()

# install fandas
#mfl = flopy.utils.MfListBudget(os.path.join(model_ws,"freyberg.lst"))
#df_flux, df_vol = mfl.get_dataframes(start_datetime="10-21-2015")
#df_flux

#groups = df_flux.groupby(lambda x:x.split('_')[-1],axis=1).groups
#df_flux_in = df_flux.loc[:,groups["IN"]]
#df_flux_in.columns = df_flux_in.columns.map(lambda x:x.split('_')[0])

#df_flux_out = df_flux.loc[:,groups["OUT"]]
#df_flux_out.columns = df_flux_out.columns.map(lambda x:x.split('_')[0])

#df_flux_delta = df_flux_in - df_flux_out
#df_flux_delta.iloc[-1,:].plot(kind="bar",figsize=(10,10),grid=True)

# you can use object plot method
mf.lpf.hk.plot(mflay=0,colorbar=True)
plt.show()

# check river boundary conditions
mf.riv.stress_period_data.data

h = flopy.utils.HeadFile(os.path.join(model_ws,"freyberg.hds"),model=mf)
#h.times # time
#h.get_data(totim = h.times[-1]) # actual head result
# check the range of head
# head.min()
# head[head < 999.0].max() #  because inactive cell val = 999.0
h.plot(totim=h.times[-1],mflay=0,contour=True, grid=True,colorbar=True ,figsize=(10,10), vmin=0,vmax=30)
plt.show()