#----------------------------------------------------------------------
# This Python script plots uocn and vocn from a HYCOM model file using
# curly vectors. PyNIO is used to read the data, and PyNGL is used to
# plot it.
#----------------------------------------------------------------------
import os,sys,numpy
from numpy import ma
import Ngl,Nio


#---Open NetCDF file
dir      = os.path.join("..","..","Data","HYCOM")
filename = "hycom-cice_ARCu0.08_035_UOA_2012010100_t000.nc"
a = Nio.open_file(os.path.join(dir,filename))

#---Get variables and their offset/scale_factor attributes
u       = a.variables["uocn"]
v       = a.variables["vocn"]
h       = a.variables["hi"]
uoffset = u.attributes['add_offset']
voffset = v.attributes['add_offset']
uscale  = u.attributes['scale_factor']
vscale  = v.attributes['scale_factor']
hoffset = h.attributes['add_offset']
hscale  = h.attributes['scale_factor']

#---Convert to float 
uf = numpy.array(u[:],'f')
vf = numpy.array(v[:],'f')
hf = numpy.array(h[:],'f')

#---Apply scale_factor and add_offset attributes to data
uf = numpy.where(uf!=u._FillValue,(uf*uscale)+uoffset,u._FillValue)
vf = numpy.where(vf!=v._FillValue,(vf*vscale)+voffset,v._FillValue)
hf = numpy.where(hf!=h._FillValue,(hf*hscale)+hoffset,h._FillValue)

#---Mask the missing values
umask = ma.masked_values(uf[0,:,:],u._FillValue)
vmask = ma.masked_values(vf[0,:,:],v._FillValue)
hmask = ma.masked_values(hf[0,:,:],h._FillValue)

#---Get lat/lon and fix the longitudes
lat     = a.variables["lat"][:]
lon     = a.variables["lon"][:]
lon     = numpy.where((lon < 0),lon+360,lon)

#---Start the graphics
wks = Ngl.open_wks("png","vector_ngl_ocn_14")   # "ps", "pdf", "png"
Ngl.define_colormap(wks,"WhiteBlue")
gray = Ngl.set_color(wks,255,0.88,0.88,0.88)       # add gray to colormap

#---Set some resources for drawing vectors
vcres                      = Ngl.Resources()

vcres.nglDraw              = False    # Don't draw plot
vcres.nglFrame             = False    # or advance frame.

vcres.vcMinFracLengthF     = 0.05     # length of min vector as fraction
vcres.vcRefLengthF         = 0.05     # default is 0.00903
vcres.vcRefMagnitudeF      = 0.30     # default is 0.723
vcres.vcMinDistanceF       = 0.008     # thin the vectors
vcres.vcGlyphStyle         = "CurlyVector"     # turn on curly vectors

vcres.vfXArray             = lon
vcres.vfYArray             = lat

vcres.vcRefAnnoOrthogonalPosF = -0.27   # Move label box up into plot.

#---Create vector plot only
vector = Ngl.vector(wks,umask,vmask,vcres)

#---Set some contour and map resources
cnres                      = Ngl.Resources()

cnres.nglDraw              = False    # Don't draw plot
cnres.nglFrame             = False    # or advance frame.

cnres.cnFillOn             = True
cnres.cnLinesOn            = False
cnres.cnLineLabelsOn       = False
cnres.cnLevelSpacingF      = 0.2
cnres.lbOrientation        = "Horizontal"
cnres.nglSpreadColorEnd    = -2
cnres.nglSpreadColorStart  = 64

#---Zoom in on map.
cnres.mpLimitMode          = "LatLon"
cnres.mpMinLatF            =   60
cnres.mpMaxLatF            =   70
cnres.mpMinLonF            = -180
cnres.mpMaxLonF            = -160

cnres.mpDataBaseVersion    = "MediumRes"      # Better resolution
cnres.mpCenterLonF         = (cnres.mpMinLonF + cnres.mpMaxLonF) / 2.

cnres.mpOutlineOn          = True
cnres.mpFillOn             = True
cnres.mpOceanFillColor     = "transparent"
cnres.mpLandFillColor      = "gray"
cnres.mpGridAndLimbOn      = False

cnres.tiMainString         = "Velocity vectors (m/s) over sea ice thickness (m)"
cnres.tiMainFont           = "Helvetica-bold"
cnres.tiMainFontHeightF    = 0.015

cnres.sfXArray             = lon
cnres.sfYArray             = lat

contour = Ngl.contour_map(wks,hmask,cnres)   # Create contours over a map
Ngl.overlay(contour,vector)                  # Overlay vectors on contour/map plot
Ngl.draw(contour)                            # Draw vectors, contours, and map
Ngl.frame(wks)                               # Page advance

Ngl.end()