from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np

from plot_utils import finish_figure, matlab_figure, set_color_order, style_axes


def F(x, y, ifun):
    if ifun == 1:
        z = 100 * (x**2 - y) ** 2 + (x - 1) ** 2
    elif ifun == 2:
        z = (x - y) ** 4 + 8 * x * y - x + y + 3
    else:
        z = 10 * x**2 + y**2
    return z


def sortrows(DD):
    return DD[np.argsort(DD[:, 0])]


def fig9_19():
    # this also contains fig9_20

    # Nelder-Mead algorithm

    # Inputs:
    #       x = starting vector
    #       xa, xb = x-interval used in contour plot
    #       ya, yb = y-interval used in contour plot
    #       tol = tolerance for stopping iteration

    # Required (end of file)
    #       F    is function F(x,y) to be minimized

    tol = 0.0001

    # Specify which example to run ifun=1,2,3 (2,3 in book)
    ifun = 3
    print("ifun =", ifun)

    fig = matlab_figure(None, 560, 420)
    ax = fig.add_subplot(111)
    set_color_order(ax)

    if ifun == 1:
        ns = 5
        x1 = np.array([-1.0, -1.0])
        x2 = np.array([-1.0, -0.72])
        x3 = np.array([-0.7, -1.0])
        xa = -1.1
        xb = 1.1
        ya = -1.1
        yb = 1.1
        xd = np.linspace(xa, xb, 200)
        yd = np.linspace(ya, yb, 200)
        X, Y = np.meshgrid(xd, yd)
        Z = 100 * (X**2 - Y) ** 2 + (X - 1) ** 2
        exact = np.array([1.0, 1.0])
        # surfc(X,Y,Z)
        # pause
        ax.set_yticks([-1, -0.5, 0, 0.5, 1])
        ax.set_xticks([-1, -0.5, 0, 0.5, 1])
    elif ifun == 2:
        ns = 5
        x1 = np.array([-1.0, -1.0])
        x2 = np.array([-1.0, -0.7])
        x3 = np.array([-0.7, -1.0])
        xa = -1.1
        xb = 1.1
        ya = -1.1
        yb = 1.1
        xd = np.linspace(xa, xb, 200)
        yd = np.linspace(ya, yb, 200)
        X, Y = np.meshgrid(xd, yd)
        Z = (X - Y) ** 4 + 8 * X * Y - X + Y + 3
        exact = np.array([0.55357993584438379685, -0.55357993584438379685])
        # surfc(X,Y,Z)
        # pause
        ax.set_yticks([-1, -0.5, 0, 0.5, 1])
        ax.set_xticks([-1, -0.5, 0, 0.5, 1])
    else:
        ns = 4
        x1 = np.array([-1.0, -1.0])
        x2 = np.array([-1.0, -0.4])
        x3 = np.array([-0.5, -1.0])
        xa = -1.1
        xb = 1.1
        ya = -1.1
        yb = 1.1
        xd = np.linspace(xa, xb, 200)
        yd = np.linspace(ya, yb, 200)
        X, Y = np.meshgrid(xd, yd)
        Z = 10 * X**2 + Y**2
        exact = np.array([0.0, 0.0])
        # surfc(X,Y,Z)
        # pause
        ax.set_yticks([-1, -0.5, 0, 0.5, 1])
        ax.set_xticks([-1, -0.5, 0, 0.5, 1])

    # DD = [F value, x value, y value] sorted by values of F (small to big)
    DD = sortrows(
        np.array(
            [
                [F(x1[0], x1[1], ifun), x1[0], x1[1]],
                [F(x2[0], x2[1], ifun), x2[0], x2[1]],
                [F(x3[0], x3[1], ifun), x3[0], x3[1]],
            ]
        )
    )

    ax.contour(X, Y, Z, levels=30, colors="b", linewidths=1)
    xpoints1 = [DD[0, 1], DD[1, 1]]
    ypoints1 = [DD[0, 2], DD[1, 2]]
    xpoints2 = [DD[1, 1], DD[2, 1]]
    ypoints2 = [DD[1, 2], DD[2, 2]]
    xpoints3 = [DD[2, 1], DD[0, 1]]
    ypoints3 = [DD[2, 2], DD[0, 2]]
    if ifun == 1:
        ax.plot(exact[0], exact[1], "r*", markersize=14, linewidth=2)
    ax.plot(xpoints1, ypoints1, "ro-", linewidth=2)
    ax.plot(xpoints2, ypoints2, "ro-", linewidth=2)
    ax.plot(xpoints3, ypoints3, "ro-", linewidth=2)

    ax.set_xlabel("x-axis")
    ax.set_ylabel("y-axis")

    say = "1"
    xx = (DD[0, 1] + DD[1, 1] + DD[2, 1]) / 3
    yy = (DD[0, 2] + DD[1, 2] + DD[2, 2]) / 3
    ax.text(xx, yy, say, fontsize=22, fontweight="bold")

    style_axes(ax, 14)

    alpha = 1
    beta = 0.5
    for is_ in range(1, ns + 1):
        DD = sortrows(DD)

        xB = DD[0, 1] - DD[2, 1]
        yB = DD[0, 2] - DD[2, 2]
        xC = DD[1, 1] - DD[2, 1]
        yC = DD[1, 2] - DD[2, 2]
        area = 0.5 * (xB * yC - xC * yB)
        if abs(area) < tol:
            break

        cen = 0.5 * np.array([DD[0, 1] + DD[1, 1], DD[0, 2] + DD[1, 2]])
        r = cen + alpha * np.array([cen[0] - DD[2, 1], cen[1] - DD[2, 2]])
        Fr = F(r[0], r[1], ifun)

        if Fr < DD[1, 0] and DD[0, 0] <= Fr:
            DD[2, :] = [Fr, r[0], r[1]]
        elif Fr < DD[0, 0]:
            rr = cen + 2 * alpha * np.array([cen[0] - DD[2, 1], cen[1] - DD[2, 2]])
            Frr = F(rr[0], rr[1], ifun)
            if Frr < Fr:
                DD[2, :] = [Frr, rr[0], rr[1]]
            else:
                DD[2, :] = [Fr, r[0], r[1]]
        elif Fr < DD[2, 0]:
            rA = cen + beta * np.array([cen[0] - DD[2, 1], cen[1] - DD[2, 2]])
            FA = F(rA[0], rA[1], ifun)
            if FA <= Fr:
                DD[2, :] = [FA, rA[0], rA[1]]
            else:
                DD[1, 1] = DD[0, 1] + beta * (DD[1, 1] - DD[0, 1])
                DD[1, 2] = DD[0, 2] + beta * (DD[1, 2] - DD[0, 2])
                DD[1, 0] = F(DD[1, 1], DD[1, 2], ifun)
                DD[2, 1] = DD[0, 1] + beta * (DD[2, 1] - DD[0, 1])
                DD[2, 2] = DD[0, 2] + beta * (DD[2, 2] - DD[0, 2])
                DD[2, 0] = F(DD[2, 1], DD[2, 2], ifun)
            # step 6
        else:
            rA = np.array([DD[2, 1], DD[2, 2]]) + beta * np.array(
                [cen[0] - DD[2, 1], cen[1] - DD[2, 2]]
            )
            FA = F(rA[0], rA[1], ifun)
            if FA < DD[2, 0]:
                DD[2, :] = [FA, rA[0], rA[1]]
            else:
                DD[1, 1] = DD[0, 1] + beta * (DD[1, 1] - DD[0, 1])
                DD[1, 2] = DD[0, 2] + beta * (DD[1, 2] - DD[0, 2])
                DD[1, 0] = F(DD[1, 1], DD[1, 2], ifun)
                DD[2, 1] = DD[0, 1] + beta * (DD[2, 1] - DD[0, 1])
                DD[2, 2] = DD[0, 2] + beta * (DD[2, 2] - DD[0, 2])
                DD[2, 0] = F(DD[2, 1], DD[2, 2], ifun)

        xpoints1 = [DD[0, 1], DD[1, 1]]
        ypoints1 = [DD[0, 2], DD[1, 2]]
        xpoints2 = [DD[1, 1], DD[2, 1]]
        ypoints2 = [DD[1, 2], DD[2, 2]]
        xpoints3 = [DD[2, 1], DD[0, 1]]
        ypoints3 = [DD[2, 2], DD[0, 2]]
        ax.plot(xpoints1, ypoints1, "ro-", linewidth=2)
        ax.plot(xpoints2, ypoints2, "ro-", linewidth=2)
        ax.plot(xpoints3, ypoints3, "ro-", linewidth=2)

        iss = is_ + 1
        say = str(iss)
        xx = (DD[0, 1] + DD[1, 1] + DD[2, 1]) / 3
        yy = (DD[0, 2] + DD[1, 2] + DD[2, 2]) / 3
        ax.text(xx, yy, say, fontsize=22, fontweight="bold")

    ax.axis([xa, xb, ya, yb])
    fig.tight_layout()

    # if ifun==3
    #     exportgraphics(gca,'/Users/mark/Desktop/nm1.eps')
    # elseif ifun==2
    #     exportgraphics(gca,'/Users/mark/Desktop/nm2.eps')
    # end

    # error=sqrt((exact(1)-xx)^2+(exact(2)-yy)^2);
    # fprintf('\n\n  Number of Steps = %i',counter)
    # fprintf('\n  Computed Solution =  %e  %e',xx,yy);
    # fprintf('\n  Error = %e\n\n',error);
    _ = exact
    finish_figure(fig, "nm1.png" if ifun == 3 else "nm2.png")
    plt.show()


if __name__ == "__main__":
    fig9_19()
