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/wrf/WRFUserARW.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/wrf/WRF_contributed.ncl"

begin
  
  ;## open work station
  timedate = systemfunc("date +$Y-%m-%d-%H-%M")
  wks = gsn_open_wks("pdf","bruntvaisala"+timedate)         
  gsn_define_colormap(wks,"WhiteYellowOrangeRed")
    
    res                       	= True 
  
  ;TICKMARK RESOURCES (ON/OFF - FONT SIZE)
	res@pmTickMarkDisplayMode = "Always" 
	res@tmXTOn				= False
	res@tmXBOn				= True
	res@tmYLOn				= False
	res@tmYROn				= False
	res@tmXTLabelFontHeightF		= 0.025
	res@tmXBLabelFontHeightF		= 0.025
	res@tmYLLabelFontHeightF		= 0.025
	res@tmYRLabelFontHeightF		= 0.025
	
  ; MAP RESOURCES.
	res@mpGridAndLimbOn			= True
	res@mpGridLineDashPattern		= 2
	res@mpGeophysicalLineColor		= "Black"
	res@mpGeophysicalLineThicknessF		= 1.5
	res@mpNationalLineColor			= "Black"	
	res@mpNationalLineThicknessF		= 1.5
	res@mpUSStateLineColor			= "Black"
	res@mpUSStateLineThicknessF		= 1.5
	res@mpPerimLineColor			= "Black"
            
  ;LABEL BAR RESOURCES
	res@lbLabelAutoStride			= True            
	res@lbOrientation			= "Vertical"
	res@lbLabelFontHeightF			= 0.025
	
  ;CONTOUR RESOURCES
	res@cnFillOn				= True             
	res@cnLinesOn				= False  
	res@cnLineColor				= "Blue"
	res@cnLineThicknessF			= 2.            
	res@cnInfoLabelOn			= False          
	res@cnLineLabelsOn			= False 
	res@cnLineLabelPerimOn			= True
	res@cnLineLabelBackgroundColor		= 0.
;	res@cnLevelSelectionMode   		= "ManualLevels"
;	res@cnMinLevelValF         		= -20.
;	res@cnMaxLevelValF         		= 30.
;	res@cnLevelSpacingF		  	= 2.
	
  ;GSN RESOURCES
	res@gsnSpreadColors			= True               
	res@gsnLeftString			= " "
	res@gsnRightString			= " "
	;res@gsnDraw				= False
	;res@gsnFrame				= False
	res@gsnMaximize				= True
	
  ;TITLE RESOURCES
	res@tiMainOn				= True
	res@tiMainPosition			= "Center"
	res@tiMainFontHeightF			= 0.025
  
  ;PROJECTION RESOURCES
	res@tfDoNDCOverlay 			= True
	res@gsnAddCyclic 			= False
 	res@mpProjection			= "Stereographic"
	res@mpLimitMode				= "Corners"
	res@mpLeftCornerLatF                  = 65.80
        res@mpLeftCornerLonF                  = 132.8014
        res@mpRightCornerLatF                 = 65.0
        res@mpRightCornerLonF                 = -45.0
  	res@mpCenterLatF 			= 90
  	res@mpCenterLonF 			= -180
  
  ;## day when integration started
  day	= "15"
  
  
  ;## hour when integration started
  hour	= "00"
 setfileoption("nc","Format","LargeFile")

    
  ;## path to wpp outfiles
  ;path	= "/wrkdir/moreira/wrf-pacman-compilation-03-22-11/WRFV3/Ensemble-Runs/run-"+day+"@"+hour+"/postprd/"
  path = "/import/wrkdir1/asemenov/WRF/WRFV3/10_25.03_res_30_10km/postprd/"
  
  ;## file name
  files	= systemfunc("ls "+path+"WRFPRS_d02.*")


  ;## create array to hold result

n		= new ( (/2,19,429,399/), "float" )
  n@_FillValue	= -99999
  n		= n@_FillValue
  n@long_name	= "brunt-vaisala frequency"
  n@units	= "s-1"
  n!0		= "time"
  n!1		= "plevels"
  n!2		= "lat"
  n!3		= "lon"
  printVarSummary(n)

  
  N2		= n
  N2		= N2@_FillValue
  
;do j 	= 0, dimsizes(files)-1
do j   = 0, dimsizes(files)-143
  
  fname	= files(j)
  

  ;## open file
  print("Opening file: "+fname+".grb")
  f     = addfile (fname+".grb", "r")              
 
  
  ;## get variables 
  print			("Getting variables...")
  g			= 9.81
  theta			= f->POT_GDS5_HYBL_10
  z			= f->HGT_GDS5_HYBL_10
  p			= f->PRES_GDS5_HYBL_10
  levels		= f->lv_HYBL3
  lat			= f->g5_lat_1
  lon			= f->g5_lon_2
  
;  print("Levels: ")
;  print(levels)
;  printVarSummary(theta)
;  printVarSummary(p)
  printVarSummary(z)




  ;## calculate n2
  print			("Calculating dTdZ and N2...")
  dThetadZ		= center_finite_diff_n(theta,z,False,0,0)
  n2			= (1/theta) * g * dThetadZ


;  printVarSummary(dThetadZ)
  printVarSummary(n2)   
 
  ;## interpolate to pressure coords
 pls			= (/ 1000.,925.,850.,800.,750.,700.,650.,600., \
 				500.,450.,400.,350.,300.,250.,200.,150.,100.,70.,50. /)				
;pls                  = (/ 1000.,950.,925.,900.,875.,850.,825.,800.,775.,750.,725.,700.,675.,650.,625.,600.,575.,550.,525., \
;                               500.,475.,450.,425.,400.,350.,300.,250.,200.,150.,100.,70.,50. /) 




  print			("Interpolating...")
  n2_pls 		= wrf_user_intrp3d(n2,p,"h",pls,0.,False)
  printVarSummary(n2_pls)   
 printMinMax(n2,True)
  printMinMax(n2_pls,True) 
  printMinMax(n,True)



  ;## get n
 print			("Calculating N...")
  n(j,:,:,:)		= where(n2_pls.gt.0,sqrt(n2_pls),n2_pls@_FillValue)
  N2(j,:,:,:)		= (/ n2_pls /)

;  printMinMax(theta,True)
;  printMinMax(z,True)
;  printMinMax(dThetadZ,True)
;  printMinMax(n2,True)
;  printMinMax(n2_pls,True) 
;  printMinMax(n,True)


  ;## plot
  print		("Plotting...")
  jj		= 2
;  res@gsnLeftString	= n@long_name+" ("+n@units+") at "+pls(jj)+" hPa at time "+j
;  res@gsnLeftString    = N2@long_name+" ("+N2@units+") at "+pls(jj)+" hPa at time "+j 

 plot		= gsn_csm_contour_map(wks,N2(j,jj,:,:),res)
  
end do

  
  ;## write to file
  print		("Writing to file...")
  a		= systemfunc("rm -f brunt.nc")
  fileout	= addfile("brunt.nc","c")
  fileout->n2=N2
  fileout->n=n
  fileout->lat=lat
  fileout->lon=lon
  fileout->plevels=pls
  
end