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"

;----------------------------------------------------------------------
; Create a lat/lon box array based on min/mat lat/lon of 
; the given shapefile outlines.
;
; Note that we are not simply doing a box with 5 points for the box.
; We need to use this box to mask against a data array, so we 
; need to define more points along each edge of the box to
; make the masking more effective.
;----------------------------------------------------------------------
function get_latlon_box(shp_filename,lat_margin,lon_margin)
local f, lat,lon, minlat, maxlat, minlon, maxlon, npts
begin
;---Read shapefile lat/lon
  f = addfile(shp_filename,"r")
  lon = f->x
  lat = f->y

;---The margins makes the lat/lon box a little bigger.
  minlat = min(lat) - lat_margin
  maxlat = max(lat) + lat_margin
  minlon = min(lon) - lon_margin
  maxlon = max(lon) + lon_margin

  npts    = 10    ; # of points per one edge of the lat/lon box
  lat_box = new(4*npts,typeof(minlat)) 
  lon_box = new(4*npts,typeof(minlon)) 

  lat_box(0:npts-1)        = minlat
  lon_box(0:npts-1)        = fspan(minlon,maxlon,npts)
  lat_box(npts:2*npts-1)   = fspan(minlat,maxlat,npts)
  lon_box(npts:2*npts-1)   = maxlon
  lat_box(2*npts:3*npts-1) = maxlat
  lon_box(2*npts:3*npts-1) = fspan(maxlon,minlon,npts)
  lat_box(3*npts:4*npts-1) = fspan(maxlat,minlat,npts)
  lon_box(3*npts:4*npts-1) = minlon
  return([/lat_box,lon_box/])
end

;----------------------------------------------------------------------
; Function to mask a data array based on a rectangle that encloses
; the given shapefile outline(s).
;----------------------------------------------------------------------
function mask_data_shapefile_rectangle(data,shp_filename,lat_margin,lon_margin)
local f, latlon, dims, ranks, lat, lon, minlat, maxlat, minlon, maxlon, lat_box, lon_box
begin
  latlon  = get_latlon_box(shp_filename,lat_margin,lon_margin)
  lat_box = latlon[0]
  lon_box = latlon[1]

;---Determine the grid type
  dims = dimsizes(data)
  rank = dimsizes(dims)

  grid_type = ""
  if(rank.eq.2.and.\
     isdimnamed(data,0).and.iscoord(data,data!0).and.\
     isdimnamed(data,1).and.iscoord(data,data!1)) then
;---Rectilinear
    lat1d = ndtooned(conform_dims(dims,data&$data!0$,0))
    lon1d = ndtooned(conform_dims(dims,data&$data!1$,1))
    grid_type = "rectilinear"
  else if(rank.eq.2.and.all(isatt(data,(/"lat2d","lon2d"/)))) then
;---Curvilinear
    lat1d = ndtooned(data@lat2d)
    lon1d = ndtooned(data@lon2d)
    if(product(dims).eq.dimsizes(lat1d).and.\
       product(dims).eq.dimsizes(lon1d)) then
      grid_type = "curvilinear"
    end if
  else if(rank.eq.1.and.all(isatt(data,(/"lat1d","lon1d"/)))) then
;---Unstructured
    lat1d = data@lat1d
    lon1d = data@lon1d
    if(dims.eq.dimsizes(lat1d).and.\
       product(dims).eq.dimsizes(lon1d)) then
      grid_type = "unstructured"
    end if
  end if
  end if
  end if

  if(grid_type.eq."") then
    print("mask_data_shapefile_rectangle: Error: not a valid rectilinear, curvilinear, or unstructured grid")
    exit
  end if
  nlatlon1d = dimsizes(lat1d)

;---Create array to hold new data mask
  data_mask_1d = new(nlatlon1d,integer)
  data_mask_1d = 0

  minlat = min(lat_box)
  maxlat = max(lat_box)
  minlon = min(lon_box)
  maxlon = max(lon_box)
  do n=0,nlatlon1d-1
    if(gc_inout(lat1d(n),lon1d(n),lat_box,lon_box)) then
      data_mask_1d(n) = 1
    end if
  end do

;
; Create a 2D data array of same size as original data,
; but with appropriate values masked. Create a missing
; value if our data doesn't have one.
;
  if(.not.isatt(data,"_FillValue")) then
    data_msg = default_fillvalue(typeof(data))
  else
    data_msg = data@_FillValue
  end if

  data_mask = where(onedtond(data_mask_1d,dims).eq.1,data,data_msg)
  copy_VarMeta(data,data_mask)      ; Copy all metadata

  if(.not.isatt(data,"_FillValue")) then
    data_mask@_FillValue = data_msg
  end if

  return(data_mask)
end

