"""Literal Python conversion of fig8_5.m."""

from pathlib import Path

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


def mat(x, n, m):
    g = np.zeros((n, m))
    for i in range(n):
        for j in range(m):
            g[i, j] = x[i] ** j
    return g


def nspline(xs, ys, x):
    """Direct SciPy replacement for the MATLAB nspline routine."""
    return CubicSpline(xs, ys, bc_type="natural")(x)


def normal_coefficients(xd, yd, m, lambda_=0.0):
    A = mat(xd, len(xd), m)
    matrix = A.T @ A + lambda_ * np.eye(m)
    coeff_ascending = (np.linalg.inv(matrix) @ A.T @ yd.reshape(-1, 1)).ravel()
    return np.fliplr(coeff_ascending.reshape(1, -1)).ravel(), A


def qr_coefficients(A, yd, m):
    Q, R = np.linalg.qr(A, mode="complete")
    Rm = R[:m, :m]
    Qy = Q.T @ yd.reshape(-1, 1)
    ym = Qy[:m, 0]
    coeff_ascending = (np.linalg.inv(Rm) @ ym.reshape(-1, 1)).ravel()
    return np.fliplr(coeff_ascending.reshape(1, -1)).ravel(), Q, R, Rm, Qy, ym


def bold_axes(ax, fontsize):
    ax.tick_params(labelsize=fontsize)
    for label in ax.get_xticklabels() + ax.get_yticklabels():
        label.set_fontweight("bold")


