;----------------------------------------------------------------------
; This script creates a mask array based on a shapefile containing an
; outline of the Mississippi River Basin shapefile, and on WRF
; data read off a WRF output file.
;
; The mask array is written to a NetCDF file, so you can use it
; again to mask other WRF variables that are on the same lat/lon grid.
;----------------------------------------------------------------------
; 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/wrf/WRFUserARW.ncl"

load "./shapefile_utils.ncl"

;----------------------------------------------------------------------
; Use shapefile_mask_data to mask the given variable against the 
; given shapefile, and return a mask array of 0s and 1s.
;
; The shapefile_mask_data function can be VERY SLOW if you have a lot
; of points in the shapefile and/or the original data variable. By
; returning a mask array of 0s and 1s, you can use this array to 
; mask other variables by using the much faster "mask" function.
;----------------------------------------------------------------------
function create_mask(var,shapefile_name)
local opt, start_time, end_time
begin
  start_time = get_cpu_time()

  opt             = True
  opt@return_mask = True    ; This forces the return of a 0s and 1s mask array
; opt@debug       = True
  var_mask        = shapefile_mask_data(var,shapefile_name,opt)

  end_time = get_cpu_time()

  print("======================================================================")
  print("It took " + (end_time-start_time) + " CPU seconds to create mask array.")
  print("======================================================================")

  return(var_mask)
end

;----------------------------------------------------------------------
; Main code
;----------------------------------------------------------------------
begin
;---Open WRF output file
  filename = "wrfout_d01_2005-12-14_13:00:00"
  a        = addfile(filename,"r")

;---Read several WRF variables at first time step
  it    = 0
  tc    = wrf_user_getvar(a,"tc",it)      ; 3D temperature
  lat2d = wrf_user_getvar(a,"lat",it)     ; 2D
  lon2d = wrf_user_getvar(a,"lon",it)     ; 2D

;---Now get the lowest (bottommost) level
  nl  = 0
  tc2 = tc(nl,:,:)
  tc2@lat2d = lat2d    ; This is needed for NCL masking and plotting only!
  tc2@lon2d = lon2d    ; Do not write these attributes to a file.

;---Plot the original data
  wks = gsn_open_wks("x11","wrf_gsn")

  res                = True
  res@gsnMaximize    = True
  res@cnFillOn       = True  
  res@lbOrientation  = "Vertical"
  res = wrf_map_resources(a,res)

  plot = gsn_csm_contour_map(wks,tc2,res)

;---Create a mask array (0s and 1s) so we can apply it to variables later.
  shp_name = "mrb.shp"
  mask_array = create_mask(tc2,shp_name)

;---Apply mask to tc2 and plot. Note how fast this is compared to "create_mask"
  start_time = get_cpu_time()
  tc2_mask   = mask(tc2,mask_array,1)
  copy_VarMeta(tc2,tc2_mask)         ; This copies metadata from tc2 to new mask variable
  end_time = get_cpu_time()

  printVarSummary(tc2)
  printVarSummary(tc2_mask)     ; tc2_mask will have lat2d/lon2d and other metadata

  print("======================================================================")
  print("It took " + (end_time-start_time) + " CPU seconds to mask tc2.")
  print("======================================================================")

  plot_mask = gsn_csm_contour_map(wks,tc2_mask,res)

;---Read slp, mask it, and plot it.
  slp       = wrf_user_getvar(a,"slp",it)  ; sea level pressure
  slp@lat2d = lat2d                        ; needed for plotting
  slp@lon2d = lon2d
  slp_mask = mask(slp,mask_array,1)
  copy_VarMeta(slp,slp_mask)               ; copy metadata

  printVarSummary(slp)
  printVarSummary(slp_mask)

  plot_orig = gsn_csm_contour_map(wks,slp,res)
  plot_mask = gsn_csm_contour_map(wks,slp_mask,res)

;---Read terrain, mask it, and plot it
  ter      = wrf_user_getvar(a,"ter",it)  ; sea level pressure
  ter@lat2d = lat2d    ; Needed for plotting
  ter@lon2d = lon2d
  ter_mask = mask(ter,mask_array,1)
  copy_VarMeta(ter,ter_mask)  ; This copies all metadata
  printVarSummary(ter)
  printVarSummary(ter_mask)

  plot_orig = gsn_csm_contour_map(wks,ter,res)
  plot_mask = gsn_csm_contour_map(wks,ter_mask,res)

;---Write mask to file so we can read it later from another script, if desired.
  mask_filename = "wrf_mask.nc"
  system("rm -f " + mask_filename)
  fout = addfile(mask_filename,"c")
  fout->WRF_mask = mask_array

end