#!/usr/bin/FEWXPYTHON

########### imports in order of appearance
from datetime import datetime			# for date time of now
NOWCIVIL=datetime.now()				# -- the actual civil time of NOW
import math                                     # for pi
import glob					# for grabbing list of files of certain naming pattern
import subprocess                               # for using system commands to get top two lines of the raw radar files (contain id and date/time)
import sys					# for use to redirect non-warning messages due to pyart import
import os                                       # for use to redirect non-warning messages due to pyart import 
orig_stdout=sys.stdout				# -- remember present state of stdout -- part of approach to supress non-warning messages due to pyart import
sys.stdout=open(os.devnull,'w')			# -- redirect stdout to null
import pyart					# for creating radar composite
sys.stdout.close()				# -- close redirect of stdout to null
sys.stdout=orig_stdout				# -- restore state of stdout to what it was before pyart import
import numpy as np				# for performing array math on the composite results
from scipy.interpolate import griddata		# for remmapping composite to lat lon grid
import xarray as xr				# for storing remapped array results into xarray data set so it can be written in grib2 format
from cfgrib.xarray_to_grib import to_grib	# for writing the xarray data set to a grib2 file


########### user settings
compositeres=2.0			# km
too_old=40                      	# minutes a radar site's data is considered too old to be current (anything older than this is excluded from the composite)
NANVAL=-15.0                    	# what to set nan values in composite routine (0.0 not great, anything less than 0 works -- the more 'less than zero' the slower the routine, so -15.0 was chosen)
GRIBFILE="natradrawcomposite.grb2"	# name of the output grib file


########## constants
nowutc = datetime.utcnow()      # UTC time of "now" to help determine how old a rad file might be in minutes
Rearth=6378137                  # radius of earth in meters
pi=math.pi                      # pi 3.1415.....
degPerMeter=360/(2*pi*Rearth)   # degrees lon per meter


########### functions
def EXITING(): exit()


### obtain list of rad data files in $LATEST/rad
raddir="/fewxops/data/latest/rad"
print(str(datetime.now()) + " - reading radar files from "+raddir+" (time stamps will be compared to current UTC time of "+str(nowutc)+")")
nexrad3files=glob.glob(raddir+"/sn.last.*")
if (len(nexrad3files)) == 0: 
	print ("No sn.last.* files found in "+ raddir)
	EXITING()


### create master list of both avairads (by id) and a master array of available radar sites -- each element of the master array holds all data per radar site (gets fed into pyart to generate composite)
availrads=[]             # array of radar ids that arent too old to use in the composite
master_array_radsites=[] # an array of each radars data -- this gets fed into the compostie routine
decluttereddbz={}        # a corresponding tuple to the master_array_radsites array of radar data that has been filtered -- this gets fed into the composite routine
key=0                    # an index that allows you to reference a specific element within the tuple

for radfile in nexrad3files: 
	try: 
		radardata=pyart.io.read_nexrad_level3(radfile)
	except:
		print(str(datetime.now())+" -   "+radfile+"being skipped -- seems corrupt")
		continue

	radtime=datetime.strptime(radardata.time['units'][14:34].replace("T"," ").replace("Z",""),"%Y-%m-%d %H:%M:%S")
	radid=subprocess.run(["sed -n 2p "+radfile+" | cut -c 4-6"],shell=True,capture_output=True,text=True).stdout.rstrip()
	minutesold=(nowutc-radtime).total_seconds()/60

	if minutesold <= too_old:
		master_array_radsites.append(radardata)
		availrads.append(radid)
		# declutter the rad field -- via source code comments in https://arm-doe.github.io/pyart/_modules/pyart/correct/despeckle.html#despeckle_field
		# the "Size" setting is a threshold number of clustered gates, below this threshold is considered noise and filtered -- we tested 1 9 25 (grid pts) and 9 seemed to be the best
		# -- why 1, 9, 25?  because likening a gate to a grid pt, at the gate of interest you have itself, or itself plus the 8 bounding gates, or itself plus the 8 bounding gates plus 16 bounding the 8
		decluttereddbz[key]=pyart.correct.despeckle_field(radardata, 'reflectivity', size=9)
		key += 1
	else:	
		print(str(datetime.now())+" -   "+radid+"  UTC "+str(radtime)+" - ignored - older than "+str(too_old)+" minutes ("+str(int(minutesold))+" minutes old)")

if len(availrads) == 0: 
	print ("There were no radar files with a time stamp less than "+str(too_old)+" minutes old")
	EXITING()

print(str(datetime.now()) + " -   "+str(len(availrads))+" radar sites will be composited:",*sorted(availrads))


