from __future__ import annotations

import os

import matplotlib.pyplot as plt
import numpy as np

from plot_utils import finish_figure, matlab_figure, set_color_order, style_axes


i_fun = 0


def F(x):
    # F=dot(f,f)
    global i_fun
    i_fun += 1
    xx = x[0]
    yy = x[1]
    f1 = 10 * (xx**2 - yy)
    f2 = xx - 1
    f = np.array([f1, f2])
    return f


def Jac(x):
    # jacobian
    global i_fun
    i_fun += 1
    xx = x[0]
    yy = x[1]
    dfx = 20 * xx
    dfy = -10
    dgx = 1
    dgy = 0
    g = np.array([[dfx, dfy], [dgx, dgy]], dtype=float)
    _ = yy
    return g


def fig9_17():
    # i_fun = number of function evaluations
    global i_fun

    # 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

    i_fun = 0

    x = np.array([3 / 4, -3 / 4], dtype=float)
    exact = np.array([1.0, 1.0])
    tol = 1e-6
    xpoints = [x[0]]
    ypoints = [x[1]]
    f = F(x)
    ff = np.linalg.norm(f, 2) ** 2
    D = np.zeros((2, 2))
    mu = 10
    alpha = 2
    beta = 2
    counter = 1
    err = 10 * tol
    max_iter = int(os.getenv("FIG9_17_MAX_ITER", "10000"))
    # while err>tol #& counter<200
    while np.linalg.norm(x - exact, np.inf) > tol and counter < max_iter:
        counter += 1
        xx = x.copy()
        ff0 = ff

        J = Jac(x)
        JJ = J.T @ J
        gk = 2 * J.T @ f
        D[0, 0] = 2 * JJ[0, 0]
        D[1, 1] = 2 * JJ[1, 1]
        S = 2 * JJ + mu * D

        # check on descent direction
        # eig(S)

        dx = -np.linalg.solve(S, gk)
        x = xx + dx
        f = F(x)
        ff = np.linalg.norm(f, 2) ** 2

        if ff < ff0:
            mu = mu / beta
            xpoints.append(x[0])
            ypoints.append(x[1])
            err = np.linalg.norm(np.linalg.solve(S, gk) / x, np.inf)

            # check on jacobian approximation
            fx = 20 * x[0]
            fy = -10
            gx = 1
            gy = 0
            fxx = 20
            fxy = 0
            fyy = 0
            gxx = 0
            gxy = 0
            gyy = 0
            check = np.array(
                [
                    (f[0] * fxx + f[1] * gxx) / (fx**2 + gx**2),
                    (f[0] * fxy + f[1] * gxy) / (fx * fy + gx * gy),
                    (f[0] * fyy + f[1] * gyy) / (fy**2 + gy**2),
                ]
            )
            print(check)
        else:
            mu = alpha * mu
            x = xx
            ff = ff0

    print("i_fun =", i_fun)
    # check on jacobian approximation

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

    x = np.array([-1.0, -0.5])
    # x=[-1.2, 1.8];
    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
    ax.plot(exact[0], exact[1], "r*", markersize=14, linewidth=2)

    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=[100 / 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, 14)
    fig.tight_layout()

    # exportgraphics(gcf,'/Users/mark/Desktop/LMM.eps')
    finish_figure(fig, "LMM.png")
    plt.show()


if __name__ == "__main__":
    fig9_17()
