;Plot tracer concentrations at different levels

; check:
; winds are the correct ones to use
; PBL average calculation
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/wrf/WRFUserARW.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"
load "wrf_contour_ggp.ncl"
load "wrf_map_resources_ggp.ncl"
load "wrf_map_ggp.ncl"
load "wrf_map_overlays_ggp.ncl"

undef("get_wrfoutfname")
function get_wrfoutfname(domain:string, yr:integer, month:integer, day:integer, hr:integer, minutes:integer)
 begin
  if(minutes .lt. 10) then
    minute = "0"+minutes
  else
    minute = minutes
  end if
  if(hr .lt. 10) then    hour = "0"+hr
  else
    hour = hr
  end if
  do while(hr .gt. 23)
    day = day + 1
    if(day .lt. 10) then
      cday = "0"+day
    else
      cday = day
    end if
    hr = hr-24
    print("   minutes "+minutes+"   hr "+hr+"   day "+day)
    if(hr .lt. 10) then
      delete(hour)
      hour = "0"+hr
    else
      hour = hr
    end if
  end do
  if(yr .lt. 10) then
    year = "000"+yr
  else
    if(yr .lt. 100) then
      year = "00"+yr
    else
      if(yr .lt. 1000) then
        year = "0"+yr
      else
        year = yr
      end if
    end if
  end if
  if(month .lt. 10) then
    mon = "0"+month
  else
    mon = month
  end if
  if(day .lt. 10) then
    cday = "0"+day
  else
    cday = day
  end if
  filename = "wrfout_"+domain+"_"+year+"-"+mon+"-"+cday+"_"+hour+":"+minute+":00"
  print(" reading from file "+filename)
  delete(year)
  delete(mon)
  delete(cday)
  delete(hour)
  delete(minute)

  return(filename)
 end
;-------------------------------------------------------------------------------------------

;****************************************************************************************************;
begin
;year = 2011
;month = 6
;day = 10
;hr = 1
minutes = 0
  if(minutes .lt. 10) then
    cmin = "0"+minutes
  else
    cmin = minutes
  end if
  if(hr .lt. 10) then
    chr = "0"+hr
  else
    chr = hr
  end if
  if(day .lt. 10) then
    cday = "0"+day
  else
    cday = day
  end if
  if(month .lt. 10) then
    mon = "0"+month
  else
    mon = month
  end if
cdate = year+"_"+mon+"-"+cday+"_"+chr+cmin

print(cdate)
dirname = "./wrfout_linked/"
print(dirname)

time = 0    ; assume one time entry per file

;set levels for plotting
levs_height_tit = (/"Sfc","PBLavg","3km","5km","8km"/)
;levs_height_tit = (/"Sfc"/)
levs_height = (/-1,0,3000.,5000.,8000./)
;levs_height = (/-1/)
nlevsplot = 5


fname1 = get_wrfoutfname("d02", year, month, day, hr, minutes) 
print(fname1)
a=addfile (dirname + fname1, "r")

lat2U = a->XLAT_U(0,:,:)
lon2U = a->XLONG_U(0,:,:)
z    = wrf_user_getvar(a,"z",time)    ; meters
ht   = wrf_user_getvar(a,"HGT",time)    ; meters
prs = wrf_user_getvar(a,"pres",time)   ;Pa
psfc = wrf_user_getvar(a,"PSFC",time)    ;Pa
Tk = wrf_user_getvar(a,"tk",time)   ;Pa
pblh= wrf_user_getvar(a,"PBLH",time) ; meters


dimz = dimsizes(z)
dims = dimsizes(lat2U)
   

windrot = wrf_user_getvar(a,"uvmet",time)
wind10rot = wrf_user_getvar(a,"uvmet10",time)

printVarSummary(windrot)
printVarSummary(wind10rot)

windU = windrot(0,:,:,:)
windV = windrot(1,:,:,:)
windU10 = wind10rot(0,:,:)
windV10 = wind10rot(1,:,:)

printVarSummary(windU10)
printVarSummary(windV10)


;windU   = wrf_user_getvar(a,"ua",time)
;windV   = wrf_user_getvar(a,"va",time)
;windU10   = wrf_user_getvar(a,"U10",time)   ; WARNING - need to rotate?
;windV10   = wrf_user_getvar(a,"V10",time)
	
type = "png"

