#!/usr/bin/python3

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

U, nu = 1., 1.  # asymptotic velocity, viscosity
max_eta = 100.  # 'infinity' = some sufficiently large eta value
g0, f0 = 0., 0. # initial conditions for g=f' and f at eta=0
N = 2000        # number of RK4 steps

def rhs(r, eta):  # the right-hand side of the Blasius equation
    f, g, h = r   # unpack the array of function values r
    return np.array([g, h, - h*f/2]) # return array of f', g', h'
    
def g_inf(h0):  # compute g(eta=max_eta) as a function of h(0)
    r0 = np.array([f0, g0, h0])
    etapoints, rpoints = rk4.rk4(rhs, 0, max_eta, r0, N)
    return rpoints[-1][1]
    
def bisection(f, a, b, delta=1.E-3): # find a zero of f in [a,b] by bisection
    if f(a) * f(b) > 0:
        print("No zero guaranteed to exist in this interval")
        return
    if a > b:
        a, b = b, a
    c = (b + a) / 2
    while b - a > delta:
        if f(a) * f(c) < 0:
            b = c
        else:
            a = c
        c = (b + a) /2
    return c

h0 = bisection(lambda h0: g_inf(h0)-1, 0, 1)  # find a zero h0 of g_inf-1
r0 = np.array([f0, g0, h0])  # use this as the missing initial condition
etapoints, rpoints = rk4.rk4(rhs, 0, max_eta, r0, N) # recompute solution
fpoints = rpoints[:, 0]
dfpoints = rpoints[:, 1]
ddfpoints = rpoints[:, 2]

def eta(x, y):
    return np.sqrt(U / nu / x) * y

def ux(x, y): # compute u_x by interpolating the numerical solution 
    return U * np.interp(eta(x, y), etapoints, dfpoints)
    
def uy(x, y): # same for u_y
    etafprime = eta(x, y) * np.interp(eta(x, y), etapoints, dfpoints)
    f = np.interp(eta(x, y), etapoints, fpoints)
    return 0.5 * np.sqrt(U * nu / x) * (etafprime - f)
    
# Plot velocity field above plate    
xpoints = np.linspace(.5, 10, 10)
ypoints = np.linspace(.5, 10, 10)
X, Y = np.meshgrid(xpoints, ypoints)
uxpoints = ux(X, Y)
uypoints = uy(X, Y)
plt.quiver(xpoints, ypoints, uxpoints, uypoints,angles='xy', scale_units='xy', scale=1)
plt.xlim(right=11)

# Plot f, f', and f''
plt.figure()
plt.ylim(top=1)
plt.xlim(right=5)
plt.plot(etapoints, fpoints, label=r"$f(\eta)$")
plt.plot(etapoints, dfpoints,label=r"$f'(\eta)$")
plt.plot(etapoints, ddfpoints,label=r"$f''(\eta)$")
plt.legend()

plt.show()