import numpy as np


def splineA(xd, yd, x, Ends=None, Ders=None):
    # Nonuniform Cubic Spline Interpolation
    # Includes any combination of first or second derivative end conditions
    # If you want the formulas for the resulting cubics, use splineB.m

    # Usage (examples at end of file)
    # For a natural cubic spline (this is the default version)
    #   y = splineA(xd, yd, x)
    # If one or both ends involve a derivative condition, then
    #   y = splineA(xd, yd, x, Ends, Ders)

    # Output:
    # If n = length of x, then y is an n-vector containing the values of the
    # cubic spline at the respective x locations.

    # Required arguments:
    #   xd, yd: data vectors with xd[i] < xd[i+1]
    #   x: point or points where y(x) is to be evaluated

    # Optional arguments (only needed if a derivative end condition is used):
    #   Ends = [a, b]: this specifies which end conditions to use
    #       a = 1: the first derivative is specified at left end
    #       a = 2: the second derivative is specified at left end
    #       b = 1: the first derivative is specified at right end
    #       b = 2: the second derivative is specified at right end
    #   Ders = [c, d]: this specifies the derivative values
    #       c = value of 1st, or 2nd, derivative at left end
    #       d = value of 1st, or 2nd, derivative at right end
    #       This vector is required if Ends is included in the calling command.

    xd = np.asarray(xd, dtype=float).reshape(-1)
    yd = np.asarray(yd, dtype=float).reshape(-1).copy()

    scalar_x = np.ndim(x) == 0
    x = np.atleast_1d(np.asarray(x, dtype=float))

    # Check for errors.
    n = len(xd)
    nd = len(yd)

    if (Ends is None) != (Ders is None):
        raise ValueError("You are missing one of the required arguments")

    if n < 2:
        raise ValueError("You need at least three data points")

    if n != nd:
        raise ValueError("xd and yd must have the same length")

    if n != len(np.unique(xd)):
        raise ValueError("the xd values can not repeat")

    # Set the defaults.
    if Ends is None:
        Ends = np.array([2, 2], dtype=int)
        Ders = np.array([0.0, 0.0])
    else:
        Ends = np.asarray(Ends, dtype=int).reshape(-1)
        Ders = np.asarray(Ders, dtype=float).reshape(-1)

        if len(Ends) != 2 or len(Ders) != 2:
            raise ValueError("Ends and Ders must each contain two values")

    # The code begins.
    # Extend nodes for B_0(x) and B_{n+1}(x).
    h = np.diff(xd)
    h1 = h[0]
    hn = h[n - 2]

    h = np.concatenate(([h1, h1, h1], h, [hn, hn, hn]))
    Xd = np.concatenate(
        (
            [xd[0] - 3 * h1, xd[0] - 2 * h1, xd[0] - h1],
            xd,
            [xd[n - 1] + hn, xd[n - 1] + 2 * hn, xd[n - 1] + 3 * hn],
        )
    )

    # Compute the coefficients for the B_i's.
    B, G = coeffs(h)

    # Interpolate data.
    Aa = np.zeros(n)
    Ab = np.zeros(n)
    Ac = np.zeros(n)

    if Ends[0] == 2:
        alpha0 = G[3, 0] / G[1, 0]
        beta0 = G[5, 0] / G[1, 0]
        gamma0 = Ders[0] / G[1, 0]
    elif Ends[0] == 1:
        alpha0 = G[2, 0] / G[0, 0]
        beta0 = G[4, 0] / G[0, 0]
        gamma0 = Ders[0] / G[0, 0]
    else:
        raise ValueError("Ends values are incorrect")

    B0 = Bi(xd[0], 0, B, Xd)
    Aa[0] = 2 / 3 - B0 * alpha0
    Ac[0] = Bi(xd[0], 2, B, Xd) - B0 * beta0
    Ab[0] = 0
    yd[0] = yd[0] - B0 * gamma0

    for ii in range(1, n - 1):
        Aa[ii] = 2 / 3
        Ac[ii] = Bi(xd[ii], ii + 2, B, Xd)
        Ab[ii] = Bi(xd[ii], ii, B, Xd)

    if Ends[1] == 2:
        alphan = G[1, 1] / G[5, 1]
        betan = G[3, 1] / G[5, 1]

        # This uses Ders[0], matching the supplied MATLAB code.
        gamman = Ders[0] / G[5, 1]
    elif Ends[1] == 1:
        alphan = G[0, 1] / G[4, 1]
        betan = G[2, 1] / G[4, 1]
        gamman = Ders[1] / G[4, 1]
    else:
        raise ValueError("End values are incorrect")

    Bn1 = Bi(xd[n - 1], n + 1, B, Xd)
    Aa[n - 1] = 2 / 3 - Bn1 * betan
    Ac[n - 1] = 0
    Ab[n - 1] = Bi(xd[n - 1], n - 1, B, Xd) - Bn1 * alphan
    yd[n - 1] = yd[n - 1] - gamman * Bn1

    w = tri(Aa, Ab, Ac, yd)
    a0 = -alpha0 * w[0] - beta0 * w[1] + gamma0
    anp1 = -alphan * w[n - 2] - betan * w[n - 1] + gamman
    w = np.concatenate(([a0], w, [anp1]))

    # Evaluate interpolation function.
    nx = len(x)
    y = np.zeros(nx)

    for ix in range(nx):
        sum = 0.0

        for k in range(n + 2):
            sum = sum + w[k] * Bi(x[ix], k, B, Xd)

        y[ix] = sum

    if scalar_x:
        return float(y[0])

    return y


