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

# konstanter
x0 = 0  # første punkt i x-retning
xf = 1  # siste punkt i x-retning
y0 = 0  # første punkt i y-retning
yf = 1  # siste punkt i y-retning
t0 = 0  # første punkt i t-retning
tf = 3   # siste punkt i t-retning
Nx = 51  # antall x-punkter 
Ny = 51  # antall y-punkter
Nt = 1001  # antall t-punkter
l = (xf-x0)/(Nx-1)  # x-gitterbredde
m = (yf-y0)/(Ny-1)  # y-gitterbredde
n = (tf-t0)/(Nt-1)  # t-gitterbredde 
A = (n**2)/(l**2)
B = (n**2)/(m**2)
x = np.linspace(x0,xf,Nx)
y = np.linspace(y0,yf,Ny)
t = np.linspace(t0,tf,Nt)



# initialkrav
def f(X,Y):
    return np.exp(-15*((X-0.5)**2+(Y-0.5)**2))

u = np.zeros((Nx,Ny,Nt))
for i in range(1,Nx-1):
    for j in range(1,Ny-1):
        u[i][j][0] = f(x[i],y[j])
# Neumann-randkrav
for i in range(Nx):
    u[i][0][0] = u[i][1][0]
    u[i][-1][0] = u[i][-2][0]
for j in range(Ny):
    u[0][j][0] = u[1][j][0]
    u[-1][j][0] = u[-2][j][0]



# neste tidssteg
def u_next(uc):
    un = np.zeros((Nx,Ny))
    for i in range(1,Nx-1):
        for j in range(1,Ny-1):
            un[i][j] = (1-A-B)*uc[i][j] + A*uc[i-1][j]/2 + A*uc[i+1][j]/2 + B*uc[i][j-1]/2 + B*uc[i][j+1]/2 
    # Mer Neumann
    for i in range(Nx):
        un[i][0] = un[i][1]
        un[i][-1] = un[i][-2]
    for j in range(Ny):
        un[0][j] = un[1][j]
        un[-1][j] = un[-2][j]
    return un

def u_next2(uc,up):
    un = np.zeros((Nx,Ny))
    for i in range(1,Nx-1):
        for j in range(1,Ny-1):
            un[i][j] = 2*(1-A-B)*uc[i][j] + A*uc[i-1][j] + A*uc[i+1][j] + B*uc[i][j-1] + B*uc[i][j+1] - up[i][j]
    # Mer Neumann
    for i in range(Nx):
        un[i][0] = un[i][1]
        un[i][-1] = un[i][-2]
    for j in range(Ny):
        un[0][j] = un[1][j]
        un[-1][j] = un[-2][j]
    return un

u[:,:,1] = u_next(u[:,:,0])
for i in range(1,Nt-1):
    u[:,:,i+1] = u_next2(u[:,:,i],u[:,:,i-1])




# animering
fig, ax = plt.subplots(subplot_kw={"projection": "3d"},dpi=200)
T_ani = 8   # tid for animasjonen i sekunder
fps = 40    # fps
total_frames = T_ani*fps
meshX, meshY = np.meshgrid(x, y)

def update(i):
    ax.clear()
    surf = ax.plot_surface(meshX, meshY, u[:,:,i])
    ax.set_xlabel("$x$")
    ax.set_ylabel("$y$")
    ax.set_zlabel("$u$")
    ax.set_xlim((x0,xf))
    ax.set_ylim((y0,yf))
    ax.set_zlim((-1.1,1.1))
    return surf




ani = animation.FuncAnimation(fig, update, repeat=True, interval=1000/fps, 
                              frames=np.linspace(0,Nt-1,total_frames).astype(int))
ani.save("bolgelikn_2d.gif")
plt.show()