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 fig9_15():
    # this include figures 9.15 and 9.16 (SDM/CGM portions)

    # 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
    #       dF  is gradient of F(x,y)

    # Specify which example to run ifun=1 (Rosen), 2 (multi-min)
    ifun = 2
    print("ifun =", ifun)

    # specify which beta choice SDM (ibeta=1), CGM-FR (ibeta=2),
    # CGM-PR (ibeta=3)
    ibeta = 3
    print("ibeta =", ibeta)

    fig = matlab_figure(None, 421, 395)
    ax = fig.add_subplot(111)
    set_color_order(ax)

    n_fun = 0
    for ic in range(1, 3):
        if ifun == 1:
            if ic == 2:
                break
            # x=[-1.0, -0.5];
            # x=[-1.2, 1.8];
            # x=round(2*rand(1,2)-1,2)

            # x=[4/5   -3/5]
            x = np.array([3 / 4, -3 / 4], dtype=float)
            print("x =", x)
            # x=[-1.2 1]
            xa = -1.1
            xb = 1.1
            ya = -1.1
            yb = 1.1
            xd = np.linspace(xa, xb, 80)
            yd = np.linspace(ya, yb, 80)
            X, Y = np.meshgrid(xd, yd)
            Z = 100 * (X**2 - Y) ** 2 + (X - 1) ** 2
            exact = np.array([1.0, 1.0])
            ax.plot(exact[0], exact[1], "r*", markersize=14, linewidth=2)
        else:
            if ic == 1:
                x = np.array([1.0, 0.8])
                # options = optimset('TolX',1e-10,'MaxIter',5000,'MaxFunEvals',5000);
                # ee=fminsearch(@(x)(x(1)-x(2))^4+8*x(1)*x(2)-x(1)+x(2)+3,x,options)
                exact = np.array([0.553579941873744, -0.553579935896837])
            else:
                x = np.array([-1.0, -0.6])
                # options = optimset('TolX',1e-10,'MaxIter',5000,'MaxFunEvals',5000);
                # exact=fminsearch(@(x)(x(1)-x(2))^4+8*x(1)*x(2)-x(1)+x(2)+3,[-0.5,0.5],options)
                exact = np.array([-0.418782723467334, 0.4187827111776])
            xa = -1
            xb = 1
            ya = -1.5
            yb = 1
            xd = np.linspace(xa, xb, 80)
            yd = np.linspace(ya, yb, 80)
            X, Y = np.meshgrid(xd, yd)
            Z = (X - Y) ** 4 + 8 * X * Y - X + Y + 3

        xpoints = [x[0]]
        ypoints = [x[1]]
        tol = 1e-6
        # n_fun=0;

        def dF(iff):
            nonlocal n_fun, x
            n_fun += 1
            if ifun == 1:
                g = np.array(
                    [
                        400 * x[0] * (x[0] ** 2 - x[1]) + 2 * (x[0] - 1),
                        -200 * (x[0] ** 2 - x[1]),
                    ]
                )
            elif ifun == 2:
                g = np.array(
                    [
                        4 * (x[0] - x[1]) ** 3 + 8 * x[1] - 1,
                        -4 * (x[0] - x[1]) ** 3 + 8 * x[0] + 1,
                    ]
                )
            else:
                g = np.array([20 * x[0], 2 * x[1]])
            return g

        g = dF(ifun)
        d = -g
        counter = 1
        alphaM = 1 / 2

        def F(s):
            nonlocal n_fun
            n_fun += 1
            xs = x + s * d
            if ifun == 1:
                z = 100 * (xs[0] ** 2 - xs[1]) ** 2 + (xs[0] - 1) ** 2
            elif ifun == 2:
                z = (xs[0] - xs[1]) ** 4 + 8 * xs[0] * xs[1] - xs[0] + xs[1] + 3
            else:
                z = 10 * xs[0] ** 2 + xs[1] ** 2
            return float(z)

        F0 = F(0)

        def arm(aM):
            nonlocal F0, alphaM
            gam = 1 / 20
            tau = 0.5
            s = aM
            Fs = F(s)
            mk = np.dot(g, d)
            gd = gam * mk

            for ir in range(20):
                for is_ in range(20):
                    if Fs < F0 and is_ > 0:
                        alphaM = 20 * s
                        break
                    s = tau**2 * s
                    Fs = F(s)

                if Fs < F0 + s * gd:
                    F0 = Fs
                    return s

                if ir == 0:
                    S = s
                    FS = Fs
                    _ = (S, FS)
                    s = tau * s
                    Fs = F(s)
                    if Fs < F0 + s * gd:
                        F0 = Fs
                        return s
                else:
                    s = tau * s
                    Fs = F(s)
                    if Fs < F0 + s * gd:
                        F0 = Fs
                        return s
            z = 0
            print("z =", z)
            return s

        while np.linalg.norm(x - exact, np.inf) > tol and counter < 500:
            counter += 1

            # alpha=fminbnd(@F,0,alphaM);
            alpha = arm(alphaM)

            # alpha=-dot(g,d)/dot(d,H(x)*d')

            # if alpha>alphaM/3
            #     alphaM=2*alphaM;
            # elseif alpha<alphaM/64
            #     alphaM=alphaM/2;
            # end

            x = x + alpha * d
            g0 = g.copy()
            g = dF(ifun)

            if ibeta == 3:
                # CGM-PR
                beta = np.dot(g, g - g0) / np.dot(g0, g0)
                if np.dot(g, g) - beta * np.dot(d, g) < 0:
                    beta = 0
            elif ibeta == 2:
                # CGM-FR
                beta = np.dot(g, g) / np.dot(g0, g0)
            else:
                beta = 0

            d = -g + beta * d
            xpoints.append(x[0])
            ypoints.append(x[1])

        print("n_fun =", n_fun)

        error = np.linalg.norm(exact - x, np.inf)
        print(f"\n  Number of Steps = {counter:d}")
        print(f"  Computed Solution =  {x[0]:e}  {x[1]:e}")
        # fprintf('\n  Error = %e \n\n',error)
        _ = error

        ax.contour(X, Y, Z, levels=40, colors="b", linewidths=1)
        ax.axis([xa - 0.004, xb, ya, yb])
        ax.set_aspect("equal", adjustable="box")
        ax.plot(xpoints, ypoints, "r-", linewidth=2)
        ax.plot(xpoints, ypoints, "o", color=[110 / 256, 0, 0], markersize=10, linewidth=2)

    ax.set_xticks([-1, -0.5, 0, 0.5, 1])
    ax.set_yticks([-1, -0.5, 0, 0.5, 1])
    ax.set_xlabel("x-axis")
    ax.set_ylabel("y-axis")
    style_axes(ax, 16)
    fig.tight_layout()

    # if ifun==2 && ibeta==1
    #     exportgraphics(gcf,'/Users/mark/Desktop/sdm11.eps')
    # elseif  ifun==1 && ibeta==1
    #     exportgraphics(gcf,'/Users/mark/Desktop/sdm21.eps')
    # elseif ifun==2 && ibeta==3
    #     exportgraphics(gcf,'/Users/mark/Desktop/cgm11.eps')
    # elseif  ifun==1 && ibeta==3
    #     exportgraphics(gcf,'/Users/mark/Desktop/cgm21.eps')
    # end
    if ifun == 2 and ibeta == 1:
        filename = "sdm11.png"
    elif ifun == 1 and ibeta == 1:
        filename = "sdm21.png"
    elif ifun == 2 and ibeta == 3:
        filename = "cgm11.png"
    else:
        filename = "cgm21.png"
    finish_figure(fig, filename)
    plt.show()


if __name__ == "__main__":
    fig9_15()
