#!python3
# Calculate the EKE of first layer
import sys

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

# 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)
# get all the time step of u and v m/s
u=ds['u'][:,0,:,:]
u2=u**2
v=ds['v'][:,0,:,:]
v2=v**2
ke=0.5*(u2+v2) ## kinetic every time step

dp=ds['dp'][:,0,:,:] # meters

# get umean and vmean to get ke of mean flow over last year (360 days)
t1=3240 ; t2=3600
#t1=0; t2=3600
um=my.np.nansum(u[t1:t2,:,:]*dp[t1:t2,:,:],axis=0)/my.np.nansum(dp[t1:t2,:,:],axis=0)
vm=my.np.nansum(v[t1:t2,:,:]*dp[t1:t2,:,:],axis=0)/my.np.nansum(dp[t1:t2,:,:],axis=0)
kem=0.5*(um**2+vm**2)

## eke=mean of ke - ke of mean flow
# mean of ke of last year
mke=my.np.nansum(ke[t1:t2,:,:]*dp[t1:t2,:,:],axis=0)/my.np.nansum(dp[t1:t2,:,:],axis=0)
# eke
eke = mke - kem # m2/s2


# colormap
cmapeke=my.mygc.get_cmapeke()
N=cmapeke.colors.shape[0]
cmap=cmapeke

vmin=0.;vmax=300.
levels_to_draw=my.np.linspace(vmin,vmax,N+1)
# plot the field
fig, axs = my.plot.subplots(nrows=1,axwidth='4.5in')
axs[0].format(title='BB86 EKE (cm2/s2) 20km  Y:10')
m = axs[0].contourf(eke*1e4, cmap=cmap, extend='neither',\
                    levels=levels_to_draw)
axs[0].colorbar(m,loc='b',ticks=50.)

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

my.plot.show() ## interactive (block=False), otherwise ()