#!/usr/bin/env python
# v0.22 4 July 2012 by Brian Fiedler
# This make the animated GIF seen at http://bunchofstuff.net/may24/blink.html
# ./nexrad.py -o refl.png -- KTLX_V03_20110524_203158.nc -180 -100 70 100
# ./nexrad.py -o radu.png -u -- KTLX_V03_20110524_203158.nc -180 -100 70 100
# convert -delay 300 -loop 0 refl.png radu.png blink.gif
# Some other options to explore
# ./nexrad.py -s 2 -k .125 -u -o mystorm.png -- KTLX_V03_20110524_203158.nc -180 -100 70 100
# Here is where the data came from:
# http://www.ncdc.noaa.gov/nexradinv/chooseday.jsp?id=ktlx
# "ordered" data, received email, then retrieved using
# wget -nH -r --reject "index.html*" -np http://ftp3.ncdc.noaa.gov/pub/has/HAS002419639/
# then used wct-viewer to export the level II data as .nc.  Only scipy.io can read
# the netcdf files, netcdf4 could not.  Apparently files exported by wct-viewer are netcdf3.
import sys,os
import numpy as np
from scipy.io import netcdf
import Image, ImageDraw, ImageFont
#### constants:
fontPath1 = "/usr/share/fonts/dejavu-lgc/DejaVuLGCSans-Bold.ttf" #redhat
fontPath2 = "/usr/share/fonts/truetype/ttf-dejavu/DejaVuSans-Bold.ttf" #ubuntu
if os.path.exists(fontPath1):
	fontPath=fontPath1
elif os.path.exists(fontPath2):
	fontPath=fontPath2
else:
	sys.exit('valid fonts not found')
outcolor="%c%c%c" % (30,30,30) # define a gray color for outside the range (meaning no data)
gatesperkm=4 # seems to be true for Level II data
relkeysize=.5 # makes the color key half the height of the image

# color for reflectivity:
#note: 0b or 1b will translate to -33 and -32.5, but should be interpreted as "no data"
noaacolors={ 75:(253,253,253), 70:(152,84,198), 65:(248,0,253), 60:(188,0,0),
55:(212,0,0), 50:(253,0,0), 45:(253,149,0), 40:(229,188,0),
35:(253,248,2), 30:(0,142,0), 25:(1,197,1), 20:(2,253,2),
15:(3,0,244), 10:(1,159,244), 5:(4,233,231), 0:(100,100,100),
-5:(153,153,102), -10:(204,204,153), -15:(102,51,102), -20:(153,102,153),
-25:(204,153,204), -30:(204,255,255), -35:(0,0,0)}
# colors for radial velocity component:
velcolors={ 55:(255,255,0), 45:(255,180,0), 35:(255,90,0), 25:(255,0,0), 15:(180,0,0), 5:(90,0,0),
 -5:(40,40,40), -15:(0,90,0), -25:(0,160,0), -35:(0,225,0), -45:(0,255,0), -55:(0,225,255), -65:(20,20,20)}

#### process command line arguments:
from optparse import OptionParser
usage="usage: %prog [options] -- infilename  [xll yll xur yur (in km)] "
parser = OptionParser(usage=usage, version="%prog 0.1")
parser.add_option("-o", dest="outputfile",type="string",help="output filename",default='myplot.png')
parser.add_option("-r", dest="reso",type="int",help="1 = 1 deg, 2=.5 deg",default=2)
parser.add_option("-s", dest="scan",type="int",help="scan (0=base)",default=0)
parser.add_option("-k", dest="kmstep",type="float",help="pixel size in km",default=0.25)
parser.add_option("-u", dest="plotu",action="store_true",help="plot U")
(opt, arg) = parser.parse_args()
try:
	infile=arg[0]
	print "will read", infile
except:
	sys.exit("quitting, need to specify file")

try:
	xll,yll,xur,yur=[float(x) for x in arg[1:]]
	plotall=False
	print "will plot from (km):",xll,yll,xur,yur
except:
	plotall=True
	print "will plot all data"

#### read data:
print "reading data file..."
ncfile = netcdf.netcdf_file(arg[0], 'r')

# This program was first designed for reflectivity data. But the variable name "refb", "refc" 
# and "refci" are now used for either reflectivity or radial velocity.
if opt.reso==1:
	if opt.plotu:
		refb=ncfile.variables['RadialVelocity'][opt.scan,:,:].view(dtype=np.uint8) # byte data
		offset=float(ncfile.variables['RadialVelocity'].add_offset)
		scale=float(ncfile.variables['RadialVelocity'].scale_factor)
		azims=ncfile.variables['azimuthV'][opt.scan,:]
		elevs=ncfile.variables['elevationV'][opt.scan,:]
		dists=ncfile.variables['distanceV'][:]
	else:
		refb=ncfile.variables['Reflectivity'][opt.scan,:,:].view(dtype=np.uint8) # byte data
		offset=float(ncfile.variables['Reflectivity'].add_offset)
		scale=float(ncfile.variables['Reflectivity'].scale_factor)
		azims=ncfile.variables['azimuthR'][opt.scan,:]
		elevs=ncfile.variables['elevationR'][opt.scan,:]
		dists=ncfile.variables['distanceR'][:]