;----------------------------------------------------------------------
; Procedure to add lat/lon box outlines to the given plot.
;----------------------------------------------------------------------
procedure add_latlon_box(wks,plot,shp_filename,lat_margin,lon_margin)
local lnres, latlon, lat_box, lon_box
begin
  latlon  = get_latlon_box(shp_filename,lat_margin,lon_margin)
  lat_box = latlon[0]
  lon_box = latlon[1]

  lnres                  = True
  lnres@gsLineColor      = "Red"
  lnres@gsLineThicknessF = 3.0            ; 3x thickness

  plot@box = gsn_add_polyline(wks, plot, lon_box, lat_box, lnres)
end

;----------------------------------------------------------------------
; Procedure to add shapefile outlines to the given plot.
;----------------------------------------------------------------------
procedure add_shp_outlines(wks,plot,shp_filename)
local lnres
begin
;---Resources for polyline
  lnres                  = True
  lnres@gsLineColor      = "Red"
  lnres@gsLineThicknessF = 3.0            ; 3x thickness

  plot@lines = gsn_add_shapefile_polylines(wks, plot, shp_filename, lnres)
end

;----------------------------------------------------------------------
; Function to create ummy rectilinear data over a map area of interest.
;----------------------------------------------------------------------
function create_dummy_array(minlat,maxlat,minlon,maxlon,nlon,nlat)
local lat1d, lon1d
begin
;---Generate some dummy data over map
  data            = generate_2d_array(10, 10, 0., 100., 0, (/nlat,nlon/))
  data@_FillValue = -9999

;---Add lat/lon coordinate array information.
  lat1d       = fspan(minlat,maxlat,nlat)
  lon1d       = fspan(minlon,maxlon,nlon)
  lat1d@units = "degrees_north"
  lon1d@units = "degrees_east"
  data!0      = "lat"
  data!1      = "lon"
  data&lat    = lat1d
  data&lon    = lon1d

  return(data)
end

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

begin
;---Create dummy array
  GRID_TYPE = "coarser"     ; "coarser"

  minlat = 20
  maxlat = 60
  minlon = -125
  maxlon = -60
  if(GRID_TYPE.eq."finer") then
    nlat = 64
    nlon = 128
  else
    nlat = 32
    nlon = 64
  end if
  data = create_dummy_array(minlat,maxlat,minlon,maxlon,nlat,nlon)

;---Start the graphics
  wks  = gsn_open_wks("png",GRID_TYPE + "_rect")

  res                     = True

  res@gsnMaximize         = True           ; maximize plot in frame
  res@gsnDraw             = False          ; don't draw plot yet
  res@gsnFrame            = False          ; don't advance frame yet

  res@cnFillOn            = True
  res@cnFillPalette       = "MPL_Set3"
 res@cnFillOpacityF      = 0.5            ; Fade the plot slightly
  res@cnLinesOn           = False
  res@cnLineLabelsOn      = False

  res@cnLevelSelectionMode = "ExplicitLevels"
  res@cnLevels             = ispan(6,96,4)
  res@lbLabelBarOn         = False

  res@gsnAddCyclic        = False

  res@mpDataBaseVersion   = "MediumRes"    ; slightly better resolution
  res@mpFillOn            = False

;---Zoom in on area of interest
  res@mpMinLatF           = minlat
  res@mpMaxLatF           = maxlat
  res@mpMinLonF           = minlon
  res@mpMaxLonF           = maxlon

  res@tiMainString        = "Original data"
  plot_orig = gsn_csm_contour_map(wks,data,res)        ; Create plot of original data

;---Use shapefile to mask this dummy array based on shapefile outline
  lat_margin  = 1
  lon_margin  = 1

  dir          = ncargpath("data") + "/shp/"
  shp_filename = dir + "mrb.shp"
  data_mask    = mask_data_shapefile_rectangle(data,shp_filename,lat_margin,lon_margin)

  res@tiMainString = "Data with rectangle mask (" + GRID_TYPE + " grid)"
  plot_mask = gsn_csm_contour_map(wks,data_mask,res)   ; Create plot of masked data

;---Add shapefile and box outlines to both plots
  add_shp_outlines(wks,plot_orig,shp_filename)
  add_shp_outlines(wks,plot_mask,shp_filename)

  add_latlon_box(wks,plot_orig,shp_filename,lat_margin,lon_margin)
  add_latlon_box(wks,plot_mask,shp_filename,lat_margin,lon_margin)

;---Add dots at the lat/lon grid locations
  mkres                 = True
  mkres@gsMarkerIndex   = 16         ; Filled dots
  mkres@gsMarkerSizeF   = 0.001      ; Make them small
  mkres@gsnCoordsAttach = True
  mkres@gsnCoordsNonMissingColor = "black"
  mkres@gsnCoordsMissingColor    = "black"
  gsn_coordinates(wks,plot_orig,data,mkres)
  gsn_coordinates(wks,plot_mask,data_mask,mkres)

  pres                  = True
  pres@gsnMaximize      = True
  pres@gsnPanelLabelBar = True
  gsn_panel(wks,(/plot_orig,plot_mask/),(/2,1/),pres)
end