'''my_second_flopy_example.py: simple modification of my_second_flopy_example.py for optimization
    we use scipy.opitmize.minimize to maximize water supply hile allowing 1 m head drop
'''

import numpy as np
import flopy
import flopy.utils.binaryfile as bf

def f(pumping_rate):
    """return a function value for groundwater supply

    :param pumping_rate: value
    :return: objval: -pumping_rate + penalty for head violation
    """

    # 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] = 0.
    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 vertical hydraulic conductivity
    lpf = flopy.modflow.ModflowLpf(mf, hk=10., vka=10., ipakcb=53)

    # well location indices
    wcol = nrow / 2 - 1  # x index for a well
    wrow = ncol / 2 - 1  # y index for a well

    # 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)

    # Add wel package to the MODFLOW model
    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)

    # Write the MODFLOW model input files
    # If we cannot (over)write input files, try to write until it succeeds
    while True:
        try:
            mf.write_input()
        except OSError as err:
            print("File writing error: %s" % (err))
        else: # if we succeed, get out of the loop
            break

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

    hds = bf.HeadFile(modelname + '.hds')
    times = hds.get_times()  # simulation time, steady state
    head = hds.get_data(totim=times[-1])
    minhead = head.min()
    hds.close() # close the file object for the next run

    objval = pumping_rate
    if minhead < -1.0:
        objval = objval + 10000.*(minhead + 1.0)**2

    #mf.remove_package(pname='WEL') # remove package from model for the next run

    print("pumpting rate: %f, objval: %f" % (pumping_rate, objval))
    return objval

import scipy.optimize as opt
# Use Nelder-Mead/Simplex method for this problem
results = opt.minimize(f,x0=-50.,method="Nelder-Mead",tol = 1.e-2)
#results = opt.minimize(f,x0=-50.,method="BFGS",tol = 1.e-2)