;*************************************************
; ce_master.ncl
; a total script that checks for dimensionality, 
; coordinate type, vectors etc. also adds
; the local time axis, the special hao labeling.
;************************************************
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/csm/shea_util.ncl" 
;load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl" 
load "$HAO_NCL_LIB/contributed.ncl.test" 
load "$HAO_NCL_LIB/hao_contrib.ncl"
load "$HAO_NCL_LIB/labeling.ncl"
load "$HAO_NCL_LIB/axis.ncl"
load "$HAO_NCL_LIB/halfzp2full.ncl"
load "$HAO_NCL_LIB/derive.ncl"
;************************************************
begin
;************************************************
; user parameters

; height overrides zplev
; can choose multiple fields
;************************************************
  ncfile   = "/chewy/d/murphys/files/astrid.nc"
  fields   = (/"ED2M3D","TN","JQR"/)    ; can be multiple
  mtime    = (/40,16,0/)    ; select time    
  zplev    = -2               ; level to plot
;  height   = 132             ; specify height (overrides zplev)
; vectors
  addvec = True              ; overlay vectors on plot
  refmag = 50.               ; magnitude of ref vector
  reflen = 0.03              ; length of ref vector
; manual contour levels
  manlev   =  False
  minlev   = -30000
  maxlev   =  30000
  int      =  5000
; convert units
;  density_units = "GM/CM3"  ; "MMR","CM3","CM3-MR","GM/CM3"
; output
  plt_type = "ps"           ; ps,eps,epsi,x11,or ncgm
  plt_name = "ce_master"    ; name of x11 window or output file
;color
; http://ngwww.ucar.edu/ncl/coltable.html
;  colormap = "BlWhRe"       ; comment out for black and white   
  colormap = "BlAqGrYeOrRevi200"
;************************************************
; no user changes after this point
;************************************************
  in = addfile(ncfile,"r")
; read in these variables for the hao labeling  
  z  = in->Z
  ut = in->ut

; if vectors requested, then process u/v only once rather than in loop
; below. this save time 
  if(addvec)then
    if(isfilevar(in,"UN"))then
      if(isfilevar(in,"VN"))then
        tmpu  = in->UN                          ; read in U
        tmpv  = in->VN                          ; read in V
        if(.not.isatt(in,"tgcmproc"))then
          tmpu   = halfzp("UN",tmpu)            ; interp to full zp
          tmpv   = halfzp("VN",tmpv)            ; interp to full zp
        end if

        if(isvar("height"))then                 ; want in km
          if(tmpu&lev@units.ne."km")then        ; not already interpolated
            if(.not.isatt(tmpu,"differences"))then ; can't interp diff files
              arg        = True
              arg        = True 
              arg@height = height
              u        = zplev_to_height(z,tmpu,arg)
              v        = zplev_to_height(z,tmpv,arg)
            else
              u = tmpu
              v = tmpv
            end if
          else
            u = tmpu
            v = tmpv
          end if
        else
          u = tmpu
          v = tmpv
        end if
      else
        print("VN not on file, can not add vectors")
        addvec = False
      end if
    else
      print("UN not on file, can not add vectors")
      addvec = False
    end if
  end if
  
; get time index from model time  
  t2plt = model_times2idx(mtime,in)            ; get time index from mtime
  

  wks = gsn_open_wks(plt_type,plt_name)        ; open a ps file
  print("creating "+plt_name+"."+plt_type)     ; so can cut & paste fm window

  res                            = True        ; plot mods desired
  res@cnLineLabelPlacementMode   = "constant"  ; constant contour labeling
  res@mpFillOn                   = False       ; don't fill in map
  res@gsnDraw                    = False       ; don't draw yet
  res@gsnFrame                   = False       ; don't advance frame yet

; color 
  if(isvar("colormap"))then
    gsn_define_colormap(wks,colormap)             ; choose colormap
    res@cnFillOn                 = True        ; turn on color
    res@gsnSpreadColors          = True        ; use full range of color map
    res@lbLabelAutoStride        = True        ; nice lb labels
  end if

; manual levels 
  if(manlev.eq.True)then
    res@cnLevelSelectionMode     ="ManualLevels"
    res@cnMinLevelValF           = minlev
    res@cnMaxLevelValF           = maxlev
    res@cnLevelSpacingF          = int
  end if

; vectors 
  if(addvec)then
    res@gsnScalarContour     = True            ; contours desired
    res@mpOutlineOn          = True            ; turn on map outlines
    res@cnFillDrawOrder      = "Draw"          ; draw on top of map
    res@vcVectorDrawOrder    = "PostDraw"      ; draw vectors last
    res@vcRefMagnitudeF      = refmag          ; define vector ref mag
    res@vcRefLengthF         = reflen          ; define length of vec ref
    res@vcGlyphStyle         = "CurlyVector"   ; turn on curley vectors
    res@vcMinDistanceF       = 0.02            ; thin vectors
    res@vcRefAnnoString2On   = True            ; turn on vec ref string
    res@vcRefAnnoString2     = u@units         ; string for ref vec
    res@vcRefAnnoArrowLineColor   = "black"    ; change ref vector color
    res@vcRefAnnoArrowUseVecColor = False      ; don't use vec color for ref
    res@vcLineArrowColor          = "white"    ; change vector color
    res@pmLabelBarOrthogonalPosF = -0.18       ; move label bar up
  end if

