; This script assumes you have a rectilinear grid. That is, your 
; data has 1D coordinate arrays. It will work with 2D (curvilinear)
; data or 1D unstructured data, but you will need to modify the 
; "mask_data_with_gadm_country" function to correctly handle the
; lat/lon data.
;----------------------------------------------------------------------

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 "./shapefile_utils.ncl"

;----------------------------------------------------------------------
; This function masks a rectilinear data array against a country 
; outline. You need to pass in the three letter country 
; abbreviation, like; "BGD" (Bangladesh), "CHN" (China), etc.
;----------------------------------------------------------------------
function mask_data_with_gadm_country(country_code,data)
begin
  mask_start_time = get_cpu_time()     ; For timing results

;---Convert rectilinear grid to 2D grid, but laid out as 1D array.
  dims  = dimsizes(data)
  lat1d = ndtooned(conform_dims(dims,data&$data!0$,0))
  lon1d = ndtooned(conform_dims(dims,data&$data!1$,1))

;---Open shapefile and read lat/lon values.
  sfilename = country_code +"_adm1.shp"
  f = addfile(sfilename,"r")

;---Read data off the shapefile
  segments = f->segments
  geometry = f->geometry
  segsDims = dimsizes(segments)
  geomDims = dimsizes(geometry)

;---Read global attributes  
  geom_segIndex = f@geom_segIndex
  geom_numSegs  = f@geom_numSegs
  segs_xyzIndex = f@segs_xyzIndex
  segs_numPnts  = f@segs_numPnts
  numFeatures   = geomDims(0)

;---Read sapefile lat/lon
  lon  = f->x
  lat  = f->y
  nlatlon = dimsizes(lat)

;
; Get the approximate lat/lon box that encloses all of the
; given country. This will make our checking go faster.
;
  min_lat = min(lat)
  max_lat = max(lat)
  min_lon = min(lon)
  max_lon = max(lon)

  print("==================================================")
  print("Shapefile : "  + sfilename)
  print("min_lat " + min_lat)
  print("max_lat " + max_lat)
  print("min_lon " + min_lon)
  print("max_lon " + max_lon)

;
; Get the approximate index values that contain the area of interest.
;
; This will make our gc_inout loop below go faster, because we're
; not checking every single lat/lon point to see if it's within
; the area of interest.
;
  ii_latlon = ind(min_lat.le.lat1d.and.lat1d.le.max_lat.and.\
                  min_lon.le.lon1d.and.lon1d.le.max_lon)
  nii = dimsizes(ii_latlon)
  print(nii + " values to check")
  print(numFeatures + " feature(s) to traverse with a maximum of " + \
        nlatlon + " points")

;---Create array to hold new data mask, and set all values to 0 initially.
  data_mask_1d = new(dimsizes(lat1d),integer)
  data_mask_1d = 0

;
; This is the loop that checks every point in lat1d/lon1d to see if it
; is inside or outside of the country. If it is inside, then data_mask_1d
; will be set to 1.
;
  ikeep = 0    ; Counter to see how many points were found inside the country
  do n=0,nii-1
    ii = ii_latlon(n)
    is_inside = False
    i = 0
    do while(.not.is_inside.and.i.lt.numFeatures)
       startSegment = geometry(i, geom_segIndex)
       numSegments  = geometry(i, geom_numSegs)
       do seg=startSegment, startSegment+numSegments-1
         startPT = segments(seg, segs_xyzIndex)
         endPT   = startPT + segments(seg, segs_numPnts) - 1
         if(data_mask_1d(ii).ne.1.and.\
            gc_inout(lat1d(ii),lon1d(ii),\
                     lat(startPT:endPT),lon(startPT:endPT))) then
           data_mask_1d(ii) = 1
           ikeep = ikeep+1
           is_inside = True
           continue
         end if
       end do
       i = i + 1
    end do
  end do
  print(ikeep + " values kept")
  print("==================================================")

;
; Create a 2D data array of same size as original data,
; but with appropriate values masked.
;
  data_mask = (where(onedtond(data_mask_1d,dims).eq.1,data,\
              data@_FillValue))
  copy_VarMeta(data,data_mask)      ; Copy all metadata

