#!/usr/bin/env python3

# Lagrange Polynomial Interpolation
#----------------------------------

import numpy as np

def my_lagrange(data_x, data_y, x):
    """
    Lagrange Interpolation

    param @data_x   : array: x values to be interpolated
    param @data_y   : array: y values to be interpolated
    param @x        : array_like: x part of the interpolation set

    returns @p      : array: y part of the interpolation set
    """
    # make sure data_x is a numpy array with nd >= 1
    x = np.array(x, ndmin=1)
    p = np.zeros(np.size(x), dtype=np.float64)


    for i, x_i in enumerate( x ):
        p[i] = np.sum( data_y * my_lagrange_polynomials(data_x, x_i) )

    return p


def my_lagrange_polynomials(data_x, x_i):
    """ 
    Determine Lagrange Polynomials
    
    param @data_x   : array_like : various data parts
    param @x_i      : float : interpolating x value

    returns @L      : array : lagrange polynomial terms for each x_i
    """

    # make sure data_x is a numpy array with nd >= 1
    data_x = np.array(data_x, ndmin=1)

    L = np.zeros(np.size(data_x), dtype=np.float64)

    L_fixed = np.prod( x_i - data_x )

    for k, x_k in enumerate( data_x ):
        if x_k != x_i:
            L[k] = L_fixed / ( x_i - x_k )
        else:
            # Avoid dividing by zero if x_i and x_k are equal
            L[k] = np.prod( x_i - np.delete(data_x, k) )

        L[k] *=  1 / np.prod( x_k - np.delete(data_x,k) )

    return L


def main(x_values = None, y_values = None, x = None, x_min = None, x_max = None, x_step = None, x_N = None):
    """
    Example implementation of my_lagrange[_polynomials]()
    """

    if x_values is None or y_values is None:
        x_values = [2,3,4,5,6]
        y_values = [2,5,5,5,6]

    if x is None:
        x_min = x_min if x_min is not None else min(x_values)
        x_max = x_max if x_max is not None else max(y_values)
        
        x_N = x_N if x_N is not None else (int(abs((x_max - x_min ) / x_step)) if x_step is not None else 1e3)
        
        x = np.linspace(x_min, x_max, x_N, endpoint=True)
    
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots()

    ax.set_title("Lagrange Interpolations")
    ax.grid()
    ax.set_xlabel("x")
    ax.set_ylabel("y")

    # Plot my Lagrange Interpolation
    ax.plot(x, my_lagrange(x_values, y_values, x), '.', label="My Lagrange Interpolation")

    # Plot Scipy's Lagrange Interpolation
    from scipy.interpolate import lagrange as scipy_lagrange
    ax.plot(x, scipy_lagrange(x_values, y_values)(x), '--', label="Scipy Interpolation")


    # Plot the data
    ax.plot(x_values, y_values, 'gd', label="Data")
    ax.legend()

    plt.show()


if __name__ == "__main__":
    main()