# Local functions used by algorithm.

# Tridiagonal solver.
def tri(a, b, c, f):
    N = len(f)
    v = np.zeros(N)
    y = np.zeros(N)

    w = a[0]
    y[0] = f[0] / w

    for i in range(1, N):
        v[i - 1] = c[i - 1] / w
        w = a[i] - b[i] * v[i - 1]
        y[i] = (f[i] - b[i] * y[i - 1]) / w

    for j in range(N - 2, -1, -1):
        y[j] = y[j] - v[j] * y[j + 1]

    return y


def coeffs(h):
    # Compute the values of the coefficients of the B_i's.
    # B = (n+2)x16 [16 = (4 subintervals for B_i)*(4 coeffs a,b,c,d)]
    # G = 6x2
    #   1st col = B_0'(x1), B0''(x1), B_1'(x1), B1''(x1),
    #             B_2'(x1), B2''(x1)
    #   2nd col = B_n-1'(xn), B_n-1''(xn), B_n-1'(xn), B_n''(xn),
    #             B_n+1'(xn), B_n+1''(xn)
    # N = n+5; number of B_i's = n+2.
    N = len(h)
    B = np.zeros((N - 3, 16))
    G = np.zeros((6, 2))

    # i = j - 1 in the MATLAB indexing.
    for j in range(N - 3):
        hm2 = h[j]
        hm1 = h[j + 1]
        h1 = h[j + 2]
        h2 = h[j + 3]

        # Find d_{i-2} and d_{i+2}.
        d2 = solver(hm2, hm1, h1, h2)

        # Interval x_{i-2} < x < x_{i-1}.
        B[j, 3] = d2[0]

        # Interval x_{i-1} < x < x_i.
        B[j, 4] = d2[0] * hm2 ** 3
        B[j, 5] = 3 * d2[0] * hm2 ** 2
        B[j, 6] = 3 * d2[0] * hm2
        Tm3 = (hm2 + hm1) ** 3 - hm1 ** 3
        B[j, 7] = (2 / 3 - Tm3 * d2[0]) / hm1 ** 3

        # Interval x_i < x < x_{i+1}.
        B[j, 8] = -d2[1] * h2 ** 3
        B[j, 9] = 3 * d2[1] * h2 ** 2
        B[j, 10] = -3 * d2[1] * h2
        T3 = (h2 + h1) ** 3 - h1 ** 3
        B[j, 11] = (-2 / 3 - T3 * d2[1]) / h1 ** 3

        # Interval x_{i+1} < x < x_{i+2}.
        B[j, 15] = d2[1]

        # B_0.
        if j == 0:
            G[0, 0] = 3 * d2[1] * h2 ** 2
            G[1, 0] = -6 * d2[1] * h2

        # B_1.
        if j == 1:
            G[2, 0] = (
                3 * d2[1] * h2 * (h2 + 2 * h1)
                + 3 * B[j, 11] * h1 ** 2
            )
            G[3, 0] = -6 * (d2[1] * h2 + B[j, 11] * h1)

        # B_2.
        if j == 2:
            G[4, 0] = 3 * d2[0] * hm2 ** 2
            G[5, 0] = 6 * d2[0] * hm2

        # B_{n-1}.
        if j == N - 6:
            G[0, 1] = 3 * d2[1] * h2 ** 2
            G[1, 1] = -6 * d2[1] * h2

        # B_n.
        if j == N - 5:
            G[2, 1] = (
                3 * d2[1] * h2 * (h2 + 2 * h1)
                + 3 * B[j, 11] * h1 ** 2
            )
            G[3, 1] = -6 * (d2[1] * h2 + B[j, 11] * h1)

        # B_{n+1}.
        if j == N - 4:
            G[4, 1] = 3 * d2[0] * hm2 ** 2
            G[5, 1] = 6 * d2[0] * hm2

    return B, G


