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)

nu=15e-6

# Channel flow
data=np.genfromtxt("channel_flow_data.dat", comments="%")

ni=199  # number of grid nodes in x_1 direction
nj=28  # number of grid nodes in x_2 direction
x1=data[:,0] #don't use this array
x2=data[:,1] #don't use this array
v1=data[:,2] #don't use this array
v2=data[:,3] #don't use this array
p=data[:,4]  #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

x1_2d=np.reshape(x1,(nj,ni)) #this is x_1 (streamwise coordinate)
x2_2d=np.reshape(x2,(nj,ni)) #this is x_2 (wall-normal coordinate)
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,(nj,ni))   #this is p   (pressure)

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


## Q1/1.

# the vorticity vector is defined as (see Eq. 1.12)
# omega_i = eps_ijk dv_k/dx_j
# omega_1 =  eps_1jk dv_k/dx_j = eps_123 dv_3/dx_2 + eps_132 dv_2/dx_3= 0
# omega_2 =  eps_2jk dv_k/dx_j = eps_231 dv_3/dx_1 + eps_213 dv_1/dx_3= 0
# omega_3 =  eps_3jk dv_k/dx_j = eps_312 dv_2/dx_1 + eps_321 dv_1/dx_2= dv_2/dx_1 - dv_1/dx_2=

dudy=np.gradient(v1_2d,x2_2d[1,:],axis=1)
dvdx=np.gradient(v2_2d,x1_2d[:,1],axis=0) 

omega_3 = dvdx-dudy


# the wall vorticty flux (see text below Eq 4.24b)
omega_3_flux=-nu*np.gradient(omega_3,x2_2d[1,:],axis=1)


#************
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
xx=0.01
i1 = (np.abs(xx-x1_2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(omega_3[i1,:],x2_2d[i1,:],'b-')
print('wall vorticiy flux at x=0.01: bottom wall: ',omega_3_flux[i1,0],' top wall: ', omega_3_flux[i1,-1])
xx=0.5
i1 = (np.abs(xx-x1_2d[:,1])).argmin()  # find index which closest fits xx
print('\nwall vorticiy flux at x=0.01: bottom wall: ',omega_3_flux[i1,0],' top wall: ', omega_3_flux[i1,-1])
plt.plot(omega_3[i1,:],x2_2d[i1,:],'r--')
plt.ylabel('$x_2$') 
plt.xlabel(r'$\omega_3$') 
plt.savefig('omega_3.eps')


## Q1/2.
# dissipation phi = 2 mu s_ij s_ij (Eq. 2.15)
# diss =  2 mu (s11**2 + s12**2 +  s21**2 +  s22**2)

dvdy=np.gradient(v2_2d,x2_2d[1,:],axis=1)
dudx=np.gradient(v1_2d,x1_2d[:,1],axis=0) 

s11=dudx
s12=0.5*(dudy+dvdx)
s21=s12
s22=dvdx


diss = 2*nu*(s11**2 + s12**2 +  s21**2 +  s22**2)

#************
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
xx=0.01
i1 = (np.abs(xx-x1_2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(diss[i1,:],x2_2d[i1,:],'b-')
xx=0.5
i1 = (np.abs(xx-x1_2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(diss[i1,:],x2_2d[i1,:],'r--')
plt.ylabel('$x_2$') 
plt.xlabel('dissipation') 
plt.savefig('diss.eps')


## Q1/3.
# pressure drop = (pressure at inlet) -  (pressure at outlet)

p_inlet=np.mean(p_2d[0,:])
p_inlet=np.trapz(p_2d[0,:],x2_2d[0,:])/x2_2d[0,-1]
p_outlet=np.trapz(p_2d[-1,:],x2_2d[0,:])/x2_2d[0,-1]

print('pressure drop: ',p_inlet-p_outlet)