'''
snow_cover_animations.py, true to its name, creates an animated GIF file showing snow
cover within a series of wrfout files.  To do this, it first writes an ncl script
with an input for the title at the top of each of the GIF frames.  The ncl script is
then run, which retrieves SNOWH (or snow depth) data from each wrfout file and creates
plots with the snow data shown in the subject area.  Then this script enters an
ImageMagick command to put each frame together into an animation and deletes all
unnecessary files that remain, including the ncl script.

Note: There are factors here which may be changed.  The first of which is the contour
levels in the snow cover plots.  Currently, they span 0-2 inches.  This can be changed
at line 79 following "opts@cnLevels".  The second thing you may wish to change is the
speed of the GIF which can be changed at line 94 following "-delay ".

Ryan Clare
November, 2016
University of Wisconsin-Madison
rclare2@wisc.edu
'''

from commands import *						#Library for executing Linux commands
from numpy import *							#Library for math

fig_title = raw_input("Figure title: ")		#Figure title input

gif_title = raw_input("GIF title: ")		#GIF title input

#NCL code for snow.ncl in two parts
ncl_content1 = '\
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"					\n\
load "$NCARG_ROOT/lib/ncarg/nclscripts/wrf/WRFUserARW.ncl"					\n\
begin																		\n\
  fList = systemfunc("ls -1 Case00*")										\n\
  nFiles = dimsizes(fList) 													\n\
  do iFile = 0, nFiles - 1													\n\
    filename = sprinti("temp_snow%03d",iFile)								\n'
ncl_content2 = '\
	wks = gsn_open_wks("png",filename) 										\n\
    gsn_define_colormap(wks,"WhiteBlue")									\n\
    a = addfile(fList(iFile),"r")											\n\
    snow = wrf_user_getvar(a,"SNOWH",-1)									\n\
    slp = wrf_user_getvar(a,"slp",-1)										\n\
    snow = snow * 39.3701 													\n\
    snow@units = "in"														\n\
    res = True																\n\
    res@MainTitle = "%s"													\n\
    res@gsnDraw = False														\n\
    res@gsnFrame = False 													\n\
    res@cnFillOn = True 													\n\
    res@cnLinesOn = False 													\n\
    res@lbLabelAutoStride = True 											\n\
    times = wrf_user_list_times(a) 											\n\
    res@ValidTime	= True 													\n\
    res@TimeLabel	= times 												\n\
    pltres	= True 															\n\
    mpres	= True 															\n\
    mpres@mpGeophysicalLineColor = "Black" 									\n\
    mpres@mpNationalLineColor = "Black" 									\n\
    mpres@mpUSStateLineColor = "Black" 										\n\
    mpres@mpGridLineColor = "Black" 										\n\
    mpres@mpLimbLineColor = "Black" 										\n\
    mpres@mpPerimLineColor = "Black" 										\n\
    mpres@mpGeophysicalLineThicknessF	= 2.0 								\n\
    mpresmpGridAndLimbOn	= False 										\n\
    mpres@mpNationalLineThicknessF = 2.0 									\n\
    mpres@mpUSStateLineThicknessF = 2.0 									\n\
    opts = res 																\n\
    opts@cnFillMode = "RasterFill" 											\n\
    opts@cnRasterSmoothingOn = True 										\n\
    pressureres = True 														\n\
    pressureres@cnLineColor = "Black" 										\n\
    pressureres@cnInfoLabelOn	= False 									\n\
    pressureres@cnMinLevelValF	= 800 										\n\
    pressureres@cnMaxLevelValF = 1200 										\n\
    pressureres@cnLevelSpacingF	= 4 										\n\
    pressureres@cnLineLabelPlacementMode = "Constant" 						\n\
    ;Explicit contours														\n\
    opts@cnLevelSelectionMode = "ExplicitLevels"							\n\
    opts@cnLevels    = (/ 0, .2, .4, .6, .8, 1., 2./)						\n\
    contour_snow	= wrf_contour(a,wks,snow(0,:,:),opts)					\n\
    contour_pres	= wrf_contour(a,wks,slp(0,:,:),pressureres)				\n\
    plot1 = wrf_map_overlays(a,wks, (/contour_snow,contour_pres/), pltres, mpres)	\n\
  end do																	\n\
end' % fig_title

ncl_content = ncl_content1 + ncl_content2	#Combine the two parts

f = open("snow.ncl",'w')					#Creates local temporary ncl script
f.write(ncl_content)						#Writes above contents into file
f.close()

getoutput('ncl snow.ncl')					#Runs temporary ncl script

getoutput('convert -loop 0 -delay 90 temp_snow* %s.gif' % gif_title)	#Runs ImageMagick

getoutput('rm temp_snow*; rm snow.ncl')		#Removes extra files