# Find d_{i-2} and d_{i+2}.
def solver(hm2, hm1, h1, h2):
    a = h1 * hm2 * (hm2 + hm1) ** 2
    b = -hm1 * h2 * (h2 + h1) ** 2
    c = h1 ** 2 * hm2 * (hm2 + hm1) * (hm2 + 2 * hm1)
    d = hm1 ** 2 * h2 * (h2 + h1) * (h2 + 2 * h1)

    A = np.array([[a, b], [c, d]])
    bb = np.array(
        [
            2 * (hm1 + h1) / 3,
            -2 * (hm1 ** 2 - h1 ** 2) / 3,
        ]
    )

    d2 = np.linalg.solve(A, bb)
    return d2


# Evaluate B_i(x).
def Bi(x, i, B, Xd):
    j = i

    if x < Xd[j] or x > Xd[j + 4]:
        g = 0.0
    elif x < Xd[j + 1]:
        g = B[j, 3] * (x - Xd[j]) ** 3
    elif x < Xd[j + 2]:
        g = (
            B[j, 4]
            + B[j, 5] * (x - Xd[j + 1])
            + B[j, 6] * (x - Xd[j + 1]) ** 2
            + B[j, 7] * (x - Xd[j + 1]) ** 3
        )
    elif x < Xd[j + 3]:
        g = (
            B[j, 8]
            + B[j, 9] * (x - Xd[j + 3])
            + B[j, 10] * (x - Xd[j + 3]) ** 2
            + B[j, 11] * (x - Xd[j + 3]) ** 3
        )
    else:
        g = B[j, 15] * (x - Xd[j + 4]) ** 3

    return g


# Examples:
#
# Generate some random data.
# n = 8
# xd = np.random.rand(n) + np.arange(n)
# yd = np.random.randint(1, n + 1, n)
# x = np.linspace(xd[0], xd[n - 1], 100)
#
# Natural spline with plot.
# y = splineA(xd, yd, x)
#
# import matplotlib.pyplot as plt
#
# plt.plot(x, y, "r", linewidth=1.6)
# plt.plot(xd, yd, "ob", markersize=8)
# plt.show()
#
# Clamped spline.
# y = splineA(xd, yd, x, [1, 1], [0.2, -1])
#
# Combination examples.
# y = splineA(xd, yd, x, [1, 2], [-0.2, 0.6])
# y = splineA(xd, yd, x, [2, 1], [-0.2, 0.6])
#
# Check on a particular value.
# y0 = splineA(xd, yd, np.pi)
# y0 = splineA(xd, yd, 7.1, [2, 1], [-0.2, 0.6])

# Comments:
# 1) The nodes xd do not need to be uniformly spaced, but it is required
#    that they are sequential, so xd[i] < xd[i+1].
# 2) This program replaces, and extends, two earlier spline programs that I
#    wrote: nspline.m and fspline.m.
# 3) The code is reasonably fast, and as an example, with len(xd)=200
#    and len(x)=1000, the time in the original MATLAB code is about 8 msec.
# 4) The algorithm uses a B-spline construction, and it is described in the
#    text "Introduction to Scientific Computing and Data Analysis"
#    (Holmes, 2016).

# Versions:
# version 1.0: May 3, 2018
# version 1.01: May 6, 2018 (made a few minor changes in the comments;
# the code was not changed)
