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

# load grid
x= np.loadtxt("x.dat")
y= np.loadtxt("y.dat")
ni=len(x)
nj=len(y)
viscos=1./950



# read data file 1
vectz= np.loadtxt("vect_uuvvwwz.dat")
ntstep=int(vectz[0])
n=len(vectz)

nn=10
iu=range(1,n,nn)
iv=range(2,n,nn)
ieps=range(3,n,nn)
iuu=range(4,n,nn)
ivv=range(5,n,nn)
iww=range(6,n,nn)
iuv=range(7,n,nn)
ik=range(8,n,nn)
ivis=range(9,n,nn)
ipksgs=range(10,n,nn)

u=vectz[iu]/ntstep
v=vectz[iv]/ntstep
eps=vectz[ieps]/ntstep
uu=vectz[iuu]/ntstep
vv=vectz[ivv]/ntstep
ww=vectz[iww]/ntstep
uv=vectz[iuv]/ntstep
rk=vectz[ik]/ntstep
vis=vectz[ivis]/ntstep
pksgs=vectz[ipksgs]/ntstep

# uu is total inst. velocity squared. Hence the resolved turbulent resolved stresses are obtained as
uu=uu-u**2
vv=vv-v**2
uv=uv-u*v

u_2d=np.reshape(u,(ni,nj))
v_2d=np.reshape(v,(ni,nj))
eps_2d=np.reshape(eps,(ni,nj))
uu_2d=np.maximum(np.reshape(uu,(ni,nj)),0.)
vv_2d=np.maximum(np.reshape(vv,(ni,nj)),0.)
uu_2d=np.maximum(np.reshape(uu,(ni,nj)),0.)
vv_2d=np.maximum(np.reshape(vv,(ni,nj)),0.)
ww_2d=np.maximum(np.reshape(ww,(ni,nj)),0.)
uv_2d=np.reshape(uv,(ni,nj))
rk2d=np.reshape(rk,(ni,nj))
vis_2d=np.reshape(vis,(ni,nj))/viscos #this is to total viscosity, i.e. vis_tot=vis+vis_turb
pksgs_2d=np.reshape(pksgs,(ni,nj)) 

# compute terms in u-ekv
dudx_2d,dudy_2d = np.gradient(u_2d,x,y)


# The arrays have have the dimension (ni,nj). The first index
# represents moving in the $x_1$ direction. For example
#
#  u_2d(:,0):  nodes at the lower wall (low y)
#  u_2d(:,nj-1): nodes at the upper wall (high y)
#  u_2d(0,:):  nodes at the inlet (low x)
#  u_2d(ni-1,:): nodes at the outlet (high x)





##########################################  uu v. y
fig1 = plt.figure("Figure 1")
plt.clf() #clear the figure

#  plot uu vs. y
i=24 # x=1.175
plt.plot(y,uu_2d[i,:]**0.5,'b-',label='x=1.175')
i=59 #  x=2.925
plt.plot(y,uu_2d[i,:]**0.5,'r--',label='x=2.925')

plt.legend()
plt.axis([0,2,0,4.5])
plt.xlabel('$y$')
plt.ylabel('$\overline{u^\prime u^\prime}$')

plt.savefig('uu_vs_y_python.eps')

##########################################  uu v. x
fig2 = plt.figure("Figure 2")
plt.clf() #clear the figure

#  plot uu vs. x
j=29 #y=0.24
plt.plot(x,uu_2d[:,j]**0.5,'b-',label='y=0.24')
j=44 #y=1.35
plt.plot(x,uu_2d[:,j]**0.5,'r--',label='y=1.35')
plt.legend()
plt.ylabel('$y$')
plt.xlabel('$\overline{u^\prime u^\prime}$')

plt.axis([0,2,0,4.5])

plt.savefig('uu_vs_x_python.eps')