import scipy.io as sio
import numpy as np
import sys
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 22})

plt.interactive(True)

# Channel flow
data=np.loadtxt("chan_data.dat", comments="%")

ni=199  # number of grid nodes in x_1 direction
nj=15  # number of grid nodes in x_2 direction
v1=data[:,0] #don't use this array
v2=data[:,1] #don't use this array
p=data[:,2]  #don't use this array

# transform the arrays from 1D fields x(n) to 2D fields x(i,j)
# the first index 'i', correponds to the x-direction
# the second index 'j', correponds to the y-direction

v1_2d=np.reshape(v1,(nj,ni)) #this is v_1 (streamwise velocity component)
v2_2d=np.reshape(v2,(nj,ni)) #this is v_1 (wall-normal velocity component)
p_2d=np.reshape(p,(ni,nj))   #this is p   (pressure)

xp=np.loadtxt("xp.dat", comments="%")
yp=np.loadtxt("yp.dat", comments="%")

v1_2d=np.transpose(v1_2d)
v2_2d=np.transpose(v2_2d)
p_2d=np.transpose(p_2d)

# Q1/1  ######################################################################
# find peak velocity at each x1
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
v1_peak=np.max(v1_2d,axis=1)
plt.plot(xp,v1_peak,'b-')
plt.xlabel('$x$')
plt.ylabel('$v_{1,peak}$')
plt.savefig('v1_peak.eps')

# Q1/2  ######################################################################
# tau_wall = nu*dv_1dv_2 
dudy=np.gradient(v1_2d,yp,axis=1) 

# viscosity for air
nu=15.2e-6
# compute the wall shear stress at the upper wall. 
tau_w = nu*np.abs(dudy[:,-1]) #Index [-1] means last index. we take abs because we know that dudy < 0 near the upper wall
# skin friction is tau_w/(0.5*U_in^2), see Eq. F.2 
# U_in is constant at the inlet
u_in=v1_2d[0,1]
cf=tau_w/(0.5*u_in**2)

# plot skin frition
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
plt.plot(xp,cf,'b-')
plt.xlabel('$x$') 
plt.ylabel('$C_f$') 
plt.axis([0, max(xp),0,max(cf)])
plt.savefig('cf.eps')


# Q1/3  ######################################################################
# in the eBook, we define fully developed flow where dv1_dx1 < 0.01 at the center
# find dv1_dv1 at the center. Center is here the first cell
dv1_dx1_c=np.gradient(v1_2d[:,0],xp)

# find where dv1_dx1_c = 0.01
xx=0.01;
i1 = (np.abs(xx-dv1_dx1_c)).argmin()  # find index which closest fits xx
# it happens at xp[i1]=0.5


# Q1/4  ######################################################################
# compute omega_1
dv1_dx1,dv1_dx2=np.gradient(v1_2d,xp,yp)
dv2_dx1,dv2_dx2=np.gradient(v2_2d,xp,yp)
omega3=dv2_dx1-dv1_dx1

# find max index. 
# we expect it to be largest along the wall
print('[i,j]', np.where(omega3 == np.amax(omega3)) )

# We get i=5, j=13 (i.e. near the wall)

# Q1/5  ######################################################################
# the flow is irrotational at the inlet where v1=v_in and v2=0. We find that 
# dv1_dx2 is indeed zero but dv2_dx1 has small (non-zero) values.