;-----------------------------------------------------------------
; Plot Divergence, PSI, and WINDS from WRF-ARW
;
;
; Version 1: Kacie N Shourd, 08/27/2017
;-----------------------------------------------------------------

load "$NCARG_ROOT/lib/ncarg/nclscripts/esmf/ESMF_regridding.ncl"
load "~/scratch/code/nclExtras/contributed640.ncl"

;-----------------------------------------------------------------
;--- Open/read WRF-ARW NetCDF output file ------------------------

  mo  = "06"
  vld = "16"
  day = "29"

  dir  = "/ScenarioA/Part1/kshourd/wrf29Jun/"
  fili = "wrfout_d03_2012-"+mo+"-"+day+"_"+vld+":00:00.nc"
  filin= dir+fili
  a    = addfile(filin,"r")

  setvalues NhlGetWorkspaceObjectId()
    "wsMaximumSize" : 300000000
  end setvalues    

;-----------------------------------------------------------------

  times  = wrf_user_getvar(a,"times",-1) ; get times in the file
  ntimes = dimsizes(times)         ; number of times in the file

  FirstTime = True

;-----------------------------------------------------------------
;--- Get variables -----------------------------------------------

    uvm = wrf_user_getvar(a,"uvmet",-1)
    u   = uvm(0,:,:,:,:)
    v   = uvm(1,:,:,:,:)

    p   = wrf_user_getvar(a,"p",-1)

    z   = wrf_user_getvar(a,"z",-1)

    T   = wrf_user_getvar(a,"tk",-1)

    lat2d = wrf_user_getvar(a,"lat",0)
    y = lat2d
    printVarSummary(lat2d)

    lon2d = wrf_user_getvar(a,"lon",0)
    x = lon2d
    printVarSummary(lon2d)

print("Calculating PSI...")
;-----------------------------------------------------------------
;--- Calculate Montgomery Stream Fcn -----------------------------

    cp = 1004.
    cp@units = "J kg-1 K-1"
    g  = 9.81
    g@units  = "m s-2"

    psi  = (cp*T)+(g*z)
    copy_VarCoords(T, psi)
    psi@units = "m2 s-2"
    psi@description = "Montgomery Stream Function (Psi)"
    printVarSummary(psi)

print("Begin regridding")
;-----------------------------------------------------------------
;--- FOR REGRIDDING w ESMF REGRID --------------------------------
;--- Available in NCL v6.1.0 and later. --------------------------

    InterpMethod= "patch"

;-----------------------------------------------------------------
;--- Retrieve either one level, or all levels. Use '-1' for all --

    sfile = a

    Opt                = True
    Opt@SrcTitle       = "WRF grid"           ;;--- optional
    Opt@WgtFileName    = "WRF_to_Rect.WgtFile_"+InterpMethod+".nc"

;-----------------------------------------------------------------
;--- Generate the names of the files containing the --------------
;--- source and destination grid descriptions --------------------

    Opt@SrcFileName    = "WRF.SCRIP_grid_description.nc"         ; Name of source and
    Opt@DstFileName    = "Rectilinear.SCRIP_grid_description.nc" ; destination files
    Opt@SrcRegional    = True

;-----------------------------------------------------------------
;--- Get the source lat/lon grid (mass grid) ---------------------

    u@lat2d  = lat2d                          ;;--- This information will be used by
    u@lon2d  = lon2d                          ;;--- ESMF_regrid for the source grid
    v@lat2d  = lat2d
    v@lon2d  = lon2d
    z@lat2d  = lat2d
    z@lon2d  = lon2d

    dims  = dimsizes(lat2d)
    nlat  = dims(0)
    nlon  = dims(1)
    plev  = dimsizes(p)
    nlev  = plev(0)

;------------------------------------------------------------------
;--- Create the destination rectilinear lat[*]/lon[*] arrays ------

    lat = fspan(min(lat2d), max(lat2d) ,nlat)
    lon = fspan(min(lon2d), max(lon2d) ,nlon)

    Opt@DstGridType    = "rectilinear"
    Opt@DstGridLat     = lat
    Opt@DstGridLon     = lon
    Opt@DstRegional    = True

;---Specify other options  
    Opt@ForceOverwrite = True
    Opt@InterpMethod   = InterpMethod

;-------------------------------------------------------------------
;---Perform the regrid: WRF ==> rectilinear (_reclin) --------------
    U_regrid  = ESMF_regrid(u, Opt)     ;;--- Do the regridding for ua
    V_regrid  = ESMF_regrid(v, Opt)

