import scipy.io as sio
import sys
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from IPython import display
plt.rcParams.update({'font.size': 22})
plt.rcParams.update({'figure.max_open_warning': 0})

plt.interactive(True)

viscos=1/395


# makes sure figures are updated when using ipython
display.clear_output(wait=True)

datax= np.loadtxt("x2d.dat")
x=datax[0:-1]
ni=int(datax[-1])
datay= np.loadtxt("y2d.dat")
y=datay[0:-1]
nj=int(datay[-1])

x2d=np.zeros((ni+1,nj+1))
y2d=np.zeros((ni+1,nj+1))

x2d=np.reshape(x,(ni+1,nj+1))
y2d=np.reshape(y,(ni+1,nj+1))

# compute cell centers
xp2d=0.25*(x2d[0:-1,0:-1]+x2d[0:-1,1:]+x2d[1:,0:-1]+x2d[1:,1:])
yp2d=0.25*(y2d[0:-1,0:-1]+y2d[0:-1,1:]+y2d[1:,0:-1]+y2d[1:,1:])

nk=32
zmax=1.6
# z grid
zmax, nk=np.loadtxt('z.dat')
zp = np.linspace(0, zmax, nk)


viscos=1/395

itstep,nk,dz=np.load('itstep.npy')
p2d=np.load('p_averaged.npy')/itstep
u2d=np.load('u_averaged.npy')/itstep
v2d=np.load('v_averaged.npy')/itstep
w2d=np.load('w_averaged.npy')/itstep
k2d=np.load('k_averaged.npy')/itstep
vis2d=np.load('vis_averaged.npy')/itstep
uu2d=np.load('uu_stress.npy')/itstep
vv2d=np.load('vv_stress.npy')/itstep
ww2d=np.load('ww_stress.npy')/itstep
uv2d=np.load('uv_stress.npy')/itstep
fk2d=np.load('fk_averaged.npy')/itstep
om2d=np.load('om_averaged.npy')/itstep

uu2d=uu2d-u2d**2
uv2d=uv2d-u2d*v2d

kres_2d=0.5*(uu2d+vv2d+ww2d)

vis2d=vis2d/viscos

ustar=(viscos*u2d[:,1]/yp2d[1,1])**0.5

two_corr=np.load('two_corr_inlet.npy')


# compute cf
cf=np.zeros(ni)
yplus2d=np.zeros((ni,nj))
for i in range (0,ni-1):
     xt=x2d[i,0]-x2d[i-1,0]
     yt=y2d[i,0]-y2d[i-1,0]
     rl=(xt**2+yt**2)**0.5
     xt=xt/rl
     yt=yt/rl
     uwall=abs(u2d[i,1]*xt+v2d[i,1]*yt)
     xn=-yt
     yn=xt
     yp=abs((yp2d[i,1]-yp2d[i,0])*yn+(xp2d[i,1]-xp2d[i,0])*xn)
     ustars=(uwall*viscos/yp)**0.5*u2d[i,1]/abs(u2d[i,1])
     cf[i]=abs(ustars)*ustars/0.5
     yplus2d[i,:]=ustar[i]*yp2d[i,:]/viscos


cf[0]=cf[1]

y_u_k_om=np.loadtxt('y_u_k_om_uv_395.dat')
y_rans=y_u_k_om[:,0]
u_rans=y_u_k_om[:,1]
# make it 2D
u_rans=np.repeat(u_rans[:,None], repeats=nk, axis=1)
k_rans=y_u_k_om[:,2]
om_rans=y_u_k_om[:,3]
uv_rans=np.abs(y_u_k_om[:,4])


data=np.loadtxt('../channel-395-DNS/y_u_k_om_uv_395_DNS-python.txt')
y_DNS=data[:,0]
u_DNS=data[:,1]
uu_DNS=data[:,2]
vv_DNS=data[:,3]
ww_DNS=data[:,4]
uv_DNS=data[:,5]

# find max
nmax=np.argmax(two_corr)
# and its value
two_max=np.max(two_corr)
# two_max is symmetric. Pick the right half and normalize
two_sym_norm=two_corr[nmax:]/two_max


########################################## Ustar
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
plt.plot(xp2d[:,1],ustar,'b-')
plt.xlabel('$x$')
plt.ylabel(r"$u_\tau$")
plt.savefig('ustar-vs-x.eps')



