%matplotlib inline
import os
import platform
import shutil
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import flopy
import pyemu


nam_file = "freyberg.nam"
org_model_ws = "freyberg_sfr_update"
temp_model_ws = "temp"
new_model_ws = "template"
# load the model, change dir and run once just to make sure everthing is working
m = flopy.modflow.Modflow.load(nam_file,model_ws=org_model_ws,check=False, exe_name="mfnwt",
                              forgive=False,verbose=True)
m.change_model_ws(temp_model_ws,reset_external=True)
m.write_input()
EXE_DIR = os.path.join("..","bin")
if "window" in platform.platform().lower():
    EXE_DIR = os.path.join(EXE_DIR,"win")
elif "darwin" in platform.platform().lower() or "macos" in platform.platform().lower():
    EXE_DIR = os.path.join(EXE_DIR,"mac")
else:
    EXE_DIR = os.path.join(EXE_DIR,"linux")

[shutil.copy2(os.path.join(EXE_DIR,f),os.path.join(temp_model_ws,f)) for f in os.listdir(EXE_DIR)]

try:
    m.run_model()
except():
    pass

"""
Explanation: "What just happened???"
Here we take an existing modflow model and setup a very complex parameterization system for arrays and boundary conditions.  All parameters are setup as multpliers: the original inputs from the modflow model are saved in separate files and during the forward run, they are multplied by the parameters to form new model inputs.  the forward run script ("forward_run.py") is also written.  And somewhat meaningful prior covariance matrix is constructed from geostatistical structures with out any additional arguements...oh yeah!
End of explanation
"""


m.get_package_list()

"""
Explanation: You want some pilot points? We got that...how about one set of recharge multiplier pilot points applied to all stress periods? and sy in layer 1?
End of explanation
"""


pp_props = [["upw.sy",0], ["rch.rech",None]]

"""
Explanation: Parameterization
End of explanation
"""


const_props = []
for iper in range(m.nper): # recharge for past and future
    const_props.append(["rch.rech",iper])
for k in range(m.nlay):
    const_props.append(["upw.hk",k])
    const_props.append(["upw.ss",k])


"""
Explanation: You want some constants (uniform value multipliers)?  We got that too....
End of explanation
"""


grid_props = [["upw.sy",0],["upw.vka",1]]
for k in range(m.nlay):
    grid_props.append(["upw.hk",k])
    

"""
Explanation: You want grid-scale parameter flexibility for hk in all layers? We got that too...and how about sy in layer 1 and vka in layer 2 while we are at it
End of explanation
"""


zn_array = np.loadtxt(os.path.join("Freyberg_Truth","hk.zones"))
plt.imshow(zn_array)

zone_props = [["upw.ss",0], ["rch.rech",0],["rch.rech",1]]
k_zone_dict = {k:zn_array for k in range(m.nlay)}


"""
Explanation: Some people like using zones...so we have those too
End of explanation
"""


bc_props = []
for iper in range(m.nper):
    bc_props.append(["wel.flux",iper])

"""
Explanation: But wait, boundary conditions are uncertain too...Can we add some parameter to represent that uncertainty?  You know it!
End of explanation
"""


# here were are building a list of stress period, layer pairs (zero-based) that we will use
# to setup obserations from every active model cell for a given pair
hds_kperk = []
for iper in range(m.nper):
    for k in range(m.nlay):
        hds_kperk.append([iper,k])

"""
Explanation: Observations
Since observations are "free", we can carry lots of them around...
End of explanation
"""


mfp_boss = pyemu.helpers.PstFromFlopyModel(nam_file,new_model_ws,org_model_ws=temp_model_ws,
                                          pp_props=pp_props,spatial_list_props=bc_props,
                                          zone_props=zone_props,grid_props=grid_props,
                                          const_props=const_props,k_zone_dict=k_zone_dict,
                                          remove_existing=True,pp_space=4,sfr_pars=True,
                                          sfr_obs=True,hds_kperk=hds_kperk)
EXE_DIR = os.path.join("..","bin")
if "window" in platform.platform().lower():
    EXE_DIR = os.path.join(EXE_DIR,"win")
elif "darwin" in platform.platform().lower():
    EXE_DIR = os.path.join(EXE_DIR,"mac")
else:
    EXE_DIR = os.path.join(EXE_DIR,"linux")

[shutil.copy2(os.path.join(EXE_DIR,f),os.path.join(new_model_ws,f)) for f in os.listdir(EXE_DIR)]


"""
Explanation: Here it goes...
Now we will use all these args to construct a complete PEST interface - template files, instruction files, control file and even the forward run script!  All parameters are setup as multiplers against the existing inputs in the modflow model - the existing inputs are extracted (with flopy) and saved in a sub directory for safe keep and for multiplying against during a forward model run.  The constructor will also write a full (covariates included) prior parameter covariance matrix, which is needed for all sorts of important analyses.|
End of explanation
"""


pst = mfp_boss.pst

pst.npar,pst.nobs

"""
Explanation: The mpf_boss instance containts a pyemu.Pst object (its already been saved to a file, but you may want to manipulate it more)
End of explanation
"""


pst.template_files

pst.instruction_files

"""
Explanation: That was crazy easy - this used to take me weeks to get a PEST interface setup with level of complexity
End of explanation
"""


cov = pyemu.Cov.from_ascii(os.path.join(new_model_ws,m.name+".pst.prior.cov"))

cov = cov.x
cov[cov==0] = np.NaN
plt.imshow(cov)

"""
Explanation: Lets look at that important prior covariance matrix
End of explanation
"""


pst.parameter_data

"""
Explanation: adjusting parameter bounds
Let's say you don't like the parameter bounds in the new control file (note you can pass a par_bounds arg to the constructor).
End of explanation
"""


par = pst.parameter_data #get a ref to the parameter data dataframe
wpars = par.pargp=="welflux_k02"
par.loc[wpars]

par.loc[wpars,"parubnd"] = 1.1
par.loc[wpars,"parlbnd"] = 0.9
pst.parameter_data

# now we need to rebuild the prior parameter covariance matrix
cov = mfp_boss.build_prior()

"""
Explanation: Let's change the welflux pars
End of explanation
"""


x = cov.x
x[x==0.0] = np.NaN
plt.imshow(x)

"""
Explanation: Boom!
End of explanation
"""