print("Calculating divergence...")
;--------------------------------------------------------------------
;--- Calculate Divergence Tendency ----------------------------------

    dv  = uv2dv_cfd (U_regrid,V_regrid,lat,lon, 2)

print("Regridding back to WRF...")
;======================================================================
; Part B: Generate Rectilinear to WRF regrid weights 
;======================================================================

;---For clarity, delete the above options and start again.
    delete(Opt)

;---Options for regridding rectilinear to WRF (curvilinear) grid
    Opt                 = True

;---The grid descriptions have been generated in 'Part A'
    Opt@SkipSrcGrid     = True
    Opt@SkipDstGrid     = True
    Opt@DstFileName     = "WRF.SCRIP_grid_description.nc"         ; Name of source and
    Opt@SrcFileName     = "Rectilinear.SCRIP_grid_description.nc" ; destination files
    Opt@ForceOverwrite  = True

    Opt@SrcTitle        = a                  ; source grid
    Opt@SrcMask2D       = where(ismissing(U_regrid),0,1)
    Opt@SrcRegional     = True

    Opt@DstTitle        = "Rectilinear_to_WRF"
    Opt@DstGridLat      = lat2d
    Opt@DstGridLon      = lon2d
    Opt@DstRegional     = True

    Opt@InterpMethod    = InterpMethod
    Opt@WgtFileName     = "Rect_to_WRF.WgtFile_"+InterpMethod+".nc"

    u_wrf               = ESMF_regrid(U_regrid,Opt)
    v_wrf               = ESMF_regrid(V_regrid,Opt)
    dv_wrf              = ESMF_regrid(dv, Opt)

;-----------------------------------------------------------------
;--- Turn on resources -------------------------------------------

    res = True
    mpres=True
    pltres=True

;-----------------------------------------------------------------
;--- Start plotting ----------------------------------------------

    print("Open WKS...")
    type = "png"
    wks = gsn_open_wks(type,"dv-ag-psi_"+times(0))

    print("Interpolating to isentropic levels...")
;-----------------------------------------------------------------
;--- Interpolate to isentropic levels ----------------------------

    thlev  = 320
    opts   = True

    u320   = wrf_user_vert_interp(a, u_wrf, "theta", thlev, opts)
    v320   = wrf_user_vert_interp(a, v_wrf, "theta", thlev, opts)
    psii   = wrf_user_vert_interp(a,   psi, "theta", thlev, opts)
    div    = wrf_user_vert_interp(a,    dv, "theta", thlev, opts)

    print("Setting resources...")
;-----------------------------------------------------------------
;--- PSI Resources -----------------------------------------------

    pres = True
    pres@cnFillOn = False
    pres@cnLineThicknessF = 3.0
    pres@cnLineColor = "black"
    pres@sfXArray = lon2d
    pres@sfYArray = lat2d
    pres@gsnContourNegLineDashPattern = True
    pres@gsnContourPosLineDashPattern = True


    p_contour  = wrf_contour(a,wks,psii(0,{320},:,:),pres)

;-----------------------------------------------------------------
;--- Divergence Resources ----------------------------------------

    dres = True
    dres@cnFillOn = False
    dres@cnLineThicknessF = 1.5
    dres@cnLineColor = "red"
    dres@sfXArray = lon2d
    dres@sfYArray = lat2d
    dres@gsnContourNegLineDashPattern = False
    dres@gsnContourPosLineDashPattern = False
    ;;dres@ContourParameters = (/0,,1/)
printVarSummary(div)

    d_contour  = wrf_contour(a, wks, div(0,{320},:,:), dres)

;-----------------------------------------------------------------
;--- Actual Wind Resources ---------------------------------------

    wopts = True          
    wopts@FieldTitle = "Wind"   ;;--- overwrite Field Title
    wopts@NumVectors = 20                    ;;--- wind barb density
    wopts@vfXArray = lon2d
    wopts@vfYArray = lat2d
    wopts@NoHeaderFooter = True ; don't want to print out info 
    wopts@vcWindBarbLineThicknessF= 2.5             ; set the wind barb thickness
    wopts@vcMonoWindBarbColor = False
    wopts@vcGlyphStyle = "WindBarb"
    wopts@vcLevelPalette = "grads_rainbow"
printVarSummary(u320)

    vector1    = wrf_vector(a,wks,u320(0,{320},:,:),v320(0,{320},:,:),wopts)

    plot = wrf_map_overlays(a, wks, (/ p_contour,d_contour,vector1 /), pltres, mpres)