;----------------------------------------------------------------------
; This function takes an average of a 2D curvilinear WRF grid over
; a given list of counties.
;
; A shapefile is used to get the lat/lon outlines for each county. 
;----------------------------------------------------------------------
undef("plot_averages_by_county")
function plot_averages_by_county(wks,plot,data,lat2d,lon2d,\
                                 fname[1]:string,levels,fill_colors)
local f, segments, geometry, segsDims, geomDims, geom_segIndex, \
geom_numSegs, segs_xyzIndex, segs_numPnts, numFeatures, i, lat, lon, \
startSegment, numSegments, seg, startPT, endPT, dims, 
minlat, maxlat, minlon, maxlon
begin
;
; Get the lat/lon area of interest so we can cull the lat/lon
; pairs that we have to check later.
;
  getvalues plot
    "mpMinLatF" : minlat
    "mpMaxLatF" : maxlat
    "mpMinLonF" : minlon
    "mpMaxLonF" : maxlon
  end getvalues
;
; Turn the 2D lat/lon arrays into 1D arrays. We will use these
; 1D arrays when checking which lat/lon pairs fall in which
; counties of the  of the given state.
;
  dims  = dimsizes(data)
  nlat  = dims(0)
  nlon  = dims(1)
  lat1d = ndtooned(lat2d)
  lon1d = ndtooned(lon2d)
  nlatlon = dimsizes(lat1d)

;---Debug prints
  printMinMax(lat1d,0)
  printMinMax(lon1d,0)

;
; 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(lat1d.ge.minlat.and.lat1d.le.maxlat.and.\
                  lon1d.ge.minlon.and.lon1d.le.maxlon)
  nii_latlon = dimsizes(ii_latlon)

;---Open the shapefile
  f = addfile(fname,"r")

;---Read data off the shapefile
  geomDims    = getfilevardimsizes(f,"geometry")
  numFeatures = geomDims(0)
  segments    = f->segments
  geometry    = f->geometry
  segsDims    = dimsizes(segments)

;---Read global attributes  
  geom_segIndex = f@geom_segIndex
  geom_numSegs  = f@geom_numSegs
  segs_xyzIndex = f@segs_xyzIndex
  segs_numPnts  = f@segs_numPnts

;---Read the lat/lon data off the shapefile
  lat = f->y
  lon = f->x

;---Read the state/county names off the file
  shape_states   = f->NAME_1
  shape_counties = f->NAME_2
  ncounties      = dimsizes(shape_counties)

;---Debug prints
  print("Number of counties         = " + ncounties)
  print("Number of lat/lon values   = " + nlatlon)
  print("Number of lat/lon values")
  print("   that will be checked    = " + nii_latlon)

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

;---Make sure "data" has a missing value.
  if(.not.isatt(data,"_FillValue")) then
    data@_FillValue = default_fillvalue(typeof(data))
  end if

;---This is for the averages computation later.
  data_1d = ndtooned(data)

;---Array to hold data averages for each county
  data_avg = new(ncounties,typeof(data),data@_FillValue)

;
; This is the loop that loops across every county of Georgia,
; and collects all the lat/lon points that are inside that
; county (data_mask_1d=1).
;
; It then takes an average of the data values for this county,
; and attaches it as a filled polygon to the given plot.
;
  gnres = True                       ; polygon resource list
  nfill = dimsizes(fill_colors)

  do i=0, ncounties-1
    if(any(shape_states(i).eq."Pennsylvania")) then
      print("--------------------------------------------------")
      print("Inspecting " + shape_states(i) + "/" + shape_counties(i) + "...")
    else
      continue
    end if

    data_mask_1d = 0    ; Be sure to reset to 0 for every county

;---Some counties have multiple segments
    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

;---Loop through each point on data grid and see if it's in this county.
      do n=0,nii_latlon-1
        nn = ii_latlon(n)     ; Get index of point we're checking
;
; This "if" statement speeds up code, by making sure we don't
; needlessly check a lat/lon point if:
; 
; - we've already found it in another county
; - it doesn't fall within the general lat/lon box that covers this county
; 
        if(skip_check(nn).or.\
           lat1d(nn).lt.min(lat(startPT:endPT)).or.\
           lat1d(nn).gt.max(lat(startPT:endPT)).or.\
           lon1d(nn).lt.min(lon(startPT:endPT)).or.\
           lon1d(nn).gt.max(lon(startPT:endPT))) 
           continue
        end if

;---Here's the check if the point is in the county.
        if(gc_inout(lat1d(nn),lon1d(nn),\
                    lat(startPT:endPT),lon(startPT:endPT))) then
          data_mask_1d(nn) = 1    ; This point is inside this county
          skip_check(nn)   = 1    ; Don't check this point again
        end if
      end do
    end do    ; End of collecting points for this county

;---Count number of points found in this county
    ndm = num(data_mask_1d.eq.1) 
    print(ndm + " data values found in this county.")
;
; If there are values found in this county, then calculate
; the average, set a fill color for this county, and do the
; fill.
;
    if(ndm.gt.0) then
      data_avg(i)  = avg(where(data_mask_1d.eq.1,data_1d,data_1d@_FillValue))
      print("Average over this county = " + data_avg(i))

;---Set the fill color for this county to the appropriate color
      iiavg := ind(data_avg(i).lt.levels)
      if(.not.any(ismissing(iiavg))) then
        gnres@gsFillColor = fill_colors(iiavg(0),:)    ; n x 4 array
      else
        gnres@gsFillColor = fill_colors(nfill-1,:)
      end if
      str = unique_string("poly")
      plot@$str$ = gsn_add_polygon(wks,plot,lon(startPT:endPT),\
                                            lat(startPT:endPT),gnres)
    end if
  end do

  return(data_avg)    ; Return data averages for each county
end