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_5():
    # SDM for 2x2 system
    #       [a b][x1] = [b1]
    #       [c d][x2] = [b2]

    fig = matlab_figure(None, 482, 477)
    ax = fig.add_subplot(111)
    set_color_order(ax)

    # specify example (3 and 1 are in book)
    iexample = 3

    # set parameters
    if iexample == 1:
        x = np.array([1.0, 1.0])
        a = 5
        b = 4.99
        c = b
        d = 5
        b1 = 1
        b2 = -1
        xL = -2
        xR = 80
        yL = -80
        yR = 5
        cxL = 0
        cxR = 100
        cy = -100
        ax.axis([xL, xR, yL, yR])
        ax.set_xlabel("x-axis", fontsize=18, fontweight="bold")
        ax.set_yticks([-80, -60, -40, -20, 0])
        nc = 10
        cx = np.linspace(cxL, cxR, nc)
        cv = np.zeros(nc)
        for i in range(nc):
            cv[i] = 0.5 * (a * cx[i] ** 2 + 2 * b * cx[i] * cy + d * cy * cy) - b1 * cx[i] - b2 * cy
    elif iexample == 2:
        x = np.array([1.0, 1.0])
        a = 2
        b = 2
        c = b
        d = 3
        b1 = 1
        b2 = -1
        xL = 0
        xR = 6
        yL = -5
        yR = 1
        cxL = 3
        cxR = 10
        cy = -2
        ax.axis([0, 4, -3, 1])
        ax.set_xlabel("x-axis")
        ax.set_ylabel("y-axis")
        cx = np.linspace(cxL, cxR, 18)
        cv = np.zeros(18)
        for i in range(18):
            cv[i] = 0.5 * (a * cx[i] ** 2 + 2 * b * cx[i] * cy + d * cy * cy) - b1 * cx[i] - b2 * cy
    elif iexample == 3:
        x = np.array([1.0, 2.0])
        a = 3
        b = -2
        c = -2
        d = 4
        b1 = 1
        b2 = -1
        xL = -0.5
        xR = 2.5
        yL = -0.5
        yR = 2.5
        cxL = 0
        cxR = 2
        cy = 1
        ax.axis([xL, xR, yL, yR])
        ax.set_xlabel("x-axis")
        ax.set_ylabel("y-axis")
        cv = [-0.13, 0.1, 0.46, 1.2, 2.3, 3.7, 6.5, 10]
        ax.set_xticks([-0.5, 1.0, 2.5])
        ax.set_yticks([-0.5, 1.0, 2.5])
    else:
        x = np.array([1.0, 1.0])
        a = 2
        b = 1
        c = 1
        d = 3
        b1 = 1
        b2 = -1
        xL = 0
        xR = 2
        yL = -1
        yR = 1
        cxL = 0
        cxR = 2
        cy = 1
        ax.axis([xL, xR, yL, yR])
        ax.set_xlabel("x-axis")
        ax.set_ylabel("y-axis")
        cv = [-0.68, -0.585, -0.3, 0.4, 1.3, 2.3, 3.5, 4.8]
    tol = 0.0001

    # generate surface data for plotting
    xd = np.linspace(xL, xR, 60)
    yd = np.linspace(yL, yR, 60)
    X, Y = np.meshgrid(xd, yd)
    Z = 0.5 * (a * X**2 + 2 * b * X * Y + d * Y**2) - b1 * X - b2 * Y
    # surf(X,Y,Z)
    # pause

    # start iteration
    xpoints = [x[0]]
    ypoints = [x[1]]
    r = np.array([b1 - a * x[0] - b * x[1], b2 - c * x[0] - d * x[1]])
    diff = np.array([1.0, 1.0])
    counter = 1
    while np.linalg.norm(diff) > tol and counter < 30:
        counter += 1
        q = np.array([a * r[0] + b * r[1], c * r[0] + d * r[1]])
        alpha = np.dot(r, r) / np.dot(r, q)
        diff = alpha * r
        x = x + diff
        r = r - alpha * q

        # fprintf('\n  %i  Computed Solution =  %e  %e \n\n',counter,x);

        xpoints.append(x[0])
        ypoints.append(x[1])
        ax.clear()
        set_color_order(ax)
        ax.set_xlim(xL, xR)
        ax.set_ylim(yL, yR)
        ax.set_aspect("equal", adjustable="box")
        ax.contour(X, Y, Z, levels=cv, colors="b", linewidths=1.6)
        ax.plot(xpoints, ypoints, "r-", linewidth=2.5)
        ax.plot(xpoints, ypoints, "o", color=[110 / 256, 0, 0], markersize=15, linewidth=2.5)
        error = np.linalg.norm(diff)
        style_axes(ax, 18)

    # 0.5*(a*x(1)^2 + 2*b*x(1)*x(2) + d*x(2)*x(2)) - b1*x(1) - b2*x(2)

    print("xpoints(2) =", xpoints[1])
    print("ypoints(2) =", ypoints[1])

    A = np.array([[a, b], [c, d]], dtype=float)
    xe = np.linalg.solve(A, np.array([b1, b2], dtype=float))
    print("length(xpoints) =", len(xpoints))
    print("error =", np.linalg.norm(x - xe, np.inf))
    print("eigenvalues =", np.linalg.eigvals(A))
    print("cond(A,inf) =", np.linalg.cond(A, np.inf))
    print("cond(A,2) =", np.linalg.cond(A, 2))

    # if iexample==3
    #     exportgraphics(gcf,'/Users/mark/Desktop/sdm01.eps')
    # elseif  iexample==1
    #     exportgraphics(gcf,'/Users/mark/Desktop/sdm02.eps')
    # end
    finish_figure(fig, "sdm01.png" if iexample == 3 else "sdm02.png")
    plt.show()


if __name__ == "__main__":
    fig9_5()