; define fields to plot
;vNames = (/"tr1","tr2","tr3","tr4","tr5","tr6","tr17_1","tr17_5","PBLH","REFL_10CM","SWDOWN","HGT"/)
;fullNames  = (/"tr_mobile","tr_area","tr_oilgas","tr_agri","tr_lnox","tr_nox_nonCO","tr_BDY","tr_STRAT","PBLH","REFL","SWDOWN","HGT"/)
;vNames = (/"tr1","tr2","tr3","tr4","tr5","tr6","tr17_1","PBLH","REFL_10CM","SWDOWN","HGT"/)
;fullNames  = (/"tr_areaNO2","tr_areaCO","tr_mobileNO2","tr_mobileCO","tr_lnox","tr_pointNO2","tr_BDY","PBLH","REFL","SWDOWN","HGT"/)
vNames = (/"Age/"
fullnames = (/"Age/"
nNames = dimsizes (vNames)   ; Number of variables on the file

; overplot main site locations
; BAO, Chatfield, DenverLCasa,FC West,NREL,Platteville
sitex   = (/  -105.0038  ,  -105.0704, -105.005, -105.1414, -105.1779, -104.726/)
sitey   = (/ 40.05,        39.5345,     39.779,    40.5927,   39.7437,   40.1827/)
nsites = dimsizes(sitex)
       
Times = a->Times
       
do n=0,nNames-1              ; Loop through each variable
    	 

  res = True                               ; Set some Basic Plot options
  res@MainTitle = ""
  res@InitTime = False                     ; Do not plot time or footers
  res@NoHeaderFooter = True
  res@gsnMaximize=True

  pltres = True                            ; Set plot options
  pltres@PanelPlot  = True   ; Tells wrf_map_overlays not to remove contours
  pltres@Maximize = True

  mpres = True                             ; Set map options
  mpres@mpGeophysicalLineColor = "Black"
  mpres@mpNationalLineColor    = "Black"
  mpres@mpUSStateLineColor     = "Black"
  mpres@mpOutlineBoundarySets = "AllBoundaries"
  mpres@mpProvincialLineThicknessF = 20
  mpres@mpGeophysicalLineThicknessF = 20

  mpres@mpNationalLineThicknessF = 20
  mpres@mpDataBaseVersion     = "Ncarg4_1"
  mpres@mpDataSetName         = "Earth..2"
  mpres@DynamicAreaGroups = 8
  mpres@mpGridAndLimbOn       = True
  mpres@mpPerimOn             = True

; get the variable field, not yet scaled

  if vNames(n) ne "Age" then 
     v  =  wrf_user_getvar(a,vNames(n),time)
  end if
  if vNames(n) eq "Age" then 
  
  tr6_a = wrf_user_getvar(a1,"tr1",time)
  tr6_b = wrf_user_getvar(a1,"tr2",time)
  tr6 = tr6_a + tr6_b
  tr5 = wrf_user_getvar(a1,"tr7",time)
  v = tr5
  v = where(tr6.gt.0.01, tau * log(tr5/tr6),  -999. )
  v@_FillValue = -999.
  print(" min/max Age  "+min(v)+"  "+max(v) )


  if (vNames(n) .ne. "PBLH" .and. vNames(n) .ne. "SWDOWN" .and. vNames(n) .ne. "HGT") then

  ;for plotting set to a low values to avoid graphs with no data in movie 
  if (max(v) .le. 0) then 
        v(:,:,:) = 0.
        v(:,150,150) = 0.001
  end if	   


; derive PBLH weighted average 
    col=v(0,:,:)
  col(:,:) = 0.
      nz=dimz(0)
        ny=dimz(1)
	nx=dimz(2)
	g=9.81
 	mw=28.966e-3/6.022e23 
  	do ix=0,nx-1
          do iy=0,ny-1
		  a1D      = ndtooned(z(:,iy,ix))
  		  dsizes_a = nz
		  limit = ndtooned(pblh(iy,ix)+ht(iy,ix))
  		  indices  = ind(a1D.le.limit)
; if PBL is lower than first level
if (a1D(0).gt.limit) then 
    indices = 0
end if
  		  dim_ida  = dimsizes(indices)    
  		  npts     = dim_ida(0)       ; number of elements 
		  dcolair = 0.
		  do iind=0,npts-1
		        vlay	= (v(indices(iind),iy,ix)+v(1+indices(iind),iy,ix))/2.   ; ppb (scale later)
		   	 dp = prs(indices(iind),iy,ix)-prs(1+indices(iind),iy,ix)
		   	 dcolair = dcolair+dp
		   	 col(iy,ix) = col(iy,ix)+vlay*dp
		  end do
    	  	  col(iy,ix) = col(iy,ix)/dcolair
		  delete(indices)
	      end do
       end do
   end if

   opts = res                ; Set plot options to be used by all variables
   opts@cnFillOn = True
   opts@cnFillMode       = "AreaFill"
  
   dimv = dimsizes(v)        ; dimension size of the variable
   rank = dimsizes(dimv)     ; rank [ie: number of dimensions]
  
   opts@cnLinesOn = False
   opts@cnLineLabelsOn = False
   opts@cnLevelSelectionMode = "ExplicitLevels"
   opts@cnFillMode = "AreaFill"

   scale= 1.	 ; for tr1-tr4
   if (vNames(n) .eq. "tr1" ) then     ; test how tracer looks if e_tr1 and e_tr2 are scaled to the max of their sum
      	       scale = 1.    ; latest runs are alreadty scaled /107.      
   end if
   if (vNames(n) .eq. "tr5" ) then
       scale = 1e9
   end if
   if (vNames(n) .eq. "PBLH" .or. vNames(n) .eq. "SWDOWN" .or. vNames(n) .eq. "REFL_10CM" .or. vNames(n) .eq. "HGT" ) then
       scale = 1
   end if

; set levels for plotting

  levels  = (/ 0.001, 0.01, 0.05,0.10, 0.25 ,0.5, 1. ,2., 3., 5., 7.5, 10., 15., 20.,100/)
 
 ;if (vNames(n) .eq. "tr3" .or. vNames(n) .eq. "tr4" ) then 
 ;  levels  = (/1., 10., 50., 100. ,200., 300., 500., 750., 1000., 1500., 2000.,3000.,4000.,6000.,10000 /)
 ; end if
 ;if (vNames(n) .eq. "tr6" .or. vNames(n) .eq. "tr17_5" ) then 
 ;  levels  = (/0.1,0.5,1.,5, 10., 50., 100. ,200., 300., 500., 750., 1000., 1500., 2000.,3000./)
 ;end if

 if (vNames(n) .eq. "Age"  ) then 
   levels = (/0.10,0.5,1,2,3,5,7,10,15,20,25,30,35,40,45,50/)
 end if

 if (vNames(n) .eq. "tr17_1"  ) then 
   levels = (/10,13,17,20,23,26,30,34,38,45,50,55,60,70,80/)
 end if
  
  if (vNames(n) .eq. "tr5" ) then 
    levels  = (/0.001,0.01,0.05,0.1, 0.15 ,0.3, 0.5, 1.0, 1.5, 2., 3., 4., 5., 7.,10./)
  end if
  if (vNames(n) .eq. "PBLH" ) then 
	      levels  = (/50., 100.,200., 350., 500. ,750., 1000., 1500., 2000., 2500., 3000., 4000., 5000.,6000.,7000. /)
   end	if
  if (vNames(n) .eq. "HGT" ) then 
	      levels  = (/0.,25., 50.,100., 175., 250. ,375., 500., 750., 1000., 1250., 1500., 2000., 2500.,3000./)
   end	if
  if (vNames(n) .eq. "SWDOWN" ) then 
	      levels  = (/50., 100.,150.,200., 250.,300., 350., 400., 450., 500., 550., 600., 650., 700., 750. /)
   end	if
 if (vNames(n) .eq. "REFL_10cm" ) then 
	      levels  = (/-40., -30., -20., -10., -5., 0., 5., 10., 20., 30., 40., 50., 60., 70., 80./)
   end	if
 
   v= v * scale

   do imap=1,1   ; loop over 2 domains 

; define plotting regions

;      if (imap .eq. 0) then
;      	 lats = (/  33.,  40. /)
;       	 lons = (/ 125., 130. /)
; 	 loc = wrf_user_ll_to_ij(a, lons, lats, True)
;  	 x_start = loc(0,0) - 1
;  	 x_end   = loc(0,1) - 1
;  	 y_start = loc(1,0) - 1
;  	 y_end   = loc(1,1) - 1	;
;	 maptit="NFRMA"
;     else	       		 
;    x_start = 5
; 	 x_end = dims(1)-5
; 	 y_start =5
;  	 y_end = dims(0)-5
    	 lats = (/  32.,  39. /)
       	 lons = (/ 123.5, 132.5 /)
 	 loc = wrf_user_ll_to_ij(a, lons, lats, True)
  	 x_start = loc(0,0) - 0
  	 x_end   = loc(0,1) - 0
  	 y_start = loc(1,0) - 0
  	 y_end   = loc(1,1) - 0	
	 maptit="Koread02"
;     end if	

     mpres1 = True
     mpres1@ZoomIn = True
     mpres1@gsnMaximize=True
     mpres1@Xstart = x_start
     mpres1@Ystart = y_start
     mpres1@Xend = x_end
     mpres1@Yend = y_end
  
; Options for Label Bar
     labstrings = sprintf("%5.0f",levels)
     opts@cnExplicitLabelBarLabelsOn = True
     opts@lbLabelStrings = labstrings
     opts@lbLabelStride = 1

 opts@lbLabelAutoStride          = True    ; Spacing of lbar labels.
  opts@lbOrientation              = "horizontal"
  opts@lbBoxMinorExtentF          = 0.13
  opts@pmLabelBarWidthF           = 0.90 
  opts@pmLabelBarHeightF          = 0.2
  opts@lbLabelFontHeightF         = 0.01
; plot 2D variables

    if ( rank .eq. 2 ) then
        plottit = cdate+"_"+vNames(n)+"_"+maptit

  	wks = gsn_open_wks(type,plottit)       ; Create a plot workstation
        gsn_define_colormap(wks,"WhiteBlueGreenYellowRed")   ; 254 colors

      	opts@FieldTitle = fullNames(n) +" : "+ chartostring((Times))  ; overwrite field name
      	opts@cnLevels = levels

printVarSummary(v)
        v_zoom  = v(y_start:y_end,x_start:x_end) 
;	    u_plane = windU(0,y_start:y_end,x_start:x_end)  ; only used because of (delete) within loop
;     	    v_plane = windV(0,y_start:y_end,x_start:x_end)
; with WRF projection:  

        contour = wrf_contour(a,wks,v_zoom(:,:),opts)
    	plot = wrf_map_overlays(a,wks,(/contour/),pltres,mpres1)

; plot as Lat/Lon:
;???? tc2@lat2d = a->XLAT(0,:,:)
;???  tc2@lon2d = a->XLONG(0,:,:)
;;opts@cnLinesOn            = False
;;plot = gsn_csm_contour_map(wks,v_zoom(:,:),opts)
	

        dum = new(2,graphic)
    	colors=(/"black","white"/)
	types = (/16,16/)
	sizes = (/0.009 ,0.007/)
	do idum=0,1 
    	   pmres = True
   	   pmres@gsMarkerColor = colors(idum)
    	   pmres@gsMarkerIndex = types(idum)
	   pmres@gsMarkerSizeF = sizes(idum)
 	   dum(idum) = gsn_add_polymarker(wks, plot, sitex, sitey, pmres)
	   delete(pmres)
         end do 

 ;---Attach the shapefile polylines using files read off gadm.org/country.
      	  usa_shp_name    = "KOR_adm1.shp"
   	  lnres                  = True
     	  lnres@gsLineColor      = "black"
      	  lnres@gsLineThicknessF = .9
      	  usa_id = gsn_add_shapefile_polylines(wks,plot,usa_shp_name,lnres)
;      	  draw(plot)       ; This will draw the map and the shapefile outlines.
;      	  frame(wks)       ; Advance the frame
     	  world_shp_name    = "countries.shp"
   	  lnres                  = True
     	  lnres@gsLineColor      = "black"
      	  lnres@gsLineThicknessF = 1.1
      	  world_id = gsn_add_shapefile_polylines(wks,plot,world_shp_name,lnres)
      	  draw(plot)       ; This will draw the map and the shapefile outlines.
      	  frame(wks)       ; Advance the frame
	  delete(v_zoom)
    end if   ; rank 2

;; 3D TRACERS, assume no time dimension
     	
    if ( rank .eq. 3 ) then
       	 do levs = 0,nlevsplot-1 ; If 4D array - plot only selected levels
	  height = levs_height(levs)
          opts@FieldTitle = fullNames(n)+" "+ levs_height_tit(levs)+ "   "+chartostring((Times))   ; overwrite field name

   	  plottit = cdate+"_"+vNames(n)+"_"+levs_height_tit(levs)+"_"+maptit

      	  wks = gsn_open_wks(type,plottit)       ; Create a plot workstation
      	  gsn_define_colormap(wks,"WhiteBlueGreenYellowRed")   ; 254 colors

          if (height .lt. 0) then
	     v_zoom  = v(0,y_start:y_end,x_start:x_end)	
 	     u_plane = windU10(y_start:y_end,x_start:x_end)  
     	     v_plane = windV10(y_start:y_end,x_start:x_end)
	  end if     
     	  if (height .gt. 0) then  
	  	u_plane  = wrf_user_intrp3d(windU(:,y_start:y_end,x_start:x_end),z(:,y_start:y_end,x_start:x_end),"h",height,0.,False)
      		v_plane  = wrf_user_intrp3d(windV(:,y_start:y_end,x_start:x_end),z(:,y_start:y_end,x_start:x_end),"h",height,0.,False)
       		v_zoom = wrf_user_intrp3d(v(:,y_start:y_end,x_start:x_end),z(:,y_start:y_end,x_start:x_end),"h",height,0.,False)
	  end if
	  if (height .eq. 0) then 	
	      v_zoom  = col(y_start:y_end,x_start:x_end)* scale	
 	      u_plane = windU10(y_start:y_end,x_start:x_end)  
     	      v_plane = windV10(y_start:y_end,x_start:x_end)
          end  if

     	  labstrings = sprintf("%4.1f",levels)
;	   	   if (levs .gt. 1) then     
;	      	      levels = levels*0.1   
;  	      	      labstrings = sprintf("%5.3f",levels)
;	           end if
 	  opts@cnLevels = levels

    	  u_plane = u_plane*1.94386     ; kts
     	  v_plane = v_plane*1.94386     ; kts
     	  u_plane@units = "kts"
     	  v_plane@units = "kts"

  ; Plotting options for Wind Vectors                 
      	  optsW = res          
      	  optsW@FieldTitle = "Wind"   ; overwrite Field Title
      	  optsW@NumVectors = 36         ;  47       ; wind barb density
      	  optsW@vcRefLengthF     = 0.02
      	  vector = wrf_vector(a,wks,u_plane,v_plane,optsW)
      	  delete(optsW)
; for entire domain
          ;contour = wrf_contour_ggp(a,lat2U,lon2U,wks,v(0,lev,:,:),opts)
          ;plot = wrf_map_overlays_ggp(a,lat2U, lon2U,wks,(/contour/),pltres,mpres)
; selected domain WRF projection
	  contour = wrf_contour(a,wks,v_zoom,opts)
	  plot = wrf_map_overlays(a,wks,(/vector,contour/),pltres,mpres1)
; plot as Lat/Lon:

;;opts@cnLinesOn            = False
;;plot = gsn_csm_contour_map(wks,v_zoom,opts)
	

          dum = new(2,graphic)
	  colors=(/"black","white"/)
	  types = (/16,16/)
	  sizes = (/0.0075,0.006/)
	  do idum=0,1 
	     pmres = True
   	     pmres@gsMarkerColor = colors(idum)
    	     pmres@gsMarkerIndex = types(idum)
	     pmres@gsMarkerSizeF = sizes(idum)
 	     dum(idum) = gsn_add_polymarker(wks, plot, sitex, sitey, pmres)
	     delete(pmres)
	  end do 

;---Attach the shapefile polylines using files read off gadm.org/country.

         usa_shp_name    = "KOR_adm1.shp";     
          lnres                  = True
          lnres@gsLineColor      = "black"
          lnres@gsLineThicknessF = .9
          usa_id = gsn_add_shapefile_polylines(wks,plot,usa_shp_name,lnres)
;          draw(plot)       ; This will draw the map and the shapefile outlines.
;          frame(wks)

    	  world_shp_name    = "countries.shp"
   	  lnres                  = True
     	  lnres@gsLineColor      = "black"
      	  lnres@gsLineThicknessF = 1.1  
      	  world_id = gsn_add_shapefile_polylines(wks,plot,world_shp_name,lnres)
      	  draw(plot)       ; This will draw the map and the shapefile outlines.
      	  frame(wks)       ; Advance the frame


	  delete(u_plane)
	  delete(v_plane)
	  delete(v_zoom)
        end do  ; levels
     end if    ; rank 3	      
   end do ; (imap)
   
   if (rank .eq. 3) then 
       delete(col)
   end if
   delete(v)
   delete(dimv)
   delete(opts)
   delete(mpres)
   delete(mpres1)
   delete(plot)
   delete(wks)

end do   ; tracers/names
delete(a)
  
end