import time

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

from splineA import splineA


def fig5_19():
    # error when fitting with cubic splines and cheb points
    # integration rule is composite trap (see first edition, exercise 6.25)

    # determine exact values
    a = -1
    b = 1
    Ne = 8000
    xp = np.linspace(a, b, Ne)
    h = xp[1] - xp[0]
    R = np.zeros(Ne)
    for i in range(Ne):
        R[i] = f(xp[i])

    N = 7

    tic = time.perf_counter()
    # spline
    nx = 5
    num = np.zeros(N, dtype=int)
    eS = np.zeros(N)
    ef = np.zeros(N)
    for in_ in range(N):
        nx = 2 * nx
        num[in_] = nx
        xd = np.linspace(a, b, nx)
        yd = np.zeros(nx)
        for iy in range(nx):
            yd[iy] = f(xd[iy])
        xp = np.linspace(a, b, Ne)
        # natural spline
        ys = splineA(xd, yd, xp)
        eS[in_] = h * np.linalg.norm(ys - R, 1)
        # clamped spline
        yf = splineA(xd, yd, xp, [1, 1], [50 / 26 ** 2, -50 / 26 ** 2])
        ef[in_] = h * np.linalg.norm(yf - R, 1)
    print(time.perf_counter() - tic)

    tic = time.perf_counter()
    # cheb
    nx = 5
    eC = np.zeros(N)
    for in_ in range(N):
        nx = 2 * nx
        num[in_] = nx
        xd = np.linspace(a, b, nx)
        xxd = np.zeros(nx)
        yyd = np.zeros(nx)
        for ix in range(nx):
            cc = (2 * (ix + 1) - 1) * np.pi / (2 * nx)
            xxd[ix] = 0.5 * (a + b + (b - a) * np.cos(cc))
            yyd[ix] = f(xxd[ix])
        data = np.column_stack((xxd, yyd))
        yc = barylag(data, xp)
        eC[in_] = h * (
            np.linalg.norm(yc - R, 1)
            - 0.5 * abs(yc[0] - R[0])
            - 0.5 * abs(yc[Ne - 1] - R[Ne - 1])
        )
    print(time.perf_counter() - tic)

    plt.close("all")
    # get(gcf)
    fig, ax = plt.subplots(figsize=(9.0, 3.2))
    fig.subplots_adjust(left=0.11, right=0.98, bottom=0.22, top=0.97)
    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],
    ])
    ax.set_prop_cycle(cycler(color=co))

    ax.loglog(num, eS, "--or", linewidth=1.6, markersize=8)
    ax.loglog(num, eC, "--*b", linewidth=1.6, markersize=8)
    # ax.loglog(num, ef, "--dk", linewidth=1.6, markersize=8)
    ax.grid(True)
    ax.set_xlim(10, 1e3)
    ax.set_ylim(1e-16, 5e-1)
    ax.set_yticks([1e-16, 1e-11, 1e-6, 1e-1])
    ax.set_xlabel("n-axis")
    ax.set_ylabel("Error")

    ax.legend(
        [" Spline", " Cheb"],
        loc="lower left",
        prop={"size": 14, "weight": "bold"},
    )
    # ax.legend([" Spline", " Cheb", " Fixed"], loc="lower left",
    #           prop={"size": 14, "weight": "bold"})

    ax.tick_params(labelsize=14)
    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/splineE.eps')


def f(x):
    g = 1 / (1 + 25 * x ** 2)
    return g


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


def barylag(data, x):
    ############################################################################
    #
    # barylag.m
    #
    # Interpolates the given data using the Barycentric
    # Lagrange Interpolation formula. Vectorized to remove all loops
    #
    # data - a two column vector where column one contains the
    #        nodes and column two contains the function value
    #        at the nodes
    # p - interpolated data. Column one is just the
    #     fine mesh x, and column two is interpolated data
    #
    # Reference:
    #
    # (1) Jean-Paul Berrut & Lloyd N. Trefethen, "Barycentric Lagrange
    #     Interpolation"
    #     http://web.comlab.ox.ac.uk/oucl/work/nick.trefethen/berrut.ps.gz
    # (2) Walter Gaustschi, "Numerical Analysis, An Introduction" (1997)
    #     pp. 94-95
    #
    #
    # Written by: Greg von Winckel       03/07/04
    # Contact:    gregvw@chtm.unm.edu
    #
    ############################################################################

    M = len(data)
    N = len(x)

    # Compute the barycentric weights
    X = np.tile(data[:, 0].reshape(-1, 1), (1, M))

    # matrix of weights
    weights = 1.0 / np.prod(X - X.T + np.eye(M), axis=0)
    W = np.tile(weights.reshape(1, -1), (N, 1))

    # Get distances between nodes and interpolation points
    xdist = np.tile(np.asarray(x).reshape(-1, 1), (1, M)) - np.tile(
        data[:, 0].reshape(1, -1), (N, 1)
    )

    # Find all of the elements where the interpolation point is on a node
    fixi, fixj = np.where(xdist == 0)

    # Use NaNs as a place-holder
    xdist[fixi, fixj] = np.nan
    H = W / xdist

    # Compute the interpolated polynomial
    p = (H @ data[:, 1]) / np.sum(H, axis=1)

    # Replace NaNs with the given exact values.
    p[fixi] = data[fixj, 1]
    return p


if __name__ == "__main__":
    fig5_19()