; main loop
  do i = 0,dimsizes(fields)-1
    if(isfilevar(in,fields(i)))then
      if(isfilevardim(in,fields(i),"mlat"))then  ; magnetic data
        if(addvec)then
          print("mag varbs can not be overlaid w/ vectors at this time")
          exit
        end if
        res@mpOutlineOn  = False                 ; off continental outline
        res@gsnAddCyclic = False                 ; mag data has cyc pt already
        x = mag2latlon(in->$fields(i)$)          ; read in data and convert
        zp = in->mlev
      else
        x   = in->$fields(i)$                    ; read in data 
        zp = in->lev
        if(.not.isatt(in,"tgcmproc"))then
          x   = halfzp(fields(i),x)                ; interp to full zp
        end if
      end if

; density conversion if necessary
      if(isvar("density_units"))then
        x@name = fields(i)
        x = den_convert_units(x,in,density_units)
      end if

; convert to height if necessary (do AFTER units conversion b/c array
; reduced by one) OR data is a tgcmproc file and may already be interpolated
; to height and we need to determine the index value of the requested 
; height to plot.
      if(dimsizes(dimsizes(x)).eq.4)then
        if(isvar("height"))then
          if(x&lev@units.eq."km")then              ; already in height
            loc        = closest_val(height,x&lev) ; a contributed function
            l2plt      = x&lev(loc)                ; coordinate subscripting
            arg        = True
            arg@height = l2plt                     ; for plot label
          else                                     ; need to interpolate
            if(.not.isatt(x,"differences"))then    ; can't interp diff files
              arg        = True
              arg@height = height                  ; for plot label
              tmp        = zplev_to_height(z,x,arg); interpolate to height
              delete(x)
              x = tmp
            else
              l2plt     = zplev
              arg       = True
              arg@zplab = zp({l2plt})
              arg@zavg  = avg(z(t2plt,{l2plt},:,:))
            end if
          end if
        else  
; plot is requested in zp coordinates
          if(zplev.gt.dimsizes(zp))then
            print("your chosen zplev of "+zplev+" is outside the range of the level array. If you are plotting both magnetic and non-magnetic variables, ensure your zplev is within the range of both lev and mlev")
            exit
          end if
          l2plt     = zplev
          arg       = True
          arg@zplab = zp({l2plt})
          if(l2plt.lt.-7 .and. isMag(x))then
            arg@zavg = "na"
          else
            arg@zavg  = avg(z(t2plt,{l2plt},:,:))
          end if
        end if
      end if
; tgcmproc files already have a cyclic point, check and turn off
      if(isatt(in,"tgcmproc"))then
        res@gsnAddCyclic = False
      end if

; assign _FillValue
      if(isatt(in,"missing_value"))then
        x@_FillValue=in@missing_value
        if(addvec)then
          v@_FillValue=in@missing_value
          u@_FillValue=in@missing_value
        end if
      end if

; turn on local time axis
      res = set_lt_axis_ceplot(ut(t2plt),isMag(x),res)
;      res@mpCenterLonF = 15*(12-ut(t2plt))+71 ; put mlt noon at center 

; get title depending on availability of info
      res   = set_title(x,fields(i),res)
           
; prepare top text arguments
      arg      = True
      arg@time = t2plt
      arg@mag  = isMag(x)
      arg@file = in
      if(isatt(x,"differences"))then
        arg@diff = True
      end if

; create plot
      if(dimsizes(dimsizes(x)).eq.3)then ; zp independent or height
        res    = textTop_ce(wks,arg,res)
        if(addvec)then
          plot = gsn_csm_vector_scalar_map_ce(wks,u(t2plt,:,:),v(t2plt,:,:),\
          x(t2plt,:,:),res)
        else
          plot = gsn_csm_contour_map_ce(wks,x(t2plt,:,:), res)
        end if
        plot   = ZeroNegDashLineContour(plot)
        textBott(wks,plot,x(t2plt,:,:),t2plt,ncfile,in,res)
        draw(plot)
        frame(wks)
      else 
        res         = textTop_ce(wks,arg,res)
        if(addvec)then
          plot      = gsn_csm_vector_scalar_map_ce(wks,u(t2plt,{l2plt},:,:),\
          v(t2plt,{l2plt},:,:),x(t2plt,{l2plt},:,:),res)
        else
          plot      = gsn_csm_contour_map_ce(wks,x(t2plt,{l2plt},:,:), res)
        end if
        plot        = ZeroNegDashLineContour(plot)
        textBott(wks,plot,x(t2plt,{l2plt},:,:),t2plt,ncfile,in,res)
        draw(plot)
        frame(wks)
      end if
      delete(x)
      delete(zp)
      delete(arg)
    else
      print("Variable "+fields(i)+" does not exist on file, skipping")
    end if
  end do
end