import numpy as np
import matplotlib.pyplot as plt
from cycler import cycler
from scipy.interpolate import CubicSpline

from splineA import splineA


def fig5_13():
    # fit cubic splines to the four data sets

    ### this requires the subaxis and parseArgs files
    ### and it requires the file splineA.m

    ##### generate the data sets
    a = -1
    b = 1
    nx = 12
    xd = np.linspace(a, b, nx)

    # 5th degree poly
    yd4 = np.zeros(nx)
    for iy in range(nx):
        yd4[iy] = (
            (xd[iy] + 0.9)
            * (xd[iy] + 0.1) ** 2
            * (xd[iy] - 0.2)
            * (xd[iy] - 0.8)
        )

    # jump function
    yd1 = np.zeros(nx)
    for iy in range(nx):
        yd1[iy] = (-1) ** (iy + 1) * 0.02
    yd1[4] = 1
    yd1[5] = 1
    yd1[6] = 1
    yd1[7] = 1

    # oscillatory function
    yd2 = np.zeros(nx)
    for iy in range(nx):
        yd2[iy] = (-1) ** (iy + 1)

    # oscillatory function
    yd3 = np.zeros(nx)
    for iy in range(nx):
        yd3[iy] = np.sqrt(1 - xd[iy] ** 2)

    #### evaluate and then plot spline function

    # ysM = MATLAB spline function
    # ys = natural spline (requires splineA.m file)

    n = 400
    xp = np.linspace(a, b, n)

    plt.close("all")
    # get(gcf)
    fig, axes = plt.subplots(2, 2, figsize=(8.8, 4.8))
    fig.subplots_adjust(left=0.09, right=0.98, bottom=0.13, top=0.99, wspace=0.18, hspace=0.18)
    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.flat:
        ax.set_prop_cycle(cycler(color=co))

    # subplot(2,2,1)
    ax = axes[0, 0]
    ys = splineA(xd, yd1, xp)
    ysM = CubicSpline(xd, yd1)(xp)
    ax.plot(xd, yd1, "or", markersize=7, linewidth=2)
    ax.plot(xp, ys, "b", linewidth=1.4)
    ax.plot(xp, ysM, "--k", linewidth=1.4)
    ax.set_xlim(-1.1, 1.1)
    ax.set_ylim(-0.2, 1.2)
    ax.set_yticks([0, 0.5, 1])
    ax.grid(True)
    ax.set_xlabel("x-axis")
    ax.tick_params(labelsize=12)

    # subplot(2,2,2)
    ax = axes[0, 1]
    ys2 = splineA(xd, yd2, xp)
    ysM2 = CubicSpline(xd, yd2)(xp)
    ax.plot(xd, yd2, "or", markersize=7, linewidth=2)
    ax.plot(xp, ys2, "b", linewidth=1.4)
    ax.plot(xp, ysM2, "--k", linewidth=1.4)
    ax.set_xlim(-1.1, 1.1)
    ax.set_ylim(-1.5, 1.5)
    ax.set_yticks([-1, 0, 1])
    ax.grid(True)
    ax.set_xlabel("x-axis")
    ax.set_ylabel("y-axis")
    ax.tick_params(labelsize=12)

    # subplot(2,2,3)
    ax = axes[1, 0]
    ys3 = splineA(xd, yd3, xp)
    ysM3 = CubicSpline(xd, yd3)(xp)
    ax.plot(xd, yd3, "or", markersize=7, linewidth=2)
    ax.plot(xp, ys3, "b", linewidth=1.4)
    ax.plot(xp, ysM3, "--k", linewidth=1.4)
    ax.set_xlim(-1.1, 1.1)
    ax.set_ylim(-0.1, 1.1)
    ax.set_yticks([0, 0.5, 1])
    ax.grid(True)
    ax.tick_params(labelsize=12)

    # subplot(2,2,4)
    ax = axes[1, 1]
    ys4 = splineA(xd, yd4, xp)
    ysM4 = CubicSpline(xd, yd4)(xp)
    ax.plot(xd, yd4, "or", markersize=7, linewidth=2)
    ax.plot(xp, ys4, "b", linewidth=1.4)
    ax.plot(xp, ysM4, "--k", linewidth=1.4)
    ax.set_xlim(-1.1, 1.1)
    ax.set_ylim(-0.2, 0.4)
    ax.set_yticks([-0.2, 0, 0.2, 0.4])
    ax.grid(True)
    ax.set_ylabel("y-axis")
    ax.tick_params(labelsize=12)

    for ax in axes.flat:
        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()


if __name__ == "__main__":
    fig5_13()
