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_mask_data.ncl"
load "./shapefile_average_data.ncl"

;----------------------------------------------------------------------
; Adds dots to the given plot at data's lat/lon grid locations.
;----------------------------------------------------------------------
undef("add_latlon_points")
procedure add_latlon_points(wks,plot,data)
local mkres
begin
  mkres                 = True
  mkres@gsMarkerIndex   = 16         ; Filled dots
  mkres@gsMarkerSizeF   = 0.001      ; Make them smaller
  mkres@gsMarkerColor   = "black"
  mkres@gsnCoordsAttach = True
  gsn_coordinates(wks,plot,data,mkres)
end

;----------------------------------------------------------------------
; Add shapefile outlines to the given plot.
;----------------------------------------------------------------------
undef("add_shapefile_outlines")
procedure add_shapefile_outlines(wks,plot,shp_filename)
local lnres
begin
;---Resources for polyline
  lnres                  = True
  lnres@gsLineColor      = "Navyblue"
; lnres@gsLineThicknessF = 3.0            ; 3x thickness

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

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

;---Get temperature (degC) and lat/lon from WRF file
  nt       = 0    ; time index
  nl       = 0    ; level index
  tc       = wrf_user_getvar(a,"tc",nt)
  tc@lat2d = wrf_user_getvar(a,"lat",nt)
  tc@lon2d = wrf_user_getvar(a,"lon",nt)

;
; Use shapefile to mask the temperature array based on shapefile outline.
; This will make the data averaging go a little faster because the code
; won't have to check as many lat/lon points.
;
  shp_filename = "GIS.OFFICIAL_CLIM_DIVISIONS.shp"
  tc_mask      = shapefile_mask_data(tc(nl,:,:),shp_filename,True)

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

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

  res@mpMinLatF        = min(tc@lat2d)-1
  res@mpMaxLatF        = max(tc@lat2d)+1
  res@mpMinLonF        = min(tc@lon2d)-1
  res@mpMaxLonF        = max(tc@lon2d)+1

  cnres                = res    ; Copy map resources before we set left/right string titles

;--Create just a map for adding filled polygons later.
  res@gsnLeftString    = tc@description + " (averaged over climate divisions)"
  res@gsnRightString   = tc@units
  plot_map = gsn_csm_map(wks,res)    

;---Create filled contour plot of original data.
  cnres@gsnAddCyclic          = False
  cnres@cnFillOn              = True
  cnres@cnLinesOn             = False
  cnres@cnLineLabelsOn        = False
  cnres@cnLevelSpacingF       = 2.
  cnres@lbLabelBarOn          = False

  plot_data = gsn_csm_contour_map(wks,tc(0,:,:),cnres)
  
;---Get the colors and levels used for contour plot
  getvalues plot_data@contour
    "cnLevels"     : levels
    "cnFillColors" : colors
  end getvalues

;---Add color-coded polygons that represent average of data in a particular climate division
  plot_averages_using_shapefile_features(wks,plot_map,tc_mask,shp_filename,levels,colors)

;---Add shapefile outlines to both plots
  add_shapefile_outlines(wks,plot_data,shp_filename)
  add_shapefile_outlines(wks,plot_map,shp_filename)

;
; Panel both plots. All the stuff we've added to both plots,
; shapefile outlines, lat/lon dots, filled polygons, will
; be seen here.
;
  pres                  = True
  pres@gsnMaximize      = True
  pres@gsnPanelLabelBar = True
  gsn_panel(wks,(/plot_data,plot_map/),(/2,1/),pres)

;---Add dots at lat/lon locations and draw panel plots again
  add_latlon_points(wks,plot_data,tc(nl,:,:))
  add_latlon_points(wks,plot_map,tc(nl,:,:))
  gsn_panel(wks,(/plot_data,plot_map/),(/2,1/),pres)

end