# Numerical Methods II, Courant Institute, NYU, spring 2018
# http://www.math.nyu.edu/faculty/goodman/teaching/NumericalMethodsII2018/index.html
#  written by Jonathan Goodman (instructor)
#  see class notes Part 3 for more discussion


import numpy as np

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.animation as animation

from TimeStep import step


#    Physical parameters

R  = 3.    # box size, linear dimension
D  = 1.    # diffusion coefficient
T  = 1.    # final time for the computation
Tf = .1    # time between movie frames (may be changed slightly)

#    numerical parameters

n   = 50            # number of grid intervals in each dimension
lam = .25           # the CFL ratio = D*dt/(dx^2)
dx  = R/n
dt  = dx*dx*lam/D   # time step may be adjusted down slightly for the movie

nf  = int(T/Tf)+1   # The number of movie frames = total time/frame time + adjustment
Tf  = T/nf          # A slightly smaller time per frame to get exactly to time T at the end
ntf = int(Tf/dt)+1  # The number of time steps per frame = (time per frame)/dt
ntf = ((ntf+1)/2 )*2
dt  = Tf/ntf

uFrames = np.ndarray([ nf, n-1, n-1]) # Store all solution at the plot times
                                      # A real computation would not do this, but
                                      # I am sick of trying to get Python to
                                      # make movies the "right" way.

u = np.ndarray([n+1,n+1])
v = np.ndarray([n+1,n+1])
for i in range(n+1):
    u[0,i] = 0.
    u[n,i] = 0.
    u[i,0] = 0.
    u[i,n] = 0.

for i in range(1,n):
    for j in range(1,n):
        u[i,j] = 1.

for i in range(n+1):
    for j in range(n+1):
        v[i,j] = u[i,j]





for frame in range(nf):
    print "frame " + str(frame)
    for i in range(ntf/2):        # there are two steps per trip through the loop

        step( v, u, D, dx, dt)     # copy u to v
        step( u, v, D, dx, dt)     # copy v back to u

    uFrames[frame, :, :] = u[1:n,1:n]



fig = plt.figure()

im = plt.imshow(uFrames[0,:,:], animated=True)


def updatefig(*args):
    global x, y
    x += np.pi / 15.
    y += np.pi / 20.
    im.set_array(f(x, y))
    return im,

ani = animation.FuncAnimation(fig, updatefig, interval=50, blit=True)
plt.show()