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

plt.interactive(True)

viscos=1/5200

y_u_k_eps_vis=np.loadtxt("y_u_k_eps_vis.dat")
x_2=y_u_k_eps_vis[:,0]
v_1=y_u_k_eps_vis[:,1]
k=y_u_k_eps_vis[:,2]
eps=y_u_k_eps_vis[:,3]
vis=y_u_k_eps_vis[:,4] # this is total viscosity, i.e. vis=viscos+vis_turb

############################   Q2/1

# wall shear stress south wall
dudy=np.gradient(v_1,x_2)
tauw=viscos*dudy[0]
u_tau=tauw**0.5
yplus=x_2*u_tau/viscos
uplus=v_1/u_tau

#************
# velocity profile plot
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
# plot the velocity profile 
plt.semilogx(yplus,v_1,'b-')
plt.axis([0.1, 5200,0, 27])
plt.ylabel("$v_1^+$")
plt.xlabel("$x_2^+$")
plt.savefig('u.eps')

# plot k
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
plt.plot(k,yplus,'b-')
plt.xlabel("$k$")
plt.ylabel("$x_2^+$")

# plot eps
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
plt.plot(eps,yplus,'b-')
plt.xlabel(r"$\varepsilon$")
plt.ylabel("$x_2^+$")
plt.savefig('eps.eps')

############################   Q2/2
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
# plot eps
plt.plot(eps,yplus,'b-')
plt.xlabel(r"$\varepsilon, \quad P^k$")
plt.ylabel("$x_2^+$")

# plot Pk
dudy=np.gradient(v_1,x_2)
vist=vis-viscos
prod=vist*dudy**2
plt.plot(prod,yplus,'r--')
plt.axis([0.1, 2000,0, 200])
plt.savefig('pk-diss.eps')

############################   Q2/3
# check b.c. for eps
eps_bc=2*viscos*k[1]/x_2[1]**2
print('eps[1],eps_bc]',eps[1],eps_bc)

############################   Q3/1
# plot fmu
cmu=0.09
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
# compute fmu backwards, see line beforw Eq. 11.159
fmu=vist/(cmu*k**2/eps)
plt.plot(fmu,yplus,'b-')
plt.xlabel(r"$f_\mu$")
plt.ylabel("$x_2^+$")
plt.axis([0,1,0, 100])
plt.savefig('fmu.eps')

############################   Q3/2
# plot large terms in v1 eq
# in Fig. 6.6 we see that the two viscous and turb. diffusion terms are large
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
uv=-vist*dudy
duvdy=np.gradient(uv,x_2)
d2udy2=np.gradient(dudy,x_2)
# this is the viscous diffusion term
visc_term=viscos*d2udy2
# this is the turbulent diffusion term
plt.plot(-duvdy,yplus,'b-')
plt.plot(visc_term,yplus,'r-')
plt.xlabel("large terms, v1 eq")
plt.ylabel("$x_2^+$")
plt.axis([-500,500,0, 100])
plt.savefig('terms-in-u-eq.eps')

############################   Q3/3
# plot large terms near the wall in the k eq
# Iin Fig. 9.1. we see that the large term at the wall is the dissipation and the viscous diffusion
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.25,bottom=0.20)
dkdy=np.gradient(k,x_2)
d2kdy2=np.gradient(dkdy,x_2)
# this is the viscous diffusion term
visc_term_k=viscos*d2kdy2
plt.plot(-eps,yplus,'b-')
plt.plot(visc_term_k,yplus,'r-')
plt.xlabel("large terms, k eq")
plt.ylabel("$x_2^+$")
plt.axis([-1000,1000,0,5])
plt.savefig('terms-in-k-eq.eps')