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
from subaxis import subaxis


def fig9_21():
    # uses 1st order symmetric differences

    # int(0.5*T*u_x^2-wu,x=0..1)
    # w=-x^r

    # n=number of internal grid points
    # n=3
    n = 20
    print("n =", n)
    r = 4

    x = np.linspace(0, 1, n + 2)
    h = x[1] - x[0]

    T = 1
    alpha = h**2 / T
    x = np.linspace(0, 1, n + 2)
    h = x[1] - x[0]

    A = np.zeros((n, n))
    A[0, 0] = 2
    A[0, 1] = -1
    for ix in range(1, n - 1):
        A[ix, ix - 1] = -1
        A[ix, ix] = 2
        A[ix, ix + 1] = -1
    A[n - 1, n - 2] = -1
    A[n - 1, n - 1] = 2

    print("cond(A,inf) =", np.linalg.cond(A, np.inf))

    b = np.zeros(n)
    for i in range(n):
        w = -x[i + 1] ** r
        b[i] = alpha * w

    u = np.linalg.solve(A, b)
    u = np.concatenate(([0], u, [0]))

    ne = 100
    xe = np.linspace(0, 1, ne)
    ue = np.zeros(ne)
    for i in range(ne):
        ue[i] = xe[i] * (xe[i] ** (r + 1) - 1) / ((r + 1) * (r + 2) * T)

    fig = matlab_figure(None, 658, 230)
    ax = subaxis(1, 1, 1, 1, "MT", 0.01, "MB", 0.18, "MR", -0.02, "ML", 0.06, "P", 0.04)
    set_color_order(ax)

    ax.plot(x, u, "o", markersize=8, linewidth=2)
    # plot(x,u,'--','MarkerSize',8,'LineWidth',2)
    ax.plot(xe, ue, "r", linewidth=1.6)
    ax.grid(True)
    ax.set_yticks([-0.02, -0.01, 0])
    # set(gca,'YTickLabel',{'0.02';'0.01';'0.0'})
    ax.set_xlabel("x-axis")
    ax.set_ylabel("u-axis")
    style_axes(ax, 14)

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


if __name__ == "__main__":
    fig9_21()
