import torch.nn as nn
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 matplotlib import ticker

plt.rcParams.update({'font.size': 22})
plt.rcParams.update({'figure.max_open_warning': 0})

plt.interactive(True)

plt.close('all')
viscos=1/550



#      y/h             y+              U+             u'+             v'+             w'+           -Om_z+          om_x'+           om_y'+           om_z'+         uv'+             uw'+           vw'+             pr'+            ps'+          psto'+            p'
# -----------------------------------------------------
DNS_mean=np.genfromtxt("Re550.dat",comments="%")
y_DNS=DNS_mean[:,0];
yplus_DNS=DNS_mean[:,1];
u_DNS=DNS_mean[:,2];
u2_DNS=DNS_mean[:,3]**2;
v2_DNS=DNS_mean[:,4]**2;
w2_DNS=DNS_mean[:,5]**2;
uv_DNS=DNS_mean[:,10];

k_DNS=0.5*(u2_DNS+v2_DNS+w2_DNS)

# %
#%       y/h               y+           dissip          produc         p-strain          p-diff          t-diff         v-diff            bal          tp-kbal
DNS_RSTE = np.genfromtxt("Re550_bal_kbal.dat",comments="%")

diss_DNS = abs(DNS_RSTE[:,2])
visc_diff =  DNS_RSTE[:,7]
diff_DNS =  DNS_RSTE[:,6]
Pk_DNS = DNS_RSTE[:,3]


diss_DNS_iso = np.maximum(diss_DNS-visc_diff,0)

########################################## k  balance zoom from pyCALC
fig1,ax1 = plt.subplots()
plt.subplots_adjust(left=0.20,bottom=0.20)
plt.plot(yplus_DNS,Pk_DNS,'b-',label= '$P^k$')
plt.plot(yplus_DNS,-diss_DNS_iso,'r-',label= r'$D^{k,\nu}-\varepsilon$')
plt.plot(yplus_DNS,diff_DNS,'k-',label= r'$D^{k,t}$')
plt.legend(loc='best',fontsize=20)
plt.ylabel('terms in $k$ eq.')
plt.xlabel("$y^+$")
plt.axis([0, 100, -0.30, 0.30])
plt.savefig('k-balance_diss-iso-550-channel-kom.png',bbox_inches='tight')