#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
from numpy import linalg as LA
import math as m
import os
import sys
from matplotlib.image import imread
import matplotlib.pyplot as plt
from matplotlib import rcParams # for changing default values
import scipy.io as sio
from scipy.optimize import minimize
import timeit
import torch
import random
import torch.nn as nn
import torch.optim as optim
from scipy.integrate import odeint
from torch.autograd import grad
import torch.optim.lr_scheduler as lr_scheduler

plt.close('all')
plt.interactive(True)
# set all fontsizes to 12
rcParams["font.size"] = 16
rcParams["axes.axisbelow"] = True # make sure grid is behind plots

viscos = 1/10000

# solve differential equation for k
# load DNS data
# %     y/h             y+             U+           u'+           v'+          w'+           uv'+         dU/dy+
DNS_mean=np.genfromtxt("P10k.dat",comments="%")
y_DNS=DNS_mean[:,0];
yplus_DNS=DNS_mean[:,1];
u_DNS=DNS_mean[:,2];
uu_DNS=DNS_mean[:,3]**2;
vv_DNS=DNS_mean[:,4]**2;
ww_DNS=DNS_mean[:,5]**2;
uv_DNS=DNS_mean[:,6];
dudy_DNS  = np.gradient(u_DNS,y_DNS)

k_DNS=0.5*(uu_DNS+vv_DNS+ww_DNS)

dkdy_DNS = np.gradient(k_DNS,y_DNS)
d2kdy2_DNS = np.gradient(dkdy_DNS,y_DNS)

# %      y/h            y+         dissip        prod         p-strain       p-diff        T-diff        V-diff
DNS_uu = np.genfromtxt("P10k.uu.dat",comments="%")
eps_DNS_uu = abs(DNS_uu[:,2])
visc_diff_uu =  DNS_uu[:,7]
diff_uu =  DNS_uu[:,6]
prod_uu =  DNS_uu[:,3]


DNS_vv = np.genfromtxt("P10k.vv.dat",comments="%")
eps_DNS_vv = abs(DNS_vv[:,2])
visc_diff_vv =  DNS_vv[:,7]
diff_vv =  DNS_vv[:,6]
prod_vv =  DNS_vv[:,3]

DNS_ww = np.genfromtxt("P10k.ww.dat",comments="%")
eps_DNS_ww = abs(DNS_ww[:,2])
visc_diff_ww =  DNS_ww[:,7]
diff_ww =  DNS_ww[:,6]
prod_ww =  DNS_ww[:,3]

diss_DNS = (eps_DNS_uu +eps_DNS_vv +eps_DNS_ww)/2
diff_DNS_visc = (visc_diff_uu +visc_diff_vv +visc_diff_ww)/2
diff_DNS = (diff_uu +diff_vv +diff_ww)/2
Pk_DNS = (prod_uu +prod_vv +prod_ww)/2

diss_DNS_iso = np.maximum(diss_DNS-diff_DNS_visc,1e-10)

vist_DNS = np.abs(uv_DNS/dudy_DNS)

########################################## 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.50, 0.52])
plt.savefig('k-balance_diss-iso-10000-channel-kom.png',bbox_inches='tight')