#!python3
# Calculate the kinetic energy of first layer
import sys

prfx='/home/abozec/PYTHON/'
sys.path.append(prfx)
import myenv as my

# get grid
kdm=2
idm=101 ; jdm=101
dy=dx=20.0e3 ## meters
# idm=201 ; jdm=201
# dy=dx=10.0e3 ## meters
mask=my.np.ones([jdm,idm]) ## no bathy
# area
area=dx*dy*mask
area_tot=my.np.nansum(area[:,:])

# read U
io='/nexsan/people/abozec/NUMO/REF/BB86/'
file='bb86_20km_beta_freeslip_1500m_visco500.nc' # m/s
#file='bb86_10km_beta_freeslip_1500m_visco250.nc' # m/s
ds=my.nc.Dataset(io+file)
u=ds['u'][:,:,:,:]
v=ds['v'][:,:,:,:]
dp=ds['dp'][:,:,:,:]
nt=dp.shape[0]

# area/volume
vol=my.np.zeros([nt,kdm,jdm,idm])
for t in my.np.arange(nt):
   for k in my.np.arange(kdm):
       vol[t,k,:,:]=area[:,:]*dp[t,k,:,:]


## define arrays
ke=my.np.zeros([nt])
ke1=my.np.zeros([nt])
for t in my.np.arange(nt):
   tmp_vol=my.np.nansum(vol[t,:,:,:])
   ke[t]=my.np.nansum(0.5*(u[t,:,:,:]**2+v[t,:,:,:]**2)*vol[t,:,:,:])/tmp_vol
   tmp_vol=my.np.nansum(vol[t,0,:,:])
   ke1[t]=my.np.nansum(0.5*(u[t,0,:,:]**2+v[t,0,:,:]**2)*vol[t,0,:,:])/tmp_vol


## plot results
time=my.np.arange(nt)/360.

fig, ax = my.plot.subplots(ncols=1,nrows=1,aspect=3, width='15cm')
ax[0].plot(time,ke*1e4, lw=2,legend='ul',labels=['KE-tot'],legend_kw={'ncols':1})
ax[0].plot(time,ke1*1e4, lw=2,legend='ul',labels=['KE-layer 1'],legend_kw={'ncols':1})
ax[0].format(title='Kinetic Energy BB86 20km', xlabel='Time (year)', \
          ylabel='K.E. (cm2/s2)',ylim=(0,15))

ps_dir='/nexsan/people/abozec/NUMO/MICHAL/VTK_REPORT2/PS/'
file_ps='time_KE_bb86_20km.pdf'
fig.savefig(ps_dir+file_ps,dpi=150,\
        facecolor='w', edgecolor='w',transparent=False) ## .pdf,.eps


my.plot.show()