;----------------------------------------------------------------------
; shapefiles_14_mask.ncl
;
; Concepts illustrated:
;   - Reading shapefiles
;   - Adding shapefile outlines to an existing WRF contour/map plot
;   - Masking a data array based on a geographical area obtained from a shapefile
;   - Drawing a WRF lat/lon grid using gsn_coordinates
;   - Zooming in on a WRF map
;----------------------------------------------------------------------
; This example shows how to use a shapefile of the United States
; to mask out all data except for over a select group of states.
;
; The "USA_adm.shp" shapefiles were downloaded from 
; http://www.gadm.org/country/
;----------------------------------------------------------------------
; These files are loaded by default in NCL V6.2.0 and newer
; load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"
; load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl"
; load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"
; load "$NCARG_ROOT/lib/ncarg/nclscripts/wrf/WRFUserARW.ncl"
;----------------------------------------------------------------------
; This script was updated Dec 2018 to use wrf_user_ll_to_xy instead of 
; wrf_user_ll_to_ij, which will be deprecated in NCL V6.6.0. The
; difference between the two routines is that wrf_user_ll_to_xy returns
; 0-based indexes, and wrf_user_ll_to_ij returns 1-based indexes.
;----------------------------------------------------------------------

load "./shapefile_utils.ncl"

begin
  wrf_filename  = "wrfout_d01_2010-07-15_20:00:00.nc"
;  wrf_filename  = "o3.month.nc"
  shp_filename1 = "gadm36_USA_1.shp"    ; State outlines
  shp_filename2 = "gadm36_USA_2.shp"    ; County outlines
  a = addfile(wrf_filename,"r")
  
;---Area to zoom in on
  lats = (/  38.5, 50.0/)
  lons = (/-111.262,-89.0/)
  loc  = wrf_user_ll_to_ij(a, lons, lats, True)

; loc(0,:) is west-east (x)
; loc(1,:) is south-north (y)
; Use "loc" values to set special resources for zooming in on map.


; wrf_user_ll_to_ij deprecated in NCL V6.6.0. 
; Index values are returned as 1 to N, so we have to
; subtract 1 for 0 to N-1 indexing that NCL requires.
;
  loc  = wrf_user_ll_to_ij(a, lons, lats, True)
  loc = loc-1       

  xstart = loc(0,0)   ; Set the zoom in coordinates
  xend   = loc(0,1)
  ystart = loc(1,0)
  yend   = loc(1,1)

;---Read "height" variable and lat/lon coordinates off WRF output file.
  nt        = 0                 ; First time step
  o3        = wrf_user_getvar(a,"o3",0)*1000            ; ozone (ppb)
  o3_lev0   = o3(0,:,:)                    ; get the first level
  o3_lev0@lat2d = a->XLAT(nt,:,:)
  o3_lev0@lon2d = a->XLONG(nt,:,:)

;---Set all tc_lev0 values to missing except for those over the select states.
  opt             = True
  opt@debug       = True
  opt@shape_var   = "NAME_1"
  opt@shape_names = (/"North Dakota","South Dakota","Nebraska","Minnesota","Montana", "Iowa", "Wyoming"/)
  o3_mask  = shapefile_mask_data(o3_lev0,shp_filename1,opt)

;---Zoom in on area of interest.
  o3_mask_zoom = o3_mask(ystart:yend,xstart:xend)
  
;---Start the graphics
  wks = gsn_open_wks("png","shapefiles_mask") ; send graphics to PNG file

  res                   = True               ; Use basic options for this field
  res@cnFillOn          = True               ; Create a color fill plot
  res@MainTitle         = "Surface Ozone"
;  res@lbTitlePosition  = "Bottom"
  res@lbLabelAutoStride = True
  res@cnFillPalette     = "WhiteBlueGreenYellowRed"
  res@lbBoxSeparatorLinesOn = "False"
;  res@lbAutoManage     = "True"
  res@lbBottomMarginF   = -0.15
  res@FieldTitle        =  "O3 (ppb)"
  res@lbTitleOn         =  False 
;  res@ValidTime         = "True"
  res@ContourParameters = (/0,90,4/)      ; Contour levels

;---Create contours of masked and zoomed in masked "tc_lev0" arrays
  contour_mask      = wrf_contour(a,wks,o3_mask,res)

  res@pmLabelBarOrthogonalPosF  = -0.10     ; Move labelbar down slightly
  contour_mask_zoom = wrf_contour(a,wks,o3_mask_zoom,res)
  

  pltres            = True   ; Set plot options
  pltres@PanelPlot  = True   ; Tells wrf_map_overlays not to remove contours

  mpres             = True   ; Set map options
  mpres@mpOutlineOn = False
  mpres@mpFillOn    = False

  plot_mask      = wrf_map_overlays(a,wks,contour_mask,pltres,mpres) 

;---Create a zoomed version
  mpres@ZoomIn = True        ; Tell wrf_map_resources we want to zoom in.
  mpres@Xstart = xstart
  mpres@Xend   = xend
  mpres@Ystart = ystart
  mpres@Yend   = yend

  plot_mask_zoom = wrf_map_overlays(a,wks,contour_mask_zoom,pltres,mpres) 
 
  

;---Attach the shapefile outlines
  lnres                  = True
  lnres@gsLineColor      = "gray25"
  lnres@gsLineThicknessF = 0.5   

;---Using USA_adm2.shp here to get county outlines
  id_mask      = gsn_add_shapefile_polylines(wks,plot_mask,shp_filename2,lnres)
  id_mask_zoom = gsn_add_shapefile_polylines(wks,plot_mask_zoom,shp_filename2,lnres)

  draw(plot_mask)
  frame(wks)

;---Mask the lat/lon values over the same area as the data.
  lat2d_mask = o3_lev0@lat2d
  lon2d_mask = o3_lev0@lon2d
  lat2d_mask@_FillValue = default_fillvalue(typeof(lat2d_mask))
  lon2d_mask@_FillValue = default_fillvalue(typeof(lon2d_mask))
  lat2d_mask = where(ismissing(o3_mask),lat2d_mask,lat2d_mask@_FillValue)
  lon2d_mask = where(ismissing(o3_mask),lon2d_mask,lon2d_mask@_FillValue)

  mkres                 = True
  mkres@gsMarkerSizeF   = 0.001
  mkres@gsMarkerColor   = "gray50"
  mkres@gsnCoordsAttach = True
  mkres@gsnCoordsLat    = lat2d_mask
  mkres@gsnCoordsLon    = lon2d_mask

;---Attach coordinates
  gsn_coordinates(wks,plot_mask_zoom,o3_mask,mkres)

  draw(plot_mask_zoom)
  frame(wks)
end