;;***************************************************************************************
;;Description: Calculates the total number of Arctic Cyclones (ACs) that are situated within a 
;;             user-specified radius (in km) of each gridpoint within a domain of interest
;;             using great circle distances and outputs the total count to a netCDF file
;;             This code only counts a AC once per grid point. For example, if a  is within 500 km
;;             of the same grid point for multiple time steps, the AC is not counted 
;;             multiple times for that grid point; the AC is only counted once for that grid point!
;;Script adapted from a similar counting script created by Philippe Papin
;;Last Modified on 17 MArch 2019 by Kevin Biernat
;;***************************************************************************************

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"   

begin

fileName_skill = "/keyserlab_rit/kbiernat/arctic_forecast_skill/GEFS_JJA/low_high_skill_aa_spread_aa_rmse_v2_files/time_all_2007to2017_JJA.nc"
inFile_skill = addfile( fileName_skill, "r" )
time_skill  = inFile_skill->time(:)
ntim_skill = dimsizes(time_skill)

inFile = addfile( "/keyserlab_rit/kbiernat/arctic_cyclone_tracks/arctic_cyclones_JJA_skill/arctic_cyclones_climo_JJA.nc","r")


all_lat = inFile->lat(:,:) ;; 2d array where values are (nTracks,nTimesTrack)
all_lon = inFile->lon(:,:) ;; 2d array where values are (nTracks,nTimesTrack)
time = inFile->time(:,:)

;;; first create lat/lon domain that is same as lat/lon domain in which TPVs are tracked
latRange = (/90,0/)  
lonRange = (/0,360/) 

;;; increments at 1.0 degree intervals
domain_lat = fspan(latRange(0),latRange(1),91)
domain_lon = fspan(lonRange(0),lonRange(1),361)

;; will hold the count of unique ACs within 500km of a grid point
count_total = new((/dimsizes(domain_lat),dimsizes(domain_lon)/),"integer")
count_total!0 = "lat"
count_total!1 = "lon"
count_total&lat = domain_lat
count_total&lon = domain_lon
count_total = 0


count_track = new((/dimsizes(domain_lat),dimsizes(domain_lon)/),"integer")
count_track!0 = "lat"
count_track!1 = "lon"
count_track&lat = domain_lat
count_track&lon = domain_lon


;; lats and lons of domain, dimensioned same as count_total
lat_tot = conform(count_total,count_total&lat,0)
lon_tot = conform(count_total,count_total&lon,1)


nTracks = dimsizes(all_lat(:,0))

;; Loop through all of the tracks
do iTrack=0,nTracks-1

  AC_timeInds := ind(.not. ismissing(all_lat(iTrack,:)))
  AC_lats_all := all_lat(iTrack,AC_timeInds)
  AC_lons_all := all_lon(iTrack,AC_timeInds)
  time_all := time(iTrack,AC_timeInds)
  
  
  count_track = 0
  
  ;; Loop through all the times in the AC track
  do iTimeTrack=0,dimsizes(AC_timeInds)-1 
  	do k=0, ntim_skill-1
    	if (time_all(iTimeTrack).eq.time_skill(k)) then

    		lat_AC = count_total&lat 
    		lon_AC = count_total&lon 

   	 		lat_AC = AC_lats_all(iTimeTrack)
    		lon_AC = AC_lons_all(iTimeTrack)

    		;;-----------------------------------------------------------------
    		;;Calculate Great Circle Distances between points
    		;;-----------------------------------------------------------------   

    		r = 6370. ; radius of spherical Earth (km)    
    		d2r = 0.0174532925199433 ; degrees to radians conversion factor
       
    		lat1 = lat_tot
    		lon1 = lon_tot
    		lat2 = conform(count_total,lat_AC,0)
    		lon2 = conform(count_total,lon_AC,1)  
       
    		lat1 = lat1*d2r
    		lon1 = lon1*d2r  
    		lat2 = lat2*d2r
    		lon2 = lon2*d2r    
  
  
    		dlat = lat2-lat1
    		dlon = lon2-lon1
    		latTerm = sin(.5*dlat) 
    		latTerm = latTerm*latTerm
    		lonTerm = sin(.5*dlon) 
    		lonTerm = lonTerm*lonTerm*cos(lat1)*cos(lat2)
    		dAngle = sqrt(latTerm+lonTerm)
    		dist = 2.*r*asin(dAngle)
    
    		;;-----------------------------------------------------------------  
	
    		rad_search = 500.0 ; radius we search outward to (km)
    		AC_count = where((dist.lt.rad_search).and.(count_track.eq.0),1,0) 
    		count_track = count_track + AC_count
    		count_total = count_total + AC_count
    	end if
    end do
  end do ; end iTimeTrack
end do ; end iTrack

printMinMax(count_total,True)


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

out_dir = "/keyserlab_rit/kbiernat/arctic_cyclone_tracks/scripts_JJA/track_density_v2/files/"
out_fil = out_dir+"track_density_ACs_climo_JJA_onlyDuringPeriods.nc"  ;-Output filename
var = "count" 

lat = domain_lat
lon = domain_lon

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

lon@units              = "degrees_east"
lon@long_name          = "longitude"
lon@grid_resolution    = "1.0_degrees"
lon@mapping            = "cylindrical_equidistant_projection_grid"
lon@coordinate_defines = "center"
lon@delta_x            = 1.0
lon@actual_range       = (/0,360/)

; File and Variable Attributes

fAtt = True

fAtt@creation_date  = systemfunc ("date")
fAtt@created_by     = "User: "+systemfunc ("whoami")
fAtt@description    = "Total AC count within specified radius of each grid point"


vAtt = 1
vAtt@long_name  = "Total AC Count"
vAtt@_FillValue = -10000000

;--------- 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 = (/ "lat", "lon" /)
dimSizes = (/dimsizes(lat), dimsizes(lon) /)
dimUnlim = (/False, False /)

chunks = (/dimsizes(lat), dimsizes(lon)  /)    ; 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, "lat", "float", "lat" )
filevardef( outFile, "lon", "float", "lon" )
filevardef( outFile, var, "integer", (/ "lat", "lon" /) )


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

; Write coordinates to record

outFile->lat = (/ lat /)
outFile->lon = (/ lon /)
outFile->$var$(:,:) = new( (/dimsizes(lat),dimsizes(lon)/), "integer")

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$(:,:) = (/count_total/)


end