########################
fig2 = plt.figure("Figure 2")
imax=10
plt.subplots_adjust(left=0.20,bottom=0.20)
plt.plot(zp[0:imax],two_sym_norm[0:imax],'ro')
plt.plot(zp[0:imax],two_sym_norm[0:imax],'r-')
plt.xlabel('zeta')
plt.title("2-p corr")
plt.savefig('two-point-corr-inlet.eps',bbox_inches='tight')


########################################## U 
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
i1 = 0
plt.semilogx(yplus2d[i1,:],u2d[i1,:],'b-')
plt.semilogx(yplus2d[i1,:],u_DNS,'o')
xx=1
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.semilogx(yplus2d[i1,:],u2d[i1,:],'r-')
xx=2
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.semilogx(yplus2d[i1,:],u2d[i1,:],'r--')
plt.xlabel("$U$")
plt.ylabel("$y^+$")
plt.axis([1, 400, 0, 28])
plt.savefig('u_log_python.eps',bbox_inches='tight')


########################################## uu 
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
i1 = 0
plt.plot(yplus2d[i1,:],uu2d[i1,:],'b-')
plt.plot(yplus2d[i1,:],uu_DNS,'o')
xx=1
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],uu2d[i1,:],'r--')
xx=2
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],uu2d[i1,:],'k-.')
plt.xlabel(r"$\overline{u'u'}$")
plt.ylabel("$y^+$")
plt.axis([1, 400, 0, 2])
plt.savefig('uu_python.eps',bbox_inches='tight')



########################################## vv 
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
i1 = 0
plt.plot(yplus2d[i1,:],vv2d[i1,:],'b-')
plt.plot(yplus2d[i1,:],vv_DNS,'o')
xx=1
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],vv2d[i1,:],'r--')
xx=2
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],vv2d[i1,:],'k-.')
plt.xlabel(r"$\overline{v'v'}$")
plt.ylabel("$y^+$")
plt.axis([1, 400, 0, 2])
plt.savefig('vv_python.eps',bbox_inches='tight')



########################################## ww 
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
i1 = 0
plt.plot(yplus2d[i1,:],ww2d[i1,:],'b-')
plt.plot(yplus2d[i1,:],ww_DNS,'o')
xx=1
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],ww2d[i1,:],'r--')
xx=2
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],ww2d[i1,:],'k-.')
plt.xlabel(r"$\overline{w'w'}$")
plt.ylabel("$y^+$")
plt.axis([1, 400, 0, 2])
plt.savefig('ww_python.eps',bbox_inches='tight')



########################################## uv 
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
i1 = 0
plt.plot(yplus2d[i1,:],uv2d[i1,:],'b-')
plt.plot(yplus2d[i1,:],uv_DNS,'o')
xx=1
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],uv2d[i1,:],'r--')
xx=2
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],uv2d[i1,:],'k-.')
plt.plot(y_rans/viscos,uv_rans,'k-.')
plt.xlabel(r"$\overline{u'v'}$")
plt.ylabel("$y^+$")
plt.axis([1, 400, -1, 0])
plt.savefig('uv_python.eps',bbox_inches='tight')



########################################## uv 
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
i1 = 0
plt.plot(yplus2d[i1,:],vis2d[i1,:],'b-')
xx=1
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],vis2d[i1,:],'r--')
xx=2
i1 = (np.abs(xx-xp2d[:,1])).argmin()  # find index which closest fits xx
plt.plot(yplus2d[i1,:],vis2d[i1,:],'k-.')
plt.xlabel(r"$\nu_t/\nu$")
plt.ylabel("$y^+$")
plt.axis([1, 400, 0,4])
plt.savefig('vis_python.eps',bbox_inches='tight')

#### auto-correlation
w_time=np.load('w_time.npy')
usynt_inlet=w_time[0,:]
auto_uu=np.correlate(usynt_inlet,usynt_inlet,'full')
nmax_auto=np.argmax(auto_uu)
# and its value
auto_max=np.max(auto_uu)
# two_max is symmetric. Pick the right half and normalize
auto_sym_norm=auto_uu[nmax_auto:]/auto_max

########################
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
plt.plot(auto_sym_norm,'r-')
plt.plot(auto_sym_norm,'ro')
plt.xlabel('time')
plt.ylabel('auto-corr')
plt.title('synt-fluct')
plt.axis([0,40,-0.2,1.1])
plt.savefig('auto_uu_synt_fluct_python.eps')


