#!/usr/bin/python
# Boundary layer flow
import numpy as np
import matplotlib.pyplot as plt


def K(t):
    """
    Forcing function as a function of time.
    """
    return (1-np.exp(-0.1*t))*np.cos(t)


s = 10  # Shear wave number
n = 50   # Number of gridpoints in vertical direction

# Grid
y = np.linspace(0, 1, n)

dy = y[1]-y[0]
dt = 0.0005

l = (dt/(s**2*dy**2))

hnu = np.exp(-np.sqrt(1j)*s*y)
# fnu=(1-1j)/s
fnu = 0


def u_ex(tn):
    """
    Exact harmonic solution of the velocity field
    """
    return (((1-hnu)/(1-fnu))*np.exp(1j*(tn))/1j).real


def u_np1(un, tn, dt):
    Kn = K(tn)
    unp1 = un
    unp1[0] = 0  # Velocity zero ver here
    unp1[1:-1] = un[1:-1]+dt*Kn+l*(un[0:-2]-2*un[1:-1]+un[2:])
    unp1[-1] = unp1[-2]  # Approximate 'infinity' bc
    return unp1


un0 = np.zeros(n, float)
t = 0
un = un0

f = plt.figure(1, figsize=(12, 6))

linefd, = plt.plot(un0, y)
linee, = plt.plot(un0, y)
plt.legend(('Finite difference', 'Periodic exact'),
           loc='lower right')
plt.ylim(0, 1)
plt.xlim(-1.5, 1.5)
plt.ylabel('Distance from wall [m]')
plt.xlabel('Velocity [m/s]')
plt.grid(True)
plt.show(block=False)
i = 0
uold = un

try:
    while(True):
        t += dt
        uold = un
        un = u_np1(uold, t, dt)
        if(i % 25 == 0):
            linefd.set_xdata(un)
            linee.set_xdata(u_ex(t))
            f.canvas.draw_idle()
            plt.pause(.000005)
        if not plt.fignum_exists(1):
            break
        i += 1
except Exception:
    pass

exit(0)