load "./shapefile_utils.ncl"
load "./s.ncl"

;----------------------------------------------------------------------
; Function to generate some dummy data over part of the United States.
; This area is being used because we want to be sure to have data
; over the state of Georgia.
;----------------------------------------------------------------------
function dummy_data(nlat,nlon,minlat,maxlat,minlon,maxlon)
local nlat, nlon, lat, lon
begin
;---Generate some dummy lat/lon data over area that covers Georgia
  lat       = fspan(minlat,maxlat,nlat)
  lon       = fspan(minlon,maxlon,nlon)
  lat@units = "degrees_north"
  lon@units = "degrees_east"

  data     = generate_2d_array(10, 25, -20, 15, 0, (/nlat,nlon/))
  data!0   = "lat"
  data!1   = "lon"
  data&lat = lat
  data&lon = lon
  
  return(data)
end

begin
  shp_west_name = "8x_west.shp"
  shp_east_name = "8x_east.shp"

  fwest = addfile(shp_west_name,"r")
  feast = addfile(shp_east_name,"r")
  lat_west = fwest->y
  lon_west = fwest->x
  lat_east = feast->y
  lon_east = feast->x

;---Get lat/lon box that covers whole domain of both shapefiles
  minlat = min((/min(lat_west),min(lat_east)/))
  maxlat = max((/max(lat_west),max(lat_east)/))
  minlon = min((/min(lon_west),min(lon_east)/))
  maxlon = max((/max(lon_west),max(lon_east)/))

;---Generate some dummy data with lat/lon coordinates.
  data = dummy_data(100,100,minlat,maxlat,minlon,maxlon)

;---Mask data against outlines in a shapefile
  opt = True
  opt@return_mask = True
  data_mask_west = shapefile_mask_data(data,shp_west_name,opt)
  data_mask_east = shapefile_mask_data(data,shp_east_name,opt)
  data_mask = where(data_mask_east.eq.1.or.data_mask_west.eq.1,data,data@_FillValue)

  copy_VarMeta(data,data_mask)

  print("----------------------------------------------------------------------")
  print("Original data")
  print("   dimensions          : " + str_join(""+dimsizes(data)," x "))
  print("   # of valid values   : " + num(.not.ismissing(data)))
  print("   # of missing values : " + num(ismissing(data)))
  print("   min/max             : " + min(data) + "/" + max(data))
  print("   average             : " + avg(data))
  print("----------------------------------------------------------------------")
  print("Masked data")
  print("   dimensions          : " + str_join(""+dimsizes(data_mask)," x "))
  print("   # of valid values   : " + num(.not.ismissing(data_mask)))
  print("   # of missing values : " + num(ismissing(data_mask)))
  print("   min/max             : " + min(data_mask) + "/" + max(data_mask))
  print("   average             : " + avg(data_mask))
  print("----------------------------------------------------------------------")

  wks = gsn_open_wks("png","data_masked_by_shapefiles")

  res                       = True

  res@gsnMaximize           = True
  res@gsnDraw               = False
  res@gsnFrame              = False
  res@gsnAddCyclic          = False

  res@mpFillOn              = False
  res@mpOutlineBoundarySets = "USStates"
  res@mpMinLatF             = minlat
  res@mpMaxLatF             = maxlat
  res@mpMinLonF             = minlon
  res@mpMaxLonF             = maxlon

  res@pmTickMarkDisplayMode = "Always"

;---Be sure to use same contour levels across both plots
  mnmxint = nice_mnmxintvl( min(data), max(data), 18, False)
  res@cnLevelSelectionMode  = "ManualLevels"
  res@cnMinLevelValF        = mnmxint(0)
  res@cnMaxLevelValF        = mnmxint(1)
  res@cnLevelSpacingF       = mnmxint(2)

  res@cnFillOn              = True
  res@cnLinesOn             = False
  res@cnLineLabelsOn        = False
  res@cnLineLabelsOn        = False
  res@lbLabelBarOn          = False
  res@gsnAddCyclic          = False

  res@tiMainFont            = "Helvetica"
  res@tiMainFontHeightF     = 0.02

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

;---Create plot of data that was masked by shapefile
  res@tiMainString = "Data masked by two shapefiles"
  plot_mask = gsn_csm_contour_map(wks,data_mask,res)

;---If desired, add shapefile outlines
  lnres = True
  id_orig_west = gsn_add_shapefile_polylines(wks,plot_orig,shp_west_name,lnres)
  id_orig_east = gsn_add_shapefile_polylines(wks,plot_orig,shp_east_name,lnres)
  id_mask_west = gsn_add_shapefile_polylines(wks,plot_mask,shp_west_name,lnres)
  id_mask_east = gsn_add_shapefile_polylines(wks,plot_mask,shp_east_name,lnres)

;---Draw both plots in a panel for comparison.
  pres                         = True
  pres@gsnMaximize             = True
  pres@gsnPanelMainString      = "Data masked by " + shp_west_name + " and " + shp_east_name
  pres@gsnPanelMainFont        = "helvetica-bold"
  pres@gsnPanelMainFontHeightF = 0.025
  pres@gsnPanelLabelBar        = True
  pres@pmLabelBarWidthF        = 0.8
  pres@lbLabelFontHeightF      = 0.015
  gsn_panel(wks,(/plot_orig,plot_mask/),(/1,2/),pres)

;---If desired, draw plot with markers at the lat/lon locations showing what locations were kept
  mkres = True
  mkres@gsMarkerIndex            = 16         ; filled tos
  mkres@gsMarkerSizeF            = 2.0        ; default is a bit large
  mkres@gsnCoordsNonMissingColor = "black"
  mkres@gsnCoordsMissingColor    = "brown"    ; use "transparent" if you don't want any markers here
  gsn_coordinates(wks,plot_mask,data_mask,mkres)
end