import copy
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from prettytable import PrettyTable

def get_interpolation_lagrange_polynomial(list_x, list_y):

    x = sp.symbols('x')

    expr = 1
    for x_i in list_x:
        expr = expr * (x - x_i)

    func_l = 0
    for x_i, y_i in zip(list_x, list_y):
        func_l = func_l + y_i * (expr / (x - x_i)) / (expr / (x - x_i)).subs(x, x_i)

    return sp.lambdify(x, func_l, "numpy")


def table_finite_difference(x, y):

    n = len(y)
    diff = [x, y]
    for diff_num in range(n - 1):
        diff.append(['' for _ in range(n)])
        for i in range(1, n - diff_num):
            diff[diff_num + 2][i - 1] = diff[diff_num + 1][i] - diff[diff_num + 1][i - 1]

    return diff


def divided_difference(X, Y):

    # dif_diff(x0..xn) = sum(f(xj) / П(i!=j)(xj - xi), j = 1..n)
    x = sp.symbols('x')

    expr = 1
    for xi in X:
        expr = expr * (x - xi)

    diff = 0
    for xi, yi in zip(X, Y):
        diff += yi / (expr / (x - xi)).subs(x, xi)

    return diff


def table_divided_difference(x, y):

    n = len(y)
    diff = [x, y]
    for order in range(1, n):
        diff.append(['' for _ in range(n)])  # empty list with n elements
        for i in range(n - order):
            diff[order + 1][i] = divided_difference(x[i:i + order + 1], y[i:i + order + 1])  # + 1 last element

    return diff


def newton_polynomial(div_diffs):

    x = sp.symbols('x')

    expr = 1
    n = len(div_diffs)  # len(x) + 2
    N = 0
    for i in range(1, n):  # div_diffs[0] = x_vector
        N += div_diffs[i][0] * expr
        expr = expr * (x - div_diffs[0][i - 1])

    return sp.lambdify(x, N, "numpy")


############################################################################################################

def linear_interpolation(list_x, list_y):

    n = len(list_x)
    a = np.empty(n - 1)
    b = np.empty(n - 1)

    for i in range(n - 1):

        # F(x) = a_i * x + b_i
        a[i] = (list_y[i + 1] - list_y[i]) / (list_x[i + 1] - list_x[i])
        b[i] = list_y[i] - a[i] * list_x[i]

        space_x = np.linspace(list_x[i], list_x[i + 1])
        plt.plot(space_x, a[i] * space_x +           b[i])

    plt.scatter(list_x, list_y)
    plt.xlabel("Axis: X")
    plt.ylabel("Axis: Y")
    plt.title("Linear Interpolation")
    plt.show()

    return a, b


def quadratic_interpolation(list_x, list_y):

    n = len(list_x)
    a = np.empty(n - 1)
    b = np.empty(n - 1)
    c = np.empty(n - 1)

    for i in range(n - 2):

        # F(x) = a_i * x + b_i + c_i * x^2
        a[i] = (list_y[i + 2] - list_y[i]) / (list_x[i + 2] - list_x[i]) / (list_x[i + 2] - list_x[i + 1])  \
             - (list_y[i + 1] - list_y[i]) / (list_x[i + 1] - list_x[i]) / (list_x[i + 2] - list_x[i + 1])

        b[i] = (list_y[i + 1] - list_y[i]) / (list_x[i + 1] - list_x[i]) - a[i] * (list_x[i + 1] + list_x[i])

        c[i] = list_y[i] - b[i] * list_x[i] - a[i] * list_x[i] ** 2

    for i in range(0, n - 2, 2):
        space_x = np.linspace(list_x[i], list_x[i + 2])
        plt.plot(space_x, a[i] * space_x ** 2 + b[i] * space_x + c[i])

    curr_len = len(list_x)
    if curr_len % 2 == 0:
        i = curr_len - 2
        ax = (list_y[i + 1] - list_y[i]) / (list_x[i + 1] - list_x[i])
        bx = list_y[i] - ax * list_x[i]

        space_x = np.linspace(list_x[i], list_x[i + 1])
        plt.plot(space_x, ax * space_x + bx)


    plt.scatter(list_x, list_y)
    plt.xlabel("Axis: X")
    plt.ylabel("Axis: Y")
    plt.title("Quadratic Interpolation")
    plt.show()

    return a, b, c


def cubic_interpolation(list_x, list_y):

    n = len(list_x)

    h_i = np.array([list_x[i] - list_x[i - 1] for i in range(1, n)])
    l_i = np.array([(list_y[i] - list_y[i - 1]) / h_i[i - 1] for i in range(1, n)])
    delta_i = np.empty(n - 2, float)
    lambda_i = np.empty(n - 2, float)

    delta_i[0] = -0.5 * h_i[1] / (h_i[0] + h_i[1])
    lambda_i[0] = 1.5 * (l_i[1] - l_i[0]) / (h_i[0] + h_i[1])

    for i in range(1, n - 2):
        delta_i[i] = - h_i[i + 1] / (2 * h_i[i] + 2 * h_i[i + 1] + h_i[i] * delta_i[i - 1])
        lambda_i[i] = (2 * l_i[i + 1] - 3 * l_i[i] - h_i[i] * lambda_i[i - 1]) / \
                      (2 * h_i[i] + 2 * h_i[i + 1] + h_i[i] * delta_i[i - 1])

    a = np.array(copy.copy(list_y)[1:])
    b = np.empty(n - 1)
    c = np.empty(n - 1)
    d = np.empty(n - 1)
    c[n - 2] = 0

    for i in range(n - 3, -1, -1):
        c[i] = delta_i[i] * c[i + 1] + lambda_i[i]
    for i in range(n - 2, -1, -1):
        b[i] = l_i[i] + 2 / 3 * c[i] * h_i[i] + 1 / 3 * h_i[i] * c[i - 1]
        d[i] = (c[i] - c[i - 1]) / (3 * h_i[i])

    for i in range(n - 1):
        space_x = np.linspace(list_x[i], list_x[i + 1])
        plt.plot(space_x, a[i] + (b[i] + (c[i] + d[i] * (space_x - list_x[i + 1])) * (space_x - list_x[i + 1])) *
            (space_x - list_x[i + 1]))

    plt.scatter(list_x, list_y)
    plt.xlabel("Axis: X")
    plt.ylabel("Axis: Y")
    plt.title("Cubic Interpolation")
    plt.show()

    return a, b, c, d


