from numpy import cos, arctan2
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

plt.rcParams["font.size"] = 10
plt.rcParams["mathtext.fontset"] = "cm"

# Minimization
def nelder_mead_step(fun, verts, alpha=1, gamma=2, rho=0.5,
                     sigma=0.5, beta=1.0):
    """Nelder-Mead iteration according to Wikipedia _[1]
    
    
    References
    ----------
     .. [1] Wikipedia contributors. "Nelder–Mead method." Wikipedia,
         The Free Encyclopedia. Wikipedia, The Free Encyclopedia,
         1 Sep. 2016. Web. 20 Sep. 2016. 
    """
    nverts, _ = verts.shape
    f = np.apply_along_axis(fun, 1, verts, beta=beta)
    # 1. Order
    order = np.argsort(f)
    verts = verts[order, :]
    f = f[order]
    # 2. Calculate xo, the centroid"
    xo = verts[:-1, :].mean(axis=0)
    # 3. Reflection
    xr = xo + alpha*(xo - verts[-1, :])
    fr = fun(xr, beta)
    if f[0]<=fr and fr<f[-2]:
        new_verts = np.vstack((verts[:-1, :], xr))
    # 4. Expansion
    elif fr<f[0]:
        xe = xo + gamma*(xr - xo)
        fe = fun(xe, beta)
        if fe < fr:
            new_verts = np.vstack((verts[:-1, :], xe))
        else:
            new_verts = np.vstack((verts[:-1, :], xe))
    # 5. Contraction
    else:
        xc = xo + rho*(verts[-1, :] - xo)
        fc = fun(xc, beta)
        if fc < f[-1]:
            new_verts = np.vstack((verts[:-1, :], xc))
    # 6. Shrink
        else:
            new_verts = np.zeros_like(verts)
            new_verts[0, :] = verts[0, :]
            for k in range(1, nverts):
                new_verts[k, :] = sigma*(verts[k,:] - verts[0,:])
 
    return new_verts

def fun(x, beta=1.0):
    """Simionescu function using log-barrier method"""
    x1, x2 = x
    if x1**2 + x2**2 < (1 + 0.2*cos(8*arctan2(x1, x2)))**2:
        barrier = -beta*np.log((1 + 0.2*cos(8*arctan2(x1, x2)))**2 - x1**2 - x2**2)
    else:
        barrier = np.inf
    return x1*x2 + barrier

# Animation

def data_gen(num):
    plt.gca().cla
    x0 = np.array([0.4, -0.6])
    x1 = np.array([-0.3, -0.6])
    x2 = np.array([0.7, 0.6])
    verts = np.vstack((x0, x1, x2))
    beta = 1.0
    for cont in range(num):
        verts = nelder_mead_step(fun, verts, beta=beta)
        beta /=2
    # Plots
    plt.cla()
    poly = plt.Polygon(verts, facecolor="none", edgecolor="k",
                       linewidth=0.5, zorder=4)
    plt.gca().add_patch(poly)
    x1, x2 = np.mgrid[-1.25:1.25:101j, -1.25:1.25:101j]
    z = x1*x2
    cons = x1**2 + x2**2 - (1 + 0.2*cos(8*arctan2(x1, x2)))**2
    z[cons > 0.02] = np.nan
    levels = np.linspace(-1, 1, 30)
    plt.contour(x1, x2, z, levels, cmap="seismic", linewidths=1)
    plt.contour(x1, x2, cons, [0], colors="black", linewidths=1)
    plt.axis("image")
    plt.xlabel(r"$x$", fontsize=14)
    plt.ylabel(r"$y$", fontsize=14)

fig = plt.figure(figsize=(5, 5))
ani = animation.FuncAnimation(fig, data_gen, range(25))
ani.save("Nelder-Mead_Simionescu.gif", writer='imagemagick', fps=2,
         dpi=200)
plt.show()