import numpy as np
import matplotlib.pyplot as plt
from cycler import cycler


def fig5_5():
    ns = np.array([4, 16, 64])
    n = 1000
    # ns = np.array([40, 100, 200]); n = 2000
    # ns = np.array([400, 800, 2000]); n = 10000

    a = -1
    b = 1
    x = np.linspace(a, b, n)

    # exact
    ye = np.zeros(n)
    for ie in range(n):
        ye[ie] = 1 / (1 + 25 * x[ie] ** 2)

    # ns(1)
    nx = ns[0] + 1
    xd = np.linspace(a, b, nx)
    # xd = np.sort(2*np.random.rand(nx)-1)
    yd = np.zeros(nx)
    for ix in range(nx):
        yd[ix] = 1 / (1 + 25 * xd[ix] ** 2)
    p = np.zeros(n)
    for ii in range(n):
        for k in range(nx):
            p[ii] = p[ii] + yd[k] * ell(k, x[ii], xd)
    print(np.max(np.abs(p)))

    # ns(2)
    nx = ns[1] + 1
    xd = np.linspace(a, b, nx)
    # xd = np.sort(2*np.random.rand(nx)-1)
    yd = np.zeros(nx)
    for ix in range(nx):
        yd[ix] = 1 / (1 + 25 * xd[ix] ** 2)
    pp = np.zeros(n)
    for ii in range(n):
        for k in range(nx):
            pp[ii] = pp[ii] + yd[k] * ell(k, x[ii], xd)

    print(np.max(np.abs(pp)))

    # ns(3)
    nx = ns[2] + 1
    xd = np.linspace(a, b, nx)
    # xd = np.sort(2*np.random.rand(nx)-1)
    yd = np.zeros(nx)
    for ix in range(nx):
        yd[ix] = 1 / (1 + 25 * xd[ix] ** 2)
    ppp = np.zeros(n)
    for ii in range(n):
        for k in range(nx):
            ppp[ii] = ppp[ii] + yd[k] * ell(k, x[ii], xd)

    print(np.max(np.abs(ppp)))

    plt.close("all")
    # get(gcf)
    fig, axes = plt.subplots(3, 1, figsize=(7.5, 5.0))
    fig.subplots_adjust(left=0.08, right=0.98, bottom=0.11, top=0.99, hspace=0.16)
    co = np.array([
        [0, 0, 1],
        [0, 0.5, 0],
        [1, 0, 0],
        [0, 0.75, 0.75],
        [0.75, 0, 0.75],
        [0.75, 0.75, 0],
        [0.25, 0.25, 0.25],
    ])
    for ax in axes:
        ax.set_prop_cycle(cycler(color=co))

    ax = axes[0]
    ax.plot(x, ye, "r", linewidth=1.4)
    ax.plot(x, p, "--b", linewidth=1.4)
    ax.grid(True)
    # ax.set_xlim(-1, 1); ax.set_ylim(-0.5, 1.5)
    ax.set_xlim(-1, 1)
    ax.set_ylim(-0.5, 1.2)
    ax.set_xticks([-1, 0, 1])
    ax.set_yticks([-0.5, 1])
    ax.legend(
        [" Runge", " n = 4"],
        loc="lower center",
        prop={"size": 11, "weight": "bold"},
    )
    ax.tick_params(labelsize=12)

    ax = axes[1]
    ax.grid(True)
    ax.plot(x, ye, "r", linewidth=1.4)
    ax.plot(x, pp, "--b", linewidth=1.8)
    # ax.set_xlim(-1, 1); ax.set_ylim(-0.5, 1.5)
    ax.set_xlim(-1, 1)
    ax.set_ylim(-15, 2)
    ax.set_xticks([-1, 0, 1])
    ax.set_yticks([-15, 2])
    ax.legend(
        [" Runge", " n = 16"],
        loc="lower center",
        prop={"size": 11, "weight": "bold"},
    )
    ax.tick_params(labelsize=12)

    ax = axes[2]
    ax.grid(True)
    ax.plot(x, ye, "r", linewidth=1.4)
    ax.plot(x, ppp, "--b", linewidth=1.8)
    # ax.set_xlim(-1, 1); ax.set_ylim(-0.5, 1.5)
    ax.set_xlim(-1, 1)
    ax.set_ylim(-15, 2)
    # ax.set_yticks([-4, -3, -2, -1, 0, 1])
    ax.set_xlabel("x-axis")
    ax.set_xticks([-1, 0, 1])
    ax.set_yticks([-15, 2])
    ax.legend(
        [" Runge", " n = 64"],
        loc="lower center",
        prop={"size": 11, "weight": "bold"},
    )
    ax.tick_params(labelsize=12)

    for ax in axes:
        for label in ax.get_xticklabels() + ax.get_yticklabels():
            label.set_fontweight("bold")
        for spine in ax.spines.values():
            spine.set_visible(True)

    plt.show()

    # plt.savefig('/Users/mark/Desktop/runge.eps')


# ell_i(x) function
def ell(i, x, xd):
    n2 = xd.size
    p = 1.0
    for j in range(n2):
        if j != i:
            p = p * (x - xd[j]) / (xd[i] - xd[j])
    return p


if __name__ == "__main__":
    fig5_5()
