;----------------------------------------------------------------------
; This script does three things:
; 1. Creates a mask array based on one variable on a WRF output file.
;
; 2. Uses the mask array to mask other variables on the same lat/lon grid 
;
; 3. Writes the masked variables and remaining variables to a new 
;    NetCDF file.
;----------------------------------------------------------------------
; 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"

;----------------------------------------------------------------------
; Main code
;----------------------------------------------------------------------

begin

;----------------------------------------------------------------------
; Step 1: Create the mask array
;----------------------------------------------------------------------

;---Open WRF output file and read a variable to use for creating a mask array
  wrf_filename = "wrfout_d01_2005-12-14_13:00:00"
  wf           = addfile(wrf_filename,"r")

  t    = wf->T           ; Time, bottom_top, south_north, west_east 
  tsub = t(0,0,:,:)      ; Subset just the lat/lon portion of the variable
                         ; NOTE: If your variable's lat/lon grid changes
                         ; for each time step, then you need to create 
                         ; a mask variable for each time step.

  tsub@lat2d = wf->XLAT(0,:,:)     ; necessary for masking against shapefile
  tsub@lon2d = wf->XLONG(0,:,:)

;---Create the mask based on the given shapefile
  print("Creating a mask array...this could take awhile...")
  shp_filename    = "mrb.shp"
  opt             = True
  opt@return_mask = True    ; This forces the return of a 0s and 1s mask array
  mask_array      = shapefile_mask_data(tsub,shp_filename,opt)
  print("Done creating mask array")

;---Open new NetCDF file to write masked variables to.
  output_filename = wrf_filename + "_MASKED.nc"
  system("rm -f " + output_filename)
  fout = addfile(output_filename,"c")

;----------------------------------------------------------------------
; Steps 2 and 3.
;
; Loop through each variable on the file and check if this is one of
; the ones we want to mask.
;
; If we do, then mask it and write to the new NetCDF file.
;
; If we don't, then write original variable to the new NetCDF file.
;----------------------------------------------------------------------
  notes_string = "This variable was masked using the " + shp_filename + " shapefile"

;---Indicate the variables you want to mask
  vars_to_mask = (/"T","QSNOW","TSLB","SST","TSK"/)

  varnames     = getfilevarnames(wf)   ; Get all variable names from the file
  nvars        = dimsizes(varnames)
;
; Loop across each variable on the file. If it's one of the variables 
; we want to mask, then mask this variable before writing it to a new 
; NetCDF file. 
;
; Otherwise, write the original variable to the NetCDF file.
;
  do nv=0,nvars-1
    vsizes := getfilevardimsizes(wf,varnames(nv))
    vnames := getfilevardims(wf,varnames(nv))
    rank    = dimsizes(vsizes)

    print("==================================================")
    print("    Variable '" + varnames(nv) + "'")
    print("    Dimension sizes: " + str_join(""+vsizes,","))
    print("    Dimension names: " + str_join(vnames,","))
;
; If this is one of the variables we want to mask, then mask
; the variable and write it to the new NetCDF file.
;
    if(any(varnames(nv).eq.(vars_to_mask))) then
;---Write masked variable to file
      print("    WILL MASK THIS VARIABLE")
      if(rank.gt.2)
        mask_array_conform := conform_dims(vsizes,mask_array,(/rank-2,rank-1/))
        var_mask           := mask(wf->$varnames(nv)$,mask_array_conform,1)
       else
        var_mask := mask(wf->$varnames(nv)$,mask_array,1)
       end if

       copy_VarMeta(wf->$varnames(nv)$,var_mask)    ; Copy metadata from var to var_mask
       var_mask@NOTES = notes_string                ; Not necessary, but useful
       fout->$varnames(nv)$ = var_mask              ; Write masked var to file
     else
;---Write original variable to file
      print("    WILL WRITE ORIGINAL VARIABLE TO FILE.")
      fout->$varnames(nv)$ = wf->$varnames(nv)$
    end if
  end do
end