import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

L = 10
delta_u = 0.1
E = 110e9
rho = 4300
N_x = 100
h = L / N_x
tau = h / np.sqrt(E/rho)
T = 1
N_t = int(T/tau)
first = 1
second = 1


# u(x,0)
def p(x):
    return (10*x - 100) * x


def f(x):
    return 0


# u'(x,0)
def q(x):
    return 0


def solver():
    x_s = np.linspace(0, L, N_x)
    a = E/rho * tau**2 / h**2
    U = np.zeros((N_t, N_x), dtype=np.float)
    U[:][0] = 0
    U[:][-1] = 0
    U[0][1:-1] = p(x_s[1:-1])
    U[1][1:-1] = p(x_s[1:-1]) + tau * q(x_s[1:-1]) + tau**2 / 2 * (f(x_s[1:-1]) + p(x_s[2:]) - 2 * p(x_s[1:-1]) + p(x_s[2:])/h**2)
    for i in range(2, N_t):
        U[i][1:-1] = -U[i-2][1:-1] + a * (U[i-1][2:] + U[i-1][:-2]) + U[i-1][1:-1] * (2 - 2*a)
    return U, x_s


def solution_drawer(u, x):
    u_lines = [u[:][i] for i in range(u.shape[1])]
    for u_x in u_lines:
        plt.plot(x, u_x)
    plt.show()


U, X = solver()
# solution_drawer(U, X)

fig = plt.figure()
ax = plt.axes(xlim=(0, 10), ylim=(-300, 300))
line, = ax.plot([], [], lw=3)
ax.grid()



def init():
    line.set_data([], [])
    return line,


def animate(i):
    line.set_data(X, U[i])
    return line,


anim = FuncAnimation(fig, animate, init_func=init,
                     frames=len(U), interval=20, blit=False)

plt.show()