import math
import numpy as np
import matplotlib.pyplot as plt


def f(x):
    """
    Function used to generate the interpolation data.
    """
    return math.sin(x)


def polynomial_interp(xx, xi, yi, npts):
    """
    Local polynomial interpolation of selected order.

    The interpolating polynomial is evaluated in Lagrange form using
    npts neighboring data points. The interpolation order is npts - 1.

    Input:
        xx   = abscissa at which the interpolation is evaluated
        xi   = array of data abscissas
        yi   = array of data ordinates
        npts = number of points used for interpolation
               (interpolation order = npts - 1)

    Output:
        interpolated value at xx

    Comments:
        The routine works for both equally and unequally spaced xi.
        A binary search locates the interval containing xx, and a local
        group of neighboring data points is selected for interpolation.
        If xx lies outside the data interval, the nearest boundary value
        is returned.

    Original implementation by Alexander Godunov, January 2010.
    Python version revised for the companion website, 2026.
    """

    ni = len(xi)

    # Check the number of interpolation points
    nuse = npts
    if nuse > ni:
        nuse = ni

    # If xx is outside the xi interval, return a boundary value
    if xx <= xi[0]:
        return yi[0]

    if xx >= xi[ni - 1]:
        return yi[ni - 1]

    # Binary search to find i such that xi[i] < xx < xi[i+1]
    i = 0
    j = ni - 1

    while j > i + 1:
        k = (i + j) // 2

        if xx < xi[k]:
            j = k
        else:
            i = k

    # Shift i so that xx lies near the middle of the selected data points
    i = i + 1 - nuse // 2

    # Keep the interpolation points inside the available data range
    if i < 0:
        i = 0

    if i + nuse - 1 > ni - 1:
        i = ni - nuse

    # Evaluate the interpolating polynomial in Lagrange form
    y = 0.0

    for js in range(i, i + nuse):
        lam = 1.0

        for jl in range(i, i + nuse):
            if jl != js:
                lam *= (xx - xi[jl]) / (xi[js] - xi[jl])

        y += yi[js] * lam

    return y


def main():
    """
    General polynomial interpolation

    Function values f(x) are calculated at n base points. A local
    polynomial of selected order is then used to interpolate the
    function at nint points across the interval. The program also
    calculates the average absolute interpolation error.

    The interpolating polynomial is evaluated in Lagrange form.
    For order = 1, 2, 3, ... the routine uses 2, 3, 4, ...
    neighboring data points, respectively.

    Original implementation by Alexander Godunov, January 2010.
    Python version revised for the companion website, 2026.
    """

    n = 11       # base points for interpolation
    nint = 21    # compute interpolation in nint points

    xmin = 0.0
    xmax = 2.0

    xi = np.zeros(n, dtype=float)
    yi = np.zeros(n, dtype=float)

    # Step 1: generate xi and yi from f(x), xmin, xmax, and n
    step = (xmax - xmin) / (n - 1)

    for i in range(n):
        xi[i] = xmin + step * i
        yi[i] = f(xi[i])

    # Step 2: choose the interpolation order
    # order = 1 -> linear
    # order = 2 -> quadratic
    # order = 3 -> cubic
    order = 3

    print("    General Polynomial Interpolation")
    print(f"     order of interpolation = {order:2d}")

    # Step 3: evaluate the interpolant at nint points
    errav = 0.0
    step = (xmax - xmin) / (nint - 1)

    xplot = np.zeros(nint, dtype=float)
    ysplot = np.zeros(nint, dtype=float)

    print(f"{'x':>12}{'interpolated':>16}{'error':>12}")

    for i in range(nint):
        x = xmin + step * i
        y = f(x)

        ys = polynomial_interp(x, xi, yi, order + 1)

        error = ys - y

        print(f"{x:12.5f}{ys:16.5f}{error:12.5f}")

        # Step 4: calculate the average absolute interpolation error
        errav += abs(y - ys) / nint

        xplot[i] = x
        ysplot[i] = ys

    print(f"{'Average error':>28}{errav:12.5f}")

    # Plot the exact function, interpolation points, and interpolated curve.
    xfine = np.linspace(xmin, xmax, 400)
    yfine = np.sin(xfine)

    # Evaluate the interpolant on a fine grid for a smooth plotted curve.
    yinterp_fine = np.zeros_like(xfine)

    for i in range(len(xfine)):
        yinterp_fine[i] = polynomial_interp(
            xfine[i], xi, yi, order + 1
        )

    plt.figure()
    plt.plot(xfine, yfine, "-", linewidth=1.2, label="Exact function")
    plt.plot(
        xfine,
        yinterp_fine,
        "--",
        linewidth=1.2,
        label="Interpolated curve",
    )
    plt.plot(
        xi,
        yi,
        "o",
        markersize=6,
        label="Interpolation points",
    )

    plt.xlabel("x")
    plt.ylabel("f(x)")
    plt.title(f"Polynomial interpolation, order = {order}")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()