##### prelim critical vars for pyart composite
print(str(datetime.now()) + " - generating "+str(compositeres)+"km res composite (at 1km res can 2+ min)")
US_centlat=38; US_centlon=-98		# US composite central lat lon
US_ew_meters = 4736                     # domain distance in km east/west across CONUS 
US_ns_meters = 3000                     # domain distance in km north/south CONUS
i=int((US_ew_meters/ compositeres) + 1) # pyart composite grid pts in x dir given US e/w domain distance and composite res
j=int((US_ns_meters/ compositeres) + 1) # pyart composite grid pts in y dir given US n/s domain distance and composite res
z=4500					# max height above ground in m to use from vol scans - for reference the highest 88d site is SLC at 3231 m -- so 1000m of the ground would be 4231 m there
vlevels=int(z/500)+1                  	# vlevels to sift through to make sure we arent getting cones of silence 500m was arbitrary lyr thickness 
#### generate composite -- result is a equal area cartesion grid - settings and more found in SOURCE CODE COMMENTS https://arm-doe.github.io/pyart/_modules/pyart/map/grid_mapper.html#grid_from_radars
pyart_compositegrid = pyart.map.grid_from_radars(  									# ingest master array of available radar arrays 
	(master_array_radsites), 											# array of radar file contents
	grid_shape=(vlevels, j, i),                                              					# desired number of pts in z, y, x for use domain
	grid_origin=(US_centlat,US_centlon),										# grid origin
	grid_limits=((1,z), (-1000*US_ns_meters/2, 1000*US_ns_meters/2), (- 1000*US_ew_meters/2, 1000*US_ew_meters/2)),	# relative to grid orgin, value is meters of the cube of the composite -- 1 to 4000m z, x -15000x
	gatefilters=(list(decluttereddbz.values())),                                                                    # insert here the decluttered object from above 
	min_radius=2000.0,                                                                                              # smaller than 1750 gives you holes, 2000 seemed best visually 
	weighting_function='barnes',                                                                                    # barnes faster, seemed better than cressman -- waaay better than nearest. cressman slowest
	)


######## isolate dbz lat lon arrays
# ideally at any point we should be using the reflectivity from the radar whose .5 deg elev slice is 1km or less agl -- and if multiple radars qualify, take the max value 
# but for now though we will revert to a poor mans method of just take the highest value at any point 
# this will minimize cones of silence by getting max dbz at each gate fm all vlevels (goes from 3d array to 2d, first axis is vlevels) AND set all nans to NANVAL since that what pyart gridding does 
# tip you can look at each level of pyart_composite.fields['reflectivity']['data'][level] but the .data looks at them all in bulk 
print(str(datetime.now()) + " -   isolating dbz data from composite")
dbz=np.max(np.nan_to_num(pyart_compositegrid.fields['reflectivity']['data'].data,nan=NANVAL),axis=0)


######## convert pyart composite grid to standard latlon via scipy prep and execution
print(str(datetime.now()) + " -   prepping dbz lat and lon data for remapping")
dbzlon,dbzlat=pyart_compositegrid.get_point_longitude_latitude(level=0)  # get_point holds the respective lat lon of the composite grid
# make 1 d arrays for scipy -- thin points by retaining only those > NANVAL.. and pair up/and transpose the lonlats into one array
srclon=dbzlon.flatten(); srclat=dbzlat.flatten(); srcdat=dbz.flatten()
srclat=srclat[srcdat>NANVAL]
srclon=srclon[srcdat>NANVAL]
srcdat=srcdat[srcdat>NANVAL]
srclonlatpaired=np.array([srclon,srclat]).T


# establish target grid 
print(str(datetime.now()) + " -   creating target lat lon grid for remapping")
minlat=np.min(srclat);maxlat=np.max(srclat);minlon=np.min(srclon);maxlon=np.max(srclon)
refl_xcoords = pyart_compositegrid.x['data']; refl_ycoords = pyart_compositegrid.y['data']  # these are the cartesian points of the composite relative to the center of the grid - not lat lons
pyart_compositegrid_xres = (((-1) * refl_xcoords[0]) - ((-1) * refl_xcoords[1])) / 1000  # should be teh same as compres 
pyart_compositegrid_yres = (((-1) * refl_ycoords[0]) - ((-1) * refl_ycoords[1])) / 1000  # because its a cartesian grid pyart produces this should be the same as xres too
targetgridlonlist=np.arange(minlon,maxlon+pyart_compositegrid_xres*1000*degPerMeter,pyart_compositegrid_xres*1000*degPerMeter)
targetgridlatlist=np.arange(minlat,maxlat+pyart_compositegrid_yres*1000*degPerMeter,pyart_compositegrid_yres*1000*degPerMeter)
targetgridlon,targetgridlat=np.meshgrid(targetgridlonlist,targetgridlatlist)
targetgridlonlatpaired=np.array([targetgridlon.flatten(),targetgridlat.flatten()]).T


### scipy remapping 
print(str(datetime.now()) + " - mapping source data to target "+str(compositeres)+"km lat lon grid (at 1km res can take 4+ min)")
latlongrid=griddata(srclonlatpaired,srcdat,targetgridlonlatpaired,method="linear")


# reshape for writing to grib2
latlongrid_reshaped=latlongrid.reshape(targetgridlon.shape)


#### write to grib2
print(str(datetime.now()) + " - writing to grib2 file "+GRIBFILE)
# create xarray array
radcomp=xr.DataArray(data=latlongrid_reshaped,dims=['latitude','longitude'],coords={'latitude':targetgridlatlist,'longitude':targetgridlonlist},)
# create xarray data set 
xwrite=radcomp.to_dataset(name="result") # name is arbitrary, literally doesnt matter
# write xarray data set to GRIBFILE
to_grib(xwrite,GRIBFILE,grib_keys={'edition':2},no_warn=True)


### exit
print(str(datetime.now())+ " - done - total time elapsed "+str(datetime.now()-NOWCIVIL) )
exit()