elif opt.reso==2:
	if opt.plotu:
		refb=ncfile.variables['RadialVelocity_HI'][opt.scan,:,:].view(dtype=np.uint8) # byte data
		offset=float(ncfile.variables['RadialVelocity_HI'].add_offset)
		scale=float(ncfile.variables['RadialVelocity_HI'].scale_factor)
		azims=ncfile.variables['azimuthV_HI'][opt.scan,:]
		elevs=ncfile.variables['elevationV_HI'][opt.scan,:]
		dists=ncfile.variables['distanceV_HI'][:]
	else:
		refb=ncfile.variables['Reflectivity_HI'][opt.scan,:,:].view(dtype=np.uint8) # byte data
		offset=float(ncfile.variables['Reflectivity_HI'].add_offset)
		scale=float(ncfile.variables['Reflectivity_HI'].scale_factor)
		azims=ncfile.variables['azimuthR_HI'][opt.scan,:]
		elevs=ncfile.variables['elevationR_HI'][opt.scan,:]
		dists=ncfile.variables['distanceR_HI'][:]
else:
	sys.exit('reso must be 1 or 2')

print "data is stored in refb, with shape", refb.shape

### make colors:
print "making colors..."
if opt.plotu:
	colors=velcolors
else:
	colors=noaacolors
levels=colors.keys() # each item is the LOWER bound to the interval for that color
levels.sort() # sorts levels from low to high
rgb=[]
for level in levels:
	c3="%c%c%c" % colors[level] # convert each integer to a byte using %c (most of such characters are not printable)
	rgb.append(c3) #store the color as 3 bytes
argb=np.array(rgb) #an array of the 3-byte character strings of the RGB color for the ~23 color levels    

#### process data, meaning turn the byte arrays into arrays of RGB colors. 
nradial,ngate=refb.shape  #the shape of the data array
print 'nradial=',nradial,'   ngate=',ngate

if plotall: # adjust km boundary to fit all data, other uses xll, yll,xur,yur from command line
	xll=-ngate/4-1
	yll=-ngate/4-1
	xur=ngate/4+1
	yur=ngate/4+1
# Make reflectivities from bytes:
print " multiply byte by scale=",scale,"  then add offset=",offset
ref=refb*scale+offset # turn the byte array into standard reflectivity values in dBZ

firstazim=float(azims[0])
# note that elevation angles are not constant:
# print "elevation able array, in degrees:", elevs.shape , elevs
avgelev=np.average(elevs) # average of all the elevation angles of the scan
print "average elevation angle:",avgelev
firstdist=float(dists[0])/1000.

# next "bin the data", converting the floating point number to a non-negative integer 
print "converting reflectivity to color index..."
if opt.plotu:
	refci=(ref+65)/10
else:
	refci=ref/5+7 # reflectivity color index
# the above has made refci, still in floating point format
refci=refci.astype(int) # convert refci to type integer
# Now refci is the data array of integers, with the integer denoting the color value that will
# be applied in a plot.

# Next make refc, which will be the same shape as refci, with the integer from refci extracting
# the RGB color bytes we put in argb. 
#  We use a very powerful feature of numpy overloading the array slicing syntax, the usual use of the []
#  but here we "feed" argb the data as an array of interger color indices, the output is the data as an array of colors
refc=argb[refci.flatten()] 
# Note the above statement only works with flattened arrays, arrays with the shape information removed.
# optional lines for debugging:
#amg = Image.fromstring('RGB', (ngate,nradial) , refc.tostring())
#amg.save('arrayrefl.png')
#del amg

### Plot the data in polar coordinates  
print "plotting in polar coordinates..."
# make a Cartesian grid, centered on the radar:
xa,ya=np.meshgrid( np.arange(xll+.5*opt.kmstep,xur,opt.kmstep) , np.arange(yur-.5*opt.kmstep,yll,-opt.kmstep) )
# next associate a radial index (azimuthal angle index) and gate index to each point in our Cartesian array:
tn=np.arctan2(xa,ya)*180./3.14159265  # array of the meteorology compass angle of grid point (0=North, 90=East, ...)
ra=np.sqrt(xa*xa+ya*ya) # array of distance of grid point from the radar
tn=(tn-firstazim)%360. # convert compass angle of grid point to angle from first data point in scan
ita=(opt.reso*tn).astype('int') # the integer azimuth index for the grid point (i.e., which radial the grid point is on)
rra=(ra-firstdist)*gatesperkm # convert distance from first "ring of data" to the gate index for the grid point 
ira=rra.astype('int') # make rra an integer
# our Cartesian array now has integers, say ("angle index","gate index") as stored with the arrays ita and ira


