#!python3
# plot vorticity in GOMb0.04 configuation
import sys
prfx='/home/abozec/PYTHON/'
sys.path.append(prfx)

from myutilities.tvplus import tvplus
from hycom.info import read_field_names,read_field_grid_names
from hycom.io import read_hycom_fields, read_hycom_grid,sub_var2
from myutilities.get_cmap import get_cmapbluered
import numpy as np
import netCDF4 as nc
import proplot as plot

from matplotlib.colors import BoundaryNorm
from scipy import ndimage

prfx='/home/abozec/'
#sys.path.append(prfx+'PYTHON/')

# get grid
iot='/Net/ugos-tops/tops_2025/IAS96_S/SCRATCH/'

year='2025'
d=2
h=0
for d in [190]:
#for d in my.np.arange(142)+225:
#for h in [12]:
#for h in my.np.arange(24):
   day='{:03d}'.format(d)
   hour='{:02d}'.format(h)
   
   io_data = iot
   file1='archm.'+year+'_'+day+'_'+hour+'.nc'
   file_data = io_data+file1
   print(file_data)
   
   # get metadata of .[ab]
   # Printing the fields available in the file
   #print(F"The fields available are: {read_field_names(file_data)}")
   fields=['uvel','vvel']
   hycom_field= nc.Dataset(file_data)

   # get grid
   plon=hycom_field['lon'][:]
   plat=hycom_field['lat'][:]
   idm=plon.shape[0] ; jdm=plat.shape[0]
   pscx=1043. ; pscy=1159. # approximate for regular grid

   
   # get u and v
   ubc=hycom_field[fields[0]][0,0,0:jdm,0:idm].data
   vbc=hycom_field[fields[1]][0,0,0:jdm,0:idm].data
   #ub=hycom_field[fields[2]][0,0:jdm,0:idm]
   #vb=hycom_field[fields[3]][0,0:jdm,0:idm]
   print(ubc.shape)
   #print(ub.shape)
   #u=ub+ubc
   #v=vb+vbc
   u=ubc
   v=vbc
   # Get mask
   u[u > 1e10] = np.nan
   v[v > 1e10] = np.nan
   
   # get vorticity
   vort=np.zeros([jdm,idm])
   vort[:,:]=(v - np.roll(v,1,axis=1))/pscx - (u - np.roll(u,1,axis=0))/pscy
   #tvplus(vort)
  
   # colorbar
   cmap_u=get_cmapbluered()
   N=255
   levels_u=np.linspace(-50,50,101)
   norm = BoundaryNorm(levels_u, ncolors=cmap_u.N, clip=True)
   # to get it right in pcolormesh... contourf does not do well with high resolution
   vort=np.nan_to_num(vort,nan=0.)
   y = ndimage.uniform_filter(vort*1e6,size=4)

   # plot the field
   # change land color
   plot.rc.landcolor=[150/256,150/256,150/256]
   plot.rc.reso = 'med'  # use higher res for zoomed in geographic features
   fig, axs = plot.subplots(nrows=1,ncols=1,proj='cyl',refwidth='12cm')
   #CS2=axs[0].contourf(plon,plat,vort*1e6, cmap=cmap_u, extend='both',\
   #              levels=levels_u)
   CS2 = axs[0].pcolormesh(plon[:], plat[:], y, cmap=cmap_u, norm=norm, \
                               shading='auto', extend='both')
   axs[0].format(latlines=2, lonlines=2,lonlim=(-99, -78), latlim=(18, 32),\
                   labels=True,land=True,\
              title='IAS96 Vorticity (1e6 s-1) Year:'+year+' Day:'+day+' Hour:'+hour)
   
   #axs[0].contour(lon1,lat1,vort*1e5, color='k', levels=levels_u[::2],lw=0.2)
   axs[0].colorbar(CS2,loc='b',ticks=10)

   
   ps_dir=prfx+'/PYTHON/IAS96/PNG/'
   file_png='IAS96_vort_'+year+'_'+day+'_'+hour+'.png'
   print(ps_dir+file_png)
   fig.savefig(ps_dir+file_png,dpi=300,\
               facecolor='w', edgecolor='w',transparent=False,bbox_inches='tight') ## .pdf,.eps
   plot.close() ##initcond
   #plot.show()