;---Print timings
  mask_end_time = get_cpu_time()
  print("Elapsed time in CPU second for 'mask_data_with_gadm_country' = " + \
         (mask_end_time-mask_start_time))

  return(data_mask)
end

;----------------------------------------------------------------------
; Main code
;----------------------------------------------------------------------
begin
;---Shapefile to use for masking. Using "Bangladesh" here (BGD).
 gadm_code = "BGD"
 gadm_name = gadm_code + "_adm"
 gadm_shp  = gadm_name + "1.shp"

;---Read lat/lon off shapefile
 s         = addfile(gadm_shp,"r")
 slat      = s->y
 slon      = s->x

;---Read NDVI data to contour and mask
  dir      = "./"
  filename = "ndvi3g_geo_v1_1985_0106.nc4"
  f        = addfile(dir+filename,"r")
  ndvi       = f->ndvi(0,:,:)
  printVarSummary(ndvi)

;---Start the graphics
  wtype          = "png"       ; "png"

 wks = gsn_open_wks(wtype,"mask_gadm_"+gadm_code)

  setvalues NhlGetWorkspaceObjectId()
    "wsMaximumSize" : 200000000
  end setvalues

  res                 = True                    ; plot mods desired

  res@cnFillOn        = True                    ; turn on color
  res@cnFillMode      = "RasterFill"
  res@cnLinesOn       = False                   ; turn off lines
  res@cnLineLabelsOn  = False                   ; turn off labels

  res@mpMinLatF       = min(ndvi&lat)
  res@mpMaxLatF       = max(ndvi&lat)
  res@mpMinLonF       = min(ndvi&lon)
  res@mpMaxLonF       = max(ndvi&lon)
  res@pmTickMarkDisplayMode = "Always"     ; nicer tickmarks

;---Create global plot of original data
  res@tiMainString    = filename
  plot = gsn_csm_contour_map(wks,ndvi, res)

;---Mask "ndvi" against GADM shapefile outlines
;---Create the mask based on the given shapefile
  opt         = True
  opt@minlat  = min(ndvi&lat)
  opt@maxlat  = max(ndvi&lat)
  opt@minlon  = min(ndvi&lon)
  opt@maxlon  = max(ndvi&lon)
  ndvi_mask  = shapefile_mask_data(ndvi,gadm_shp,opt)
  printVarSummary(ndvi_mask)
  copy_VarCoords(ndvi,ndvi_mask)

;---Set some additional resources for the second set of plots

  res@gsnDraw         = False                   ; don't draw yet
  res@gsnFrame        = False                   ; don't advance frame yet

;---Make sure both plots have the same contour levels.
  mnmxint = nice_mnmxintvl( min(ndvi_mask), max(ndvi_mask), 18, False)
  res@cnLevelSelectionMode = "ManualLevels"
  res@cnMinLevelValF  = mnmxint(0)
  res@cnMaxLevelValF  = mnmxint(1)
  res@cnLevelSpacingF = mnmxint(2)/2.   ; twice as many levels

  res@mpMinLatF       := min(slat)
  res@mpMaxLatF       := max(slat)
  res@mpMinLonF       := min(slon)
  res@mpMaxLonF       := max(slon)

  res@gsnRightString  = ""
  res@gsnLeftString   = ""

  res@lbLabelBarOn    = False

;---Create plot of original data
  res@tiMainString    = "original data zoomed in"
  plot_orig = gsn_csm_contour_map(wks,ndvi, res)

;---Create plot of masked data
  res@tiMainString = "with country mask"
  plot_mask = gsn_csm_contour_map(wks,ndvi_mask, res)

;---Add shapefile outlines to both plots
  lnres             = True
  lnres@gsLineColor = "NavyBlue"

  id_orig = gsn_add_shapefile_polylines(wks,plot_orig,gadm_shp,lnres)
  id_mask = gsn_add_shapefile_polylines(wks,plot_mask,gadm_shp,lnres)

;---Panel both plots on one page
  pres                  = True
  pres@txString         = ndvi@long_name + " (" + ndvi@units + ")"
 pres@gsnMaximize      = True
 pres@gsnPanelLabelBar = True
 gsn_panel(wks,(/plot_orig,plot_mask/),(/2,1/),pres)
end