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"

;----------------------------------------------------------------------
; Procedure to attach colored markers to an NCL map, given 1D
; data, and lat/lon arrays, a set of levels, and a color map.
;----------------------------------------------------------------------
procedure add_obs_markers_to_map(wks,plot,levels,colormap,data_1d,lat_1d,lon_1d)
local mkres, nlevels, colors, nlevels, n, ii
begin

;---Set resources for the markers
  mkres = True
  mkres@gsMarkerIndex = NhlNewMarker(wks, "y", 35, 0.0, 0.0, 1., 0.5, 0.)     ; filled square
;
; Based on the levels we had for the contour plot and the colormap used, 
; get an array of colors that span the colormap. This uses the same
; algorithm that the contour used, so the colors will be identical.
;
  nlevels = dimsizes(levels)
  colors = span_color_rgba(colormap,nlevels+1)   ; Need one more color than number of levels

;
; Loop through each level, gather all the data values in the given level
; range, and add colored markers to the existing map.
;
  do n=0,nlevels

;---These "if" statements are how color contours are handled in NCL.
    if(n.eq.0) then
      ii := ind(data_1d.lt.levels(n))
    else if(n.eq.nlevels) then
      ii := ind(data_1d.ge.levels(n-1))
    else
      ii := ind(data_1d.ge.levels(n-1).and.data_1d.lt.levels(n))
    end if
    end if
    if(any(ismissing(ii))) then
      continue
    end if

;---Add the markers
    mkres@gsMarkerColor = colors(n,:)    ; colors is an N x 4 array
    plot@$unique_string("markers")$ = gsn_add_polymarker(wks,plot,lon_1d(ii),lat_1d(ii),mkres)
  end do
end

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

begin

;---Open WRF output file and read data
  dir      = "./"
  filename = "wrfout_d01_2003-07-15_00:00:00.nc"
  a = addfile(dir + filename,"r")

  it    = 0                              ; first time step
  tc    = wrf_user_getvar(a,"tc",it)     ; temperature (degC)
  lat2d = wrf_user_getvar(a,"XLAT",it)   ; latitude/longitude
  lon2d = wrf_user_getvar(a,"XLONG",it)

  tc@lat2d = lat2d    ; for plotting
  tc@lon2d = lon2d

;---Will use this for contours and filled markers
  colormap = "BlAqGrYeOrReVi200"
  levels   = ispan(10,38,2)

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

;---Common resources shared by both contour plot and marker plot
  res                        = True
  res@gsnMaximize            = True
  res@gsnLeftString          = ""
  res@gsnRightString         = ""

  res@mpDataBaseVersion      = "MediumRes"
  res@mpFillOn               = False
  res@mpMinLatF              = min(tc@lat2d)
  res@mpMaxLatF              = max(tc@lat2d)
  res@mpMinLonF              = min(tc@lon2d)
  res@mpMaxLonF              = max(tc@lon2d)
  res@pmTickMarkDisplayMode  = "Always"      ; better map tickmarks

;---Resources for filled contour plot
  cnres                      = res
  cnres@cnFillOn             = True
  cnres@cnLevelSelectionMode = "ExplicitLevels"
  cnres@cnLevels             = levels
  cnres@cnFillPalette        = colormap
  cnres@cnLinesOn            = False
  cnres@cnLineLabelsOn       = False
  cnres@lbLabelBarOn         = False   ; will add in panel later
  cnres@pmTitleZone          = 2       ; move title down
  cnres@gsnAddCyclic         = False

;---Create smooth contour plot
  nl = 0
  cnres@tiMainString = "WRF model data"
  plot_wrf = gsn_csm_contour_map(wks,tc(nl,:,:),cnres)

;----------------------------------------------------------------------
; Recreate similar plot, but by using filled markers drawn over a map.
;----------------------------------------------------------------------

;
; NORMALLY YOU WANT TO USE A REAL OBSERVATIONAL DATASET HERE.
;
; We will instead use the same data, but convert it to 1D so it 'mimics"
; an observational dataset. We will randomly grab 1/8 of the points
; just for something more "observation data" like.
;
; Using the same dataset is a good debugging tool to make sure our marker
; code is working properly.
;

;---Randomly grab about an eighth of the points
  nlatlon  = product(dimsizes(lat2d))
  npts     = toint(nlatlon/8)         
  indexes  = get_unique_values(toint(random_uniform(0,nlatlon-1,npts)))

  tc1d  = ndtooned(tc(nl,:,:))
  lat1d = ndtooned(lat2d)
  lon1d = ndtooned(lon2d)

  tc1d_subset  = tc1d(indexes)
  lat1d_subset = lat1d(indexes)
  lon1d_subset = lon1d(indexes)

  delete([/tc1d,lat1d,lon1d/])    ; No longer needed

;---Create a map plot for adding markers to.
  res@gsnDraw         = False
  res@gsnFrame        = False
  res@tfPolyDrawOrder = "Draw"    ; Necessary to make sure filled markers drawn *under* map outlines
  res@tiMainString    = "'Observational' data"

  plot_obs = gsn_csm_map(wks,res)

;---Add colored markers to the map
  add_obs_markers_to_map(wks,plot_obs,levels,colormap,tc1d_subset,lat1d_subset,lon1d_subset)

;---This draws the map with the filled markers
  draw(plot_obs)
  frame(wks)

;---Panel both plots for comparison.
  pres                    = True
  pres@gsnMaximize        = True
  pres@gsnPanelLabelBar   = True
  pres@pmLabelBarWidthF   = 0.8
  pres@lbLabelFontHeightF = 0.01 
  gsn_panel(wks,(/plot_wrf,plot_obs/),(/1,2/),pres)
end