## Let me show you something by example in the python interpreter.
## Suppose refc had only 2 radials and 3 gates, and the items were one-byte letters, rather than 3-byte RGB:
# >>> refc
# array([['a', 'b', 'c'],
#        ['d', 'e', 'f']], 
#       dtype='|S1')
## Next we show this example how to find the data value using the radial and gate index, in the orginally 
## shaped array, and in the flattened array (here named 'flat')
# >>> flat=refc.flatten()
# >>> flat
# array(['a', 'b', 'c', 'd', 'e', 'f'], 
#       dtype='|S1')
# >>> refc[1,2]
# 'f'
# >>> flat[1*3+2]
# 'f'

# To find the data in the flattend array we take the "angle index" times the length of a data row
# (which is ngate) plus the "gate index" to find the integer where we obtain the data value in the 
# flattend refc array.  The array of such integers is here called fta:
fta=(ita*ngate+ira).flatten()

lbd=refc.size # max valid index is lbd-1, integer values in fta that are greater than lbd are outside the range of valid data
# we append one more color to the array of data colors, with the intent to grab that color when needed as the outcolor
verysmallarray=np.array(outcolor).reshape(1) # reshpae(1) is tricky,  otherwise shape=(), which makes concatenation impossible 
refc=np.concatenate((refc,verysmallarray)) # the outcolor is now appended to the data array

# numpy arrays can accept a logical array of the same shape within the []. Very convenient!
fta[ira.flatten()>=ngate] =lbd # If the Cartesian array had "gate index" >=ngate, then associate it with the outcolor
fta[rra.flatten()<0 ]=lbd # if the Cartisian grid point was inside the first data ring, then associate it with outcolor 

# Finally assign colors to our Cartesian array of grid points
print "now do the the amazing numpy step..."
# CartColorized will be the same shape as fta, with the integer from fta extracting the color bytes from refc
CartColorized=refc[fta] # a short but powerful numpy statement! The Cartesian grid is now colorized


# We could output CartColorized as a ppm, but instead write it as a PNG, with annotations,
# using the Python Imaging Library (PIL) module.
print "making PNG image..."
yras,xras=xa.shape # xa, ya, ra, are all of the shape of the unflattened CartColorized array...
img = Image.fromstring('RGB', (xras,yras) , CartColorized.tostring())

# delay saving image as PNG file until other annotations are added:

d=int(relkeysize*yras/len(rgb))
efontsize=int(.8*d) #estimate font size for color key
fontsize=efontsize #only fontsize >=16 or =14 or =10 don't clip descenders
if efontsize<16: fontsize=14
if efontsize<13: fontsize=10
print "keysquare size=", d,"   fontsize=",fontsize

# In the following, [::-1] reverses the array, so the legend ppm will have high values at the top:
kmg = Image.fromstring('RGB', (1,len(rgb)) , argb[::-1].tostring() ).resize( (d,d*len(rgb)) )
keytop=yras-d*(len(rgb)+1)- 40
img.paste(kmg,(xras-2*d,keytop)) # paste color key onto image

### make text annotations:
myfont =  ImageFont.truetype ( fontPath, fontsize )
draw  =  ImageDraw.Draw ( img )

# label the color key with numbers and units:
for i in range(len(levels))[1:]:
	lev= "%3d" % (levels[i]) # number label
	draw.text ( (xras-4*d,yras-40-(2+i)*d+d/2), lev, font=myfont, fill="white" )
if opt.plotu:
	draw.text ( (xras-2*d,keytop-d), "m/s", font=myfont, fill="white" )
else:
	draw.text ( (xras-2*d,keytop-d), "dBZ", font=myfont, fill="white" )


myfont =  ImageFont.truetype ( fontPath, 10 ) # reduce font size to 10 for other text

# plot length scale key:
draw.line( [(xras-105,yras-25),(xras-5,yras-25)], fill='white',width=14)
scalelab= '%4.1f km' % (100*opt.kmstep) # length scale
draw.text( (xras-65,yras-30), scalelab , font=myfont, fill="black" )

# print the options and arguments used to make the plot:
scaninfo ="avg elev = %5.2f deg, -s %d -r %d -k %5.3f, %d %d %d %d" % (
avgelev, opt.scan, opt.reso, opt.kmstep, xll,yll,xur,yur) 
draw.text( (5,yras-15), scaninfo , font=myfont, fill="white" )

# write file name on image:
filename=arg[0].split('/')[-1]
labsize = myfont.getsize(filename)
draw.text( (xras-labsize[0]-5,yras-15), filename , font=myfont, fill="white" )

# save output file:
print "saving",opt.outputfile
img.save(opt.outputfile)

print "Done! To see, do: eog "+opt.outputfile