def linear_spline_function(x_value, x_points, a, b):
    interval_index = len(x_points) - 2
    for i in range(1, len(x_points)):
        if x_value < x_points[i]:
            interval_index = i - 1
            break

    return a[interval_index] * x_value + b[interval_index]


def quadratic_spline_function(x_value, x_points, a, b, c):
    interval_index = len(x_points) - 3
    for i in range(1, len(x_points) - 1):
        if x_value < x_points[i]:
            interval_index = i - 1
            break

    return c[interval_index] + (b[interval_index] + a[interval_index] * x_value) * x_value


def cubic_spline_function(x_value, x_points, a, b, c, d):
    ii = len(x_points) - 2  # ii == interval_index
    for i in range(1, len(x_points)):
        if x_value < x_points[i]:
            ii = i - 1
            break

    return a[ii] + (b[ii] + (c[ii] + d[ii] * (x_value - x_points[ii + 1])) * (x_value - x_points[ii + 1])) * (
            x_value - x_points[ii + 1])

def main():

    x = [0.357, 0.871, 1.567, 2.032, 2.628, 3.228]
    y = [0.548, 1.012, 1.159, 0.694, -0.503, 0.703]

    table = PrettyTable()
    table.add_column("#", [i for i in range(len(x))])
    table.add_column("x", x)
    table.add_column("y", y)

#region Task_1

    print("\nTable Values:")
    print(table)

    func_l = get_interpolation_lagrange_polynomial(x, y)
    table_for_func_l = PrettyTable()
    table_for_func_l.add_column("#", [i for i in range(len(x))])
    table_for_func_l.add_column("x", x)
    table_for_func_l.add_column("y", [func_l(x_i) for x_i in x])

    print("\nLagrangian polynomial:")
    print(table_for_func_l)

    print("\nValue of L(x1 + x2):")
    print(func_l(x[0] + x[1]))

    func_l_range = np.linspace(0, x[-1])
    plt.plot(func_l_range, func_l(func_l_range))
    plt.scatter(x, y)
    plt.xlabel("Axis: X")
    plt.ylabel("Axis: Y")
    plt.title("Lagrangian polynomial")
    plt.grid()
    plt.show()

#endregion

#region Task_2

    din_diffs = table_finite_difference(x, y)

    print("\nFinite difference:")
    finite_table = PrettyTable()
    finite_table.add_column("#", [i for i in range(len(x))])
    finite_table.add_column("x", x)
    finite_table.add_column("y", y)

    curr_len = len(x)
    for i in range(len(din_diffs) - 1):
        add = ["None"]
        for j in range(1, curr_len):
            add.append(din_diffs[j][i])
        finite_table.add_column("del_" + str(i + 1), add)
    print(finite_table)

    print("\nDivide difference:")
    div_diffs = table_divided_difference(x, y)
    div_table = PrettyTable()
    div_table.add_column("#", [i for i in range(len(x))])
    div_table.add_column("x", x)
    div_table.add_column("y", y)
    for i in range(len(div_diffs) - 1):
        add = ["None"]
        for j in range(1, curr_len):
            add.append(div_diffs[j][i])
        div_table.add_column("del_" + str(i + 1), add)
    print(div_table)



    f_N = newton_polynomial(div_diffs)
    print("\nValue of N(x1 + x2):")
    print(f_N(x[0] + x[1]))

    range_x = np.linspace(0, x[-1])
    plt.plot(range_x, f_N(range_x))
    plt.scatter(x, y)
    plt.xlabel("Axis X")
    plt.ylabel("Axis Y")
    plt.title("Newton polynomial")
    plt.show()

#endregion

    linear_coefficients = linear_interpolation(x, y)
    quadratic_coefficients = quadratic_interpolation(x, y)
    cubic_coefficients = cubic_interpolation(x, y)

    space = np.linspace(0, x[-1])
    plt.xlabel("Axis: X")
    plt.ylabel("Axis: Y")
    plt.plot(space, func_l(space), label="Lagrange")
    plt.plot(space, f_N(space), label="Newton")
    plt.plot(space, [linear_spline_function(x_value, x, *linear_coefficients) for x_value in space], label="Linear")
    plt.plot(space, [quadratic_spline_function(x_value, x, *quadratic_coefficients) for x_value in space],label="Quadratic")
    plt.plot(space, [cubic_spline_function(x_value, x, *cubic_coefficients) for x_value in space], label="Cubic")
    plt.scatter(x, y)
    plt.legend()
    plt.show()

if __name__ == '__main__':
    main()