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/contributed.ncl"    ; Needed for dim_avg_n_Wrap
load "$NCARG_ROOT/lib/ncarg/nclscripts/contrib/ut_string.ncl" 
load "/home11/grad/2011/abentley/ncl/lib/shea_util.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/esmf/ESMF_regridding.ncl"
load "/home/kb823469/nclscripts/tTick.ncl"
load "/home/kb823469/nclscripts/latlon.ncl"



; Author: Jeremy Berman, SUNY Albany
; Updated: 06.22.18
; Description: Compute area-average of a variable
;
; Notes: Modify the area-averaging radius (e.g., 300 km) and the variable you want to do it for (read it in)
; the subroutine "area_average" will compute the area-average of the quantity by passing it to a fortran program
; the fortran subroutine, "circulation_llgrid" , computes the area-average by passing the variable, lat, lon vectors, 
; first 3 dimensions (reversed order!), and the output variable
; My subroutine "area_average" reorders the array to be passed to the fortran program. But double-check for any inconsistencies
; the third dimension, can be ensemble members, time, vertical levels, etc. Doesn't really matter as long as other two are lat and lon


external AREAAVG "./calc_area_average_fix_2.so" ; change the file path to the area-average fortran object

;-----------------------------------------------------------------------------------------
; Function to area average a quantity
;-----------------------------------------------------------------------------------------
function area_average(var_std,circrad)
begin 

dim_3 = dimsizes(var_std&lon)
dim_2 = dimsizes(var_std&lat)
dim_1 = dimsizes(var_std&$var_std!0$)

test_std_avg = var_std(lat|:,lon|:,$var_std!0$|:) ; reorganize the array
std_avg = new(dimsizes(test_std_avg),typeof(test_std_avg))

AREAAVG :: circulation_llgrid(test_std_avg, var_std&lat, var_std&lon, circrad, dim_3, dim_2, dim_1, std_avg)

std_avg!0 ="lat"
std_avg&lat=var_std&lat
std_avg!1 ="lon"
std_avg&lon=var_std&lon
std_avg!2 = var_std!0 ; time dimension ;;;Jeremy had an "ensemble0" dimension
std_avg&$var_std!0$ = var_std&$var_std!0$

final_array = std_avg($var_std!0$|:,lat|:,lon|:) ; reorganize the array to original
copy_VarMeta(var_std,final_array)
return(final_array)
end

;-----------------------------------------------------------------------------------------
; main script

; read in the variable you want to area average (add code lines below for your dataset!)
; now do area-averaging

;begin


yr = 2019
 syyyy = yr    ; Start Date
   smm = 07
   sdd = 15
   shh = 00

 eyyyy = yr    ; End Date
   emm = 08
   edd = 15
   ehh = 18

timeUnits = "hours since 1800-01-01 00:00:00"   ; Time is represented as "hours since 1800-01-01 00:00:00"

startTime = ut_inv_calendar(syyyy,smm,sdd,shh,00,00,timeUnits,0)   ; Convert start time to the above mentioned units
endTime   = ut_inv_calendar(eyyyy,emm,edd,ehh,00,00,timeUnits,0)   ; Convert end time to the above mentioned units

ntimes = toint(((endTime - startTime)/6) + 1)   ; Find the number of times from the start time to the end time.
                                                ; E.g., if startTime = 700024 and endTime = 700048, and the data is 
                                                ; available every 6 hours, we'd expect ntimes = 5. 
                                                ; We see that (700048 - 700024)/6 + 1 = 4+1 = 5

timestepfactor = 4 ;in 6hr intervals

times = new( (/ntimes/), "float")  ; Create a new 1-dimensional float array of size nTimes
times@units = timeUnits            ; Give the new times array the same units as timeUnits (in this case: 
                                   ; "hours since 1800-01-01 00:00:00")

do n=0,ntimes-1                           ; This do loop fills the times array with the times from the startTime to the endTime
   times(n) = tofloat(startTime + (6*n))  ; in 6 hour increments, in units of "hours since 1800-01-01 00:00:00"
end do



levs_irro = (/300,250,200/)

fileName = "/cfsr/data/"+yr+"/u."+yr+".0p5.anl.nc"
inFile_u = addfile( fileName, "r" ) 

fileName = "/cfsr/data/"+yr+"/v."+yr+".0p5.anl.nc"
inFile_v = addfile( fileName, "r" ) 

fileName = "/cfsr/data/"+yr+"/t."+yr+".0p5.anl.nc"
inFile_t = addfile( fileName, "r" )   

u     = inFile_u->u({times},{levs_irro},:,:) 
v     = inFile_v->v({times},{levs_irro},:,:)  

time = u&time
nlat = dimsizes(u&lat)
nlon = dimsizes(u&lon)

  
  lev_pv  = (/300,250,200/)
  lev_pv_pa  = (/30000,25000,20000/)  
  upv = inFile_u->u({times},{lev_pv},:,:)
  vpv = inFile_v->v({times},{lev_pv},:,:)
  tpv = inFile_t->t({times},{lev_pv},:,:)

  lev_pv_pa@units = "Pa"
  pv_all = PotVortIsobaric(lev_pv_pa,upv,vpv,tpv,tpv&lat,1,0)
  pv_avg = dim_avg_n_Wrap(pv_all,1)
  pv_avg = pv_avg*(10^6)
 printMinMax(pv_avg, True)
 


;***************calculate divergent (irrotational) wind components*************** 
div = uv2dvF_Wrap(u,v)  ; calculate divergence from the u and v wind components

