#!/bin/env python

# This script was wrote using anconda3 as of date 20 March 2020
# py-art 1.12dev https://arm-doe.github.io/pyart/#
#

import matplotlib.pyplot as plt
import pyart
import glob
from netCDF4 import Dataset
import numpy as np
import os
import math
import sys
import signal
import time

start_time = time.time()
# Velocity display is set to nyquist value
#
# dBz levels min to max, the MMR uses 0 as _FillValue 
zmin, zmax = 0, 60
maskoutside = "False"
# Plot size
xsize, ysize = 6, 4

def plotdata(i):
	# read the data and create the display object
	radar = pyart.io.read_cfradial(i)
	display = pyart.graph.RadarDisplay(radar)
	fh = Dataset(i)
	nyquist = radar.get_nyquist_vel(0)
	#nyquist = ((fh.variables['num_doppler_cells'][0] * fh.variables['doppler_cell_size'][0]) / 2)
	#fields to plot and ranges
	#radar.info(level='c')  # c = compact, s = standard, f = full
	dist = round(fh.variables['range'][-1]/1000)
	nm_dist = round(dist/1.852) # Convert to NM
	print(i, ": Max Range:", nm_dist,"NM,", dist, "KM,", "Nyquist Velocity:",nyquist, "M/S")
	fields_to_plot = ['DBZ', 'VEL']
	ranges = [(zmin, zmax), (nyquist*-1, nyquist)]
	colormaps = ['pyart_NWSRef', 'pyart_NWSVel']
	# plot the data
	nplots = len(fields_to_plot)
	plt.figure(figsize=[xsize * nplots, ysize])
	# plot each field
	for plot_num in range(nplots):
		field = fields_to_plot[plot_num]
		vmin, vmax = ranges[plot_num]
		colormap = colormaps[plot_num]
		plt.subplot(1, nplots, plot_num + 1)
		display.plot(field, 0, vmin=vmin, vmax=vmax,cmap=colormap, title_flag=True, mask_outside=maskoutside)
		display.set_limits(ylim=[dist*-1, dist], xlim=[dist*-1, dist])
	plt.show()

def saveimage(i):
	radar = pyart.io.read_cfradial(i)
	display = pyart.graph.RadarDisplay(radar)
	fh = Dataset(i)
	# fields to plot and ranges
	nyquist = radar.get_nyquist_vel(0)
	#nyquist = ((fh.variables['num_doppler_cells'][0] * fh.variables['doppler_cell_size'][0]) / 2)
	dist = round(fh.variables['range'][-1]/1000)
	nm_dist = round(dist/1.852) # Convert to NM
	fields_to_plot = ['DBZ', 'VEL']
	print(i, ": Max Range:", nm_dist,"NM,", dist, "KM,", "Nyquist Velocity:",nyquist, "M/S")
	ranges = [(zmin, zmax), (nyquist*-1, nyquist)]
	colormaps = ['pyart_NWSRef', 'pyart_NWSVel']
	# plot the data
	nplots = len(fields_to_plot)
	plt.figure(figsize=[xsize * nplots, ysize])
	# plot each field
	for plot_num in range(nplots):
		field = fields_to_plot[plot_num]
		vmin, vmax = ranges[plot_num]
		colormap = colormaps[plot_num]
		plt.subplot(1, nplots, plot_num + 1)
		display.plot(field, 0, vmin=vmin, vmax=vmax,cmap=colormap, title_flag=True, mask_outside=maskoutside)
		display.set_limits(ylim=[dist*-1, dist], xlim=[dist*-1, dist])
	plt.savefig(os.path.splitext(i)[0]+'.png', pad_inches=0.0, transparent=False)
	plt.close()

mylist = [f for f in glob.glob("HWX*.nc")]

try:
	var = sys.argv[1]
	if 'plot' in var:
		for i in mylist:
			try:
				plotdata(i)
			except KeyboardInterrupt:
				exit()
			except:
				pass
	if 'save' in var:
		for i in mylist:
			try:
				saveimage(i)
				#print("Elapsed Time:", time.time() - start_time)
			except KeyboardInterrupt:
				exit()
			except:
				pass
	else:
		print("Please enter: 2-plot-mmr.py save to create png's")
		print("or 2-plot-mmr.py plot to view the data as an image")
		exit() 
except IndexError as e:
	#print(e)
	print("Please enter: 2-plot-mmr.py save or 2-plot-mmr.py plot") 
	exit()