undef("get_color_indexes")
;;;   extends the get_color_index function to retrieve multiple indices for multiple values
;;;   same code just extended for an array. 
function get_color_indexes(cmapt,cnlvls[*]:numeric,value[*]:numeric)
local cmap, dims, ncn, nclr, color, n, col_indexes, ncoli
begin

 if(isstring(cmapt)) then
    cmap = read_colormap_file(cmapt)
 else if(isnumeric(cmapt)) then
   dims = dimsizes(cmapt)
   if(dimsizes(dims).ne.2.or.dims(0).lt.3.or.dims(0).gt.256.or.\
       .not.any(dims(1).ne.(/3,4/))) then
     print ("Error: get_color_index: cmap must be an n x 3 or n x 4 array of RGB or RGBA values, or a valid color map name")
     return(new(3,"float"))    ; return a missing value
   end if
   cmap = cmapt
 else
   print ("Error: get_color_index: cmap must be an n x 3 or n x 4 array of RGB or RGBA values, or a valid color map name")
 end if
 end if

 ncn  = dimsizes (cnlvls)
 nclr = dimsizes (cmap(:,0))
 imsg = new(1,integer)          ; missing value

 if (nclr-2 .lt. ncn+1) then
   print ("Warning: get_color_index: Not enough colors in colormap for number of contour levels")
   print ("         Colors will be repeated")
 end if
 if (all(ismissing(value))) then
   print ("Error: get_color_index: Input value is missing")
   return (imsg)
 end if
 if (any(ismissing(cnlvls))) then
   print ("Error: get_color_index: One or more input contour levels are missing")
   return (imsg)
 end if


;---Get nice span of indexes throughout the color map
 col_indexes = span_color_indexes(cmap,dimsizes(cnlvls)+1)
 ncoli       = dimsizes(col_indexes)    ; should be ncn+1

n = ind_nearest_coord(value, cnlvls, 0)   ;;;; brammer -- retrieve levels for values
n = where(ismissing(n), 0, n)             ;;;; brammer -- use index 0 if not found


 if(isstring(cmapt)) then
   return(col_indexes(n)+2)    ; Account for 0/1 index being dropped
  else
   return(col_indexes(n))
 end if
end



undef("add_colored_line")
;``` Draws segmented colored polyline based on values provided 
;    Adds line to plot using attribute and incrementing the attribute count ```
procedure add_colored_line( wks, plot, inX, inY, inZ, levels, colors,gsres)
local wks, plot, inX, inY, inZ, levels, colors,gsres \
      nmi, x, y, z,newx, newy, newz, lres, cmap, icolors,cts
begin
nmi = ind( (.not. ismissing(inX)) .and. (.not. ismissing(inY)) .and. (.not. ismissing(inZ)) )
x =  inX(nmi)
y = inY(nmi)
z = inZ(nmi)


;;; Make coloured line segments shorter so values are represented more accurately. 
;;; linear interpolation is ugly though, could be replaced with a cubic interp routine 
if(dimsizes(x) .lt. 999)
x := linint1( fspan(0, 100, dimsizes(x)),  x, False, fspan(0, 100, 1001), 0) 
y := linint1( fspan(0, 100, dimsizes(y)),  y, False, fspan(0, 100, 1001), 0) 
z := linint1( fspan(0, 100, dimsizes(z)),  z, False, fspan(0, 100, 1001), 0) 
end if

;; segments need a value for the start and end of each segment.  
;; therefore create an array where every value it repeated. 
newx = new(dimsizes(x)*2, typeof(x))
newx(::2) = x
newx(1::2) = x
newy = new(dimsizes(y)*2, typeof(y))
newy(::2) = y
newy(1::2) = y
newz = new(dimsizes(z)*2, typeof(z))
newz(::2) = z
newz(1::2) = z


;; Set some default values for the line res. 
;; probably a more complete manner of getting line res' from gsres
lres = True
lres@gsLineThicknessF = get_res_value(gsres,"gsLineThicknessF", 2.)
lres@gsSegments = ispan(1, dimsizes(newy)-2,2)
lres@gsEdgesOn   = get_res_value(gsres, "gsEdgesOn", True )   ; Turn on edges
lres@gsEdgeColor = get_res_value(gsres, "gsEdgeColor", "white")
lres@gsFillIndex = get_res_value(gsres, "gsFillIndex", 0)      ; Solid fill, the default


;;;; retrieve colortable and get colors for each segment
if(isnumeric(colors))
    cmap = colors
else
    cmap = read_colormap_file(colors)
end if
icolors = get_color_indexes( cmap, levels, newz)
lres@gsColors = cmap(icolors,:)

;;; retrieve existing lines on plot and increment count
;;; using unique_string becomes unbearably slow after multiple iterations
;;; this may slow over time but not noticeably. 
cts = str_sub_str(str_match(getvaratts(plot),"colored_line"),"colored_line","" )
if(all(ismissing(cts)))
    ct = 0
else
    ct = max(toint(cts))+1
end if 

plot@$tostring("colored_line"+ct)$ = gsn_add_polyline(wks, plot, newx,  newy, lres)

end




undef("colored_line_demo")
procedure colored_line_dumo()
begin

wks = gsn_open_wks("x11","")

x := fspan(0, 10, 21)
y := sin(x)
;z := cos(x) ;; we won't actually use this but you could switch out 
z := y

res = True
res@trXMinF := min(x)
res@trXMaxF := max(x)
res@trYMinF := 1.2*min(y)
res@trYMaxF := 1.2*max(y)
res@gsnDraw = False
res@gsnFrame = False

levels = fspan(min(z), max(z), 11)
colors := read_colormap_file("matlab_jet")

line_res := True
line_res@gsLineThicknessF = 4.

plot := gsn_csm_blank_plot(wks, res)
add_colored_line(wks, plot, x, y, z, levels, colors, line_res)
draw(plot)
frame(wks)
end