def fig8_5(icase=1, lambda_=1e-4):
    #  poly fit: normal vs QR, regularization, and overfitting plots
    #  also include figures 8.6, 8.7, and 8.8
    #  also inclues table 8.2

    # use the given training and testing data:  icase=1
    # use given training data, but create and save testing data set:  icase=2
    # create, and save, training and testing data:  icase=3

    print(f"icase = {icase}")
    print(f"lambda = {lambda_}")

    base = Path(__file__).resolve().parent

    if icase == 1:
        Xdata = np.loadtxt(base / "training_data.txt")
        xd = Xdata[:, 0].T; yd = Xdata[:, 1].T
        nd = len(xd)

        X2data = np.loadtxt(base / "testing_data.txt")
        xd2 = X2data[:, 0].T; yd2 = X2data[:, 1].T
        nd2 = len(xd2)

    elif icase == 2:
        Xdata = np.loadtxt(base / "training_data.txt")
        xd = Xdata[:, 0].T; yd = Xdata[:, 1].T
        nd = len(xd)

        #  generate testing data
        n = 12
        xs = np.linspace(0, 1, n)
        ys = np.array([2, 8, 8, 8, 3, 2, 1, 1, 1, 5, 5, 5], dtype=float)
        xd2 = np.sort(np.random.rand(nd))
        yd2 = nspline(xs, ys, xd2)
        dis = 1 + 0.5 * (np.random.rand(nd) - 0.5)
        yd2 = dis * yd2
        X2data = np.column_stack((xd2, yd2))
        np.savetxt(base / "testing_data.txt", X2data)
        nd2 = len(xd2)

    elif icase == 3:
        #  generate spline
        n = 12
        xs = np.linspace(0, 1, n)
        ys = np.array([2, 8, 8, 8, 3, 2, 1, 1, 1, 5, 5, 5], dtype=float)

        #  generate data for regression
        nd = 20
        xd = np.sort(np.random.rand(nd))
        yd = nspline(xs, ys, xd)
        dis = 1 + 0.2 * (np.random.rand(nd) - 0.5)
        yd = dis * yd
        Xdata = np.column_stack((xd, yd))
        np.savetxt(base / "training_data.txt", Xdata)

        #  generate testing data
        xd2 = np.sort(np.random.rand(nd))
        yd2 = nspline(xs, ys, xd2)
        dis = 1 + 0.5 * (np.random.rand(nd) - 0.5)
        yd2 = dis * yd2
        X2data = np.column_stack((xd2, yd2))
        np.savetxt(base / "testing_data.txt", X2data)
        nd2 = len(xd2)

    else:
        raise ValueError("icase must be 1, 2, or 3")

    # pseudo-inverse and Rm
    m = 4
    c4, A = normal_coefficients(xd, yd, m)
    r4, Q, R, Rm, Qy, ym = qr_coefficients(A, yd, m)
    condc4 = np.linalg.cond(A.T @ A, np.inf)
    condr4 = np.linalg.cond(Rm, np.inf)

    m = 8
    c8, A = normal_coefficients(xd, yd, m)
    r8, Q, R, Rm, Qy, ym = qr_coefficients(A, yd, m)
    condc8 = np.linalg.cond(A.T @ A, np.inf)
    condr8 = np.linalg.cond(Rm, np.inf)

    m = 12
    c12, A = normal_coefficients(xd, yd, m)
    r12, Q, R, Rm, Qy, ym = qr_coefficients(A, yd, m)
    condc12 = np.linalg.cond(A.T @ A, np.inf)
    condr12 = np.linalg.cond(Rm, np.inf)

    cR12, _ = normal_coefficients(xd, yd, 12, lambda_)
    condcR12 = np.linalg.cond(A.T @ A + lambda_ * np.eye(12), np.inf)

    m = 16
    A = mat(xd, nd, m)
    cc16 = (np.linalg.inv(A.T @ A) @ A.T @ yd.reshape(-1, 1)).ravel()
    c16 = np.fliplr(cc16.reshape(1, -1)).ravel()
    Q, R = np.linalg.qr(A, mode="complete")
    Rm = R[:m, :m]
    Qy = Q.T @ yd.reshape(-1, 1)
    ym = Qy[:m, 0]
    np.linalg.cond(Rm, np.inf)
    rr16 = (np.linalg.inv(Rm) @ ym.reshape(-1, 1)).ravel()
    r16 = np.fliplr(rr16.reshape(1, -1)).ravel()
    # The MATLAB source prints condc16 before assigning it.  The intended
    # normal-equation condition number is assigned here so the script runs.
    condc16 = np.linalg.cond(A.T @ A, np.inf)
    condr16 = np.linalg.cond(Rm, np.inf)

    print(cc16[:4])
    print(rr16[:4])

    cR16, _ = normal_coefficients(xd, yd, 16, lambda_)
    condcR16 = np.linalg.cond(A.T @ A + lambda_ * np.eye(16), np.inf)

    #  nn = number of points plotted
    nn = 200
    X = np.linspace(xd[0], xd[nd - 1], nn)
    Y4 = np.polyval(c4, X)
    R4 = np.polyval(r4, X)
    Y8 = np.polyval(c8, X)
    R8 = np.polyval(r8, X)
    Y12 = np.polyval(c12, X)
    R12 = np.polyval(r12, X)
    Y16 = np.polyval(c16, X)
    R16 = np.polyval(r16, X)

    YR12 = np.polyval(cR12, X)
    YR16 = np.polyval(cR16, X)

    # condition numbers
    print(f"{4:d} &  {condc4:5.2e}  & {condr4:5.2e}")
    print(f"{8:d} &  {condc8:5.2e}  & {condr8:5.2e}")
    print(f"{12:d} &  {condc12:5.2e}  & {condr12:5.2e}")
    print(f"{16:d} &  {condc16:5.2e}  & {condr16:5.2e}")

    plt.close("all")
    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],
    ])

    # get(gcf)
    # MATLAB Position [27 978 666 367], scaled to fit the plot window.
    fig1, axes = plt.subplots(2, 2, figsize=(6.66, 4.8), layout="constrained")
    marker = 18
    m_size = 14
    xc = 0.12; yc = 0

    # plot interpolation fn's and exact result
    ax = axes[0, 0]
    ax.set_prop_cycle(color=co)
    ax.plot(X, R4, "r", linewidth=1.6)
    ax.plot(X, Y4, color=(0, 0.5, 0), linewidth=1.6)
    ax.text(xc, yc, "m = 4", fontsize=m_size, fontweight="bold")
    ax.plot(xd, yd, ".b", markersize=marker)
    ax.axis([0, 1, -3, 12])
    # ax.set_xlabel('x-axis')
    ax.set_ylabel("y-axis", fontsize=12, fontweight="bold")
    ax.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
    ax.set_yticks([0, 6, 12])
    ax.grid(True)
    bold_axes(ax, 12)

    ax = axes[0, 1]
    ax.set_prop_cycle(color=co)
    ax.plot(X, R8, "r", linewidth=1.6)
    ax.plot(X, Y8, color=(0, 0.5, 0), linewidth=1.6)
    ax.plot(xd, yd, ".b", markersize=marker)
    ax.text(xc, yc, "m = 8", fontsize=m_size, fontweight="bold")
    # ax.set_xlabel('x-axis')
    ax.axis([0, 1, -3, 12])
    ax.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
    ax.set_yticks([0, 6, 12])
    ax.grid(True)
    bold_axes(ax, 12)

    ax = axes[1, 0]
    ax.set_prop_cycle(color=co)
    ax.plot(X, R12, "r", linewidth=1.6)
    ax.plot(X, Y12, color=(0, 0.5, 0), linewidth=1.6)
    ax.plot(xd, yd, ".b", markersize=marker)
    ax.text(xc, yc, "m = 12", fontsize=m_size, fontweight="bold")
    ax.set_xlabel("x-axis", fontsize=12, fontweight="bold")
    ax.axis([0, 1, -3, 12])
    ax.set_ylabel("y-axis", fontsize=12, fontweight="bold")
    ax.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
    ax.set_yticks([0, 6, 12])
    ax.grid(True)
    bold_axes(ax, 12)

    ax = axes[1, 1]
    ax.set_prop_cycle(color=co)
    ax.plot(X, R16, "r", linewidth=1.6)
    ax.plot(X, Y16, color=(0, 0.5, 0), linewidth=1.6)
    ax.plot(xd, yd, ".b", markersize=marker)
    ax.text(xc, yc, "m = 16", fontsize=m_size, fontweight="bold")
    ax.set_xlabel("x-axis", fontsize=12, fontweight="bold")
    ax.axis([0, 1, -3, 12])
    ax.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
    ax.set_yticks([0, 6, 12])
    ax.grid(True)
    bold_axes(ax, 12)

    # fig1.savefig('/Users/mark/Desktop/p1.eps')

    ELS_T = np.zeros(15)
    ELS_R = np.zeros(15)
    for im in range(2, 17):
        m = im
        A = mat(xd, nd, m)
        cR_ascending = (
            np.linalg.inv(A.T @ A + lambda_ * np.eye(im)) @ A.T @ yd.reshape(-1, 1)
        ).ravel()
        cR = np.fliplr(cR_ascending.reshape(1, -1)).ravel()
        ELS_T[im - 2] = np.linalg.norm(yd - np.polyval(cR, xd), 2) ** 2
        ELS_R[im - 2] = np.linalg.norm(yd2 - np.polyval(cR, xd2), 2) ** 2

        # check on differences in x^{12} coefficient
        if im == 12 or im == 16:
            print(cR)
            print(cR[m - 12])

    # error for training set
    E = np.zeros(15)
    for im in range(1, 16):
        E[im - 1] = np.linalg.norm(yd - np.polyval(np.polyfit(xd, yd, im), xd), 2) ** 2

    # get(gcf)
    # MATLAB Position [27 674 658 230], scaled to fit the plot window.
    fig2, ax2 = plt.subplots(figsize=(6.58, 3.0), layout="constrained")
    ax2.plot(range(2, 17), ELS_T, "--ob", linewidth=1.6, markersize=8)
    ax2.plot(range(2, 17), ELS_R, "--or", linewidth=1.6, markersize=8)
    ax2.legend(
        [" Training", " Testing"],
        loc="upper center",
        prop={"size": 14, "weight": "bold"},
    )
    ax2.set_xlabel("m-axis", fontsize=14, fontweight="bold")
    ax2.set_ylabel("Error", fontsize=14, fontweight="bold")
    ax2.grid(True)
    bold_axes(ax2, 14)
    # fig2.savefig('/Users/mark/Desktop/p2.eps')

    # error for training set
    for im in range(1, 16):
        E[im - 1] = np.linalg.norm(yd - np.polyval(np.polyfit(xd, yd, im), xd), 2) ** 2

    # error for testing set
    EE = np.zeros(15)
    for im in range(1, 16):
        EE[im - 1] = np.linalg.norm(yd2 - np.polyval(np.polyfit(xd, yd, im), xd2), 2) ** 2
    marker = 20

    # MATLAB Position [27 232 666 367], scaled to fit the plot window.
    fig3, axes3 = plt.subplots(2, 1, figsize=(6.66, 5.0), layout="constrained")
    ax = axes3[0]
    ax.plot(xd, yd, ".b", markersize=20)
    ax.plot(xd2, yd2, "+r", markersize=12, linewidth=2)
    ax.legend(
        [" Training Data", " Testing Data"],
        loc="upper right",
        prop={"size": 14, "weight": "bold"},
    )
    ax.set_xlabel("x-axis", fontsize=12, fontweight="bold")
    ax.set_ylabel("y-axis", fontsize=12, fontweight="bold")
    ax.grid(True)
    bold_axes(ax, 12)

    ax = axes3[1]
    ax.plot(range(2, 17), E, "-ob", linewidth=1.6, markersize=6)
    ax.plot(range(2, 17), EE, "-+r", linewidth=1.6, markersize=10)
    ax.legend(
        [" Training", " Testing"],
        loc="upper right",
        prop={"size": 14, "weight": "bold"},
    )
    ax.axis([2, 16, 0, 200])
    ax.set_xlabel("m-axis", fontsize=12, fontweight="bold")
    ax.set_ylabel("Error", fontsize=12, fontweight="bold")
    # ax.set_xticks([0,0.2,0.4,0.6,0.8,1.0])
    ax.grid(True)
    bold_axes(ax, 12)
    # fig3.savefig('/Users/mark/Desktop/p3.eps')

    # get(gcf)
    # MATLAB Position [695 1115 658 230], scaled to fit the plot window.
    fig4, axes4 = plt.subplots(1, 2, figsize=(6.58, 3.2), layout="constrained")
    marker = 18
    m_size = 14
    xc = 0.12; yc = 0

    # plot interpolation fn's and exact result
    ax = axes4[0]
    ax.set_prop_cycle(color=co)
    ax.plot(X, YR12, "k", linewidth=1.6)
    ax.plot(X, R12, "r", linewidth=1.6)
    ax.text(xc, yc, "m = 12", fontsize=m_size, fontweight="bold")
    ax.plot(xd, yd, ".b", markersize=marker)
    ax.axis([0, 1, -3, 12])
    ax.set_xlabel("x-axis", fontsize=12, fontweight="bold")
    ax.set_ylabel("y-axis", fontsize=12, fontweight="bold")
    ax.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
    ax.set_yticks([0, 6, 12])
    ax.grid(True)
    bold_axes(ax, 12)

    ax = axes4[1]
    ax.set_prop_cycle(color=co)
    ax.plot(X, YR16, "k", linewidth=1.6)
    ax.plot(X, R16, "r", linewidth=1.6)
    ax.text(xc, yc, "m = 16", fontsize=m_size, fontweight="bold")
    ax.plot(xd, yd, ".b", markersize=marker)
    ax.axis([0, 1, -3, 12])
    ax.set_xlabel("x-axis", fontsize=12, fontweight="bold")
    # ax.set_ylabel('y-axis')
    ax.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
    ax.set_yticks([0, 6, 12])
    ax.grid(True)
    bold_axes(ax, 12)

    # fig4.savefig('/Users/mark/Desktop/p4.eps')
    plt.show()


if __name__ == "__main__":
    fig8_5()