ud    = new ( dimsizes(u), typeof(u), u@_FillValue )   ; Create new NCL variable for u-component of divergent wind
vd    = new ( dimsizes(v), typeof(v), v@_FillValue )   ; Create new NCL variable for v-component of divergent wind

dv2uvf(div,ud,vd)    ; calculates the divergent  wind components

;***Give metadata to the divergent wind components***
copy_VarCoords(u, ud ) 
copy_VarCoords(u, vd ) 
ud@long_name  = "Zonal Divergent Wind"
ud@units      = u@units
vd@long_name  = "Meridional Divergent Wind"
vd@units      = v@units

ud_avg = dim_avg_n_Wrap(ud,1) ; Calculate layer-average u-component of divergent wind using second dimension (1) of ud (the levels: 30000 - 20000 pa)
vd_avg = dim_avg_n_Wrap(vd,1) ; Calculate layer-average v-component of divergent wind using second dimension (1) of vd (the levels: 30000 - 20000 pa)

printMinMax(ud_avg, True)
printMinMax(vd_avg, True)


lat = u&lat
lon = u&lon
rad    = 4.0*atan(1.0)/180.0

nlat = dimsizes(lat)
dlon = (lon(2)-lon(1))*rad 
dlat = (lat(2)-lat(1))

pvgrad_lat = pv_all
pvgrad_lon = pv_all
	
do nl=0,nlat-1                      ; loop over each latitude
	dX_double = 6378388.*cos(0.0174533*lat(nl))*dlon  ; constant at this latitude
	dX = tofloat(dX_double)
	pvgrad_lon(:,:,nl,:) = center_finite_diff (pv_all(:,:,nl,:), dX , False,0)
end do

dY_double = dlat*111200.0 ;m/deg lat
dY = tofloat(dY_double)

pvgrad_lat = center_finite_diff_n (pv_all, dY , False,0,2)

pv_adv = pv_all
pv_adv = -(ud*pvgrad_lon + vd*pvgrad_lat)
copy_VarCoords(pv_all, pv_adv ) 
pv_adv = pv_adv*86400*10^6 
pv_adv@units  = "PVU/day"
pv_adv@long_name  = "PV Advection by the Irrotational Wind"
pv_adv_avg = dim_avg_n_Wrap(pv_adv,1)

 
 circrad = 300.0 ; select your area-average radius (here in km)
 
 irro_pvadv_areaavg = area_average(pv_adv_avg,circrad) ; this is your area-averaged quantity
 printVarSummary(irro_pvadv_areaavg)
 

 ;;##################################################    
;;Output count to netCDF file
;;##################################################   

out_dir = "/keyserlab_rit/kbiernat/July_Aug_2019_Greenland_event/area_average_files/"
out_fil = out_dir+"irro_pvadv_areaavg_300km.nc"  ;-Output filename
var = "irro_pvadv_areaavg"



lat@units              = "degrees_north"
lat@long_name          = "latitude"
lat@grid_resolution    = "0.5_degrees"
lat@mapping            = "cylindrical_equidistant_projection_grid"
lat@coordinate_defines = "center"
lat@delta_y            = 0.5
lat@actual_range       = (/-90,90/)

lon@units              = "degrees_east"
lon@long_name          = "longitude"
lon@grid_resolution    = "0.5_degrees"
lon@mapping            = "cylindrical_equidistant_projection_grid"
lon@coordinate_defines = "center"
lon@delta_x            = 0.5
lon@actual_range       = (/-180,179.5/)


; File and Variable Attributes

fAtt = True

fAtt@creation_date  = systemfunc ("date")
fAtt@created_by     = "User: "+systemfunc ("whoami")
fAtt@description    = "Area average relative vorticity"

vAtt = 1
vAtt@long_name  = "Area average PV adv by nondiv wind"
;Att@_FillValue = 1e+20


;--------- Output a compressed netCDF ------

setfileoption("nc","Format","NetCDF4Classic")
setfileoption("nc","CompressionLevel",1)

;--------- Initialize the netCDF file ------

system("/bin/rm -f "+out_fil)
outFile = addfile(out_fil, "c" )

;fileattdef( outFile, fAtt )            ; Set file attributes

; Specify dimension coordinates

dimNames = (/ "time", "lat","lon"/)
dimSizes = (/ntimes,nlat,nlon/)
dimUnlim = (/False, False, False/)

chunks = (/1,nlat,nlon/)    ; Not as neccessary in this line but may as well force it to be user friendly.


filedimdef( outFile, dimNames, dimSizes, dimUnlim )
filechunkdimdef(outFile,dimNames,chunks,dimUnlim)

filevardef( outFile, "time", "double", "time" )
filevardef( outFile, "lat", "float", "lat" )
filevardef( outFile, "lon", "float", "lon" )
filevardef( outFile, var, "float", (/ "time", "lat","lon"/) )

filevarattdef( outFile, "time", time )
filevarattdef( outFile, "lat", lat )
filevarattdef( outFile, "lon", lon )
filevarattdef( outFile, var, vAtt )


; Write coordinates to record

outFile->time = (/ time /)
outFile->lat = (/ lat /)
outFile->lon = (/ lon /)
outFile->$var$(:,:,:) = new( (/ntimes,nlat,nlon/), "float")

print("netCDF file for "+var+" initialized on "+systemfunc("date"))

delete(outFile)
delete(dimNames)
delete(dimSizes)
delete(dimUnlim)
delete(chunks)

outfile = addfile(out_fil,"w")
outfile->$var$(:,:,:) = (/irro_pvadv_areaavg/)
    
  

 

;end