'''
groundwater supply optimization example using scipy.optimize.basinhopping
Note that basinhopping internally uses local optimization method
and I used Nelder-Mead, gradient free downhill simplex algorithm
Single precision MODFLOW results - hydraulic head seem not smooth and differentiable,
thus typical local optimization did not work well for this problem.
Maybe this is because the minimum head is always at the pumping well.
'''
from mymf_v3 import mymf
import numpy as np
from scipy.optimize import basinhopping

def f(Q):
    ''' objective function of groundwater supply

    :param Q: pumping rate
    :return objval: cost (pumping rate + drawdown penalty if head < 0)
    '''
    init_head = 1.  # initial head
    model = mymf(init_head=init_head)
    well_rcs = [[4, 4]]  # center
    model.run(well_rcs, Q)
    head = model.head() # note this array is 3D!
    minhead = model.minhead()

    objval = Q
    if minhead < 0.0:
        objval = objval + 10000000.*(minhead)**2
    if Q > 0.0: # don't allow injection
        objval = 1000000.0

    print("Q: %f, objval: %f" % (Q, objval))
    return objval

x0 = -10.0

# basinhopping internally uses local optimization method
# here I used Nelder-Mead, gradient free downhill simplex algorithm
minimizer_kwargs = {'method':"Nelder-Mead"}
ret = basinhopping(f, x0, minimizer_kwargs=minimizer_kwargs, niter = 10, disp=True)


print(ret)