import numpy as np, matplotlib.pyplot as plt, matplotlib.animation as ani, time
from numpy import sin, cos, pi, e, exp, log

def m_print(m):
    for row in m:
        for x in row:
            x = round(x,2)
            print(x, end=' '*(5-len(str(x))))
        print()
    print()

def v_print(v):
    for x in v:
        x = round(x,2)
        print(x, end=' '*(5-len(str(x))))
    print()

t0 = time.process_time()

x0,xn = 0,1
tn = 0.4

n = 100
h = (xn-x0)/n
k = 0.5*(h**2)

x_points = np.linspace(x0,xn,n)

A = np.array([np.zeros(n-2) for i in range(n-2)])

A[0][0] = 2
A[0][1] = -1
A[-1][-1] = 2
A[-1][-2] = -1
for i in range(1,len(A)-1):
    A[i][i-1] = -1
    A[i][i] = 2
    A[i][i+1] = -1

M = np.identity(n-2) - (k/h**2)*A

U = np.array([np.zeros(n) for i in range(int(tn/k))])
u0 = sin(pi*x_points)
U[0] = u0

for i in range(len(U)-1):
    U[i+1][1:-1] = np.dot(U[i][1:-1], M)

###############################
#alt under her er bare plotting
###############################

print('eksp_tid: ', time.process_time()- t0)

nyU = np.array([np.zeros(n) for i in range(int(len(U)/16))])
for i in range(0,len(U),16):
    try:
        nyU[int(i/16)] = U[i]
    except IndexError:
        pass

k=0
fig, ax = plt.subplots()

print(len(U), len(nyU))

def animate(i):
    ferdig = False
    global k

    try:
        u = nyU[k]
    except IndexError:
        ferdig = True
        u = [x for x in x_points]
        ux = [1-x for x in x_points]

    k+=1

    ax.clear()
    plt.plot(x_points,u)

    if ferdig:
        plt.plot(x_points, ux, color='b')

    plt.ylim(0,1)
    plt.xlim(0,1)


ani = ani.FuncAnimation(fig, animate, frames=360, interval=1)
plt.show()