#!/usr/bin/python
#
# slicer.py
# 	-- reads slice files, converts to netcdf, and plots them
#
# 31 January 2008
###############################################################################

import os.path
import sys
import gc
import numpy as N
from pylab import *
from matplotlib.toolkits.basemap import cm

try:
	afile = sys.argv[1]
	s = open(afile,'r')
except:
	print 'Using default file: slicer.dat'
	s = open('slicer.dat','r')

times = []
data = {}
cnt = 0
dpi = 200
res = 20		# resolution in meters
xyl  = [129]	# x,y location
zloc = [0]		# z locations (filled later)
npts = 100		# max height including non-physical points (meters)
end = '.grd'
zht = [res*zloc[0]]
gpath = 'plots/'
gtype = '.png'


vars = { 'u' : r'$\it{x}$-component velocity (m s$^{-1}$)',
		 'v' : r'$\it{y}$-component velocity (m s$^{-1}$)',
		 'w' : r'$\it{z}$-component velocity (m s$^{-1}$)',
		 'T' : 'Potential temperature (K)',
		 'q' : 'Specific humidity (g/g)',
		 'E' : r'Subgrid kinetic energy (m$^2$ s$^{-2}$)' }

xlabs = { 'X' : r'$\it{y}$ (m)',
		  'Y' : r'$\it{x}$ (m)',
		  'Z' : r'$\it{x}$ (m)' }

ylabs = { 'X' : r'$\it{z}$ (m)',
		  'Y' : r'$\it{z}$ (m)',
		  'Z' : r'$\it{y}$ (m)' }
		 
splts = { 'u' : (3,2,1),
		  'v' : (3,2,2),
		  'w' : (3,2,5),
		  'T' : (3,2,3),
		  'q' : (3,2,4),
		  'E' : (3,2,6) }

clims = { 'u' : [-5,20],
		  'v' : [-5,5],
		  'w' : [-5,5],
		  'T' : [294,300],
		  'q' : [4.5e-3,6e-3],
		  'E' : [0,0.5] }

cuts = { 'X' : 'x: ',
		 'Y' : 'y: ',
		 'Z' : 'z: ' }

locs = { 'X' : xyl,
		 'Y' : xyl,
		 'Z' : zloc }


for line in s :
	tmp = line.split()
	times.append( [int(tmp[0]), float(tmp[1]), float(tmp[2])] )
	cnt += 1

F = figure(1)

for i in range(len(times)) :
	step = times[i][0]
	zloc[0] = int(round( times[i][2] * npts ))
	zht[0]  = res * zloc[0]
	
	for cut in sorted( cuts.keys() ) :

		if cut == 'X' or cut == 'Y' :
			F.set_size_inches(16.5,10.5)
		elif cut == 'Z' :
			F.set_size_inches(13.5,16.5)
		
		for var in sorted( vars.keys() ) :

			loc = locs[cut][0]
			
			fname = '%s/%s%s%03d%07d%s' % (var, var, cut, loc, step, end)

			if os.path.exists(fname) :
				print 'Processing: ', fname,
				tmp = []
		
				f = open(fname,'r')
				f.readline()
				rows,cols = map(int, f.readline().split())
				minDim1,maxDim1 = map(float, f.readline().split())
				minDim2,maxDim2 = map(float, f.readline().split())
				minVal, maxVal  = map(float, f.readline().split())
						
				print '-- rows: %03d   cols: %03d   minVal: %5.2f   maxVal: %5.2f' % (rows,cols,minVal,maxVal)
		
				for line in f :
					tmp.append( float(line) )
		
				tmp = N.array(tmp,order='C')
				tmp = N.reshape(tmp, (cols,rows))
				data[var] = tmp
				x = N.arange( minDim1+res, maxDim1, res, dtype=int )
				y = N.arange( minDim2+res, maxDim2, res, dtype=int )
		
				subplot( splts[var][0], splts[var][1], splts[var][2] )
				pcolormesh( x,y,data[var][1:-1, 1:-1], cmap=cm.GMT_polar )
				clim( clims[var][0], clims[var][1] )

				tck = clims[var][1] + (clims[var][1]-clims[var][0])/5.
				colorbar( ticks=arange(clims[var][0],tck,(clims[var][1]-clims[var][0])/5.) )

				axis('image')
				title( vars[var] )
				xlabel( xlabs[cut] )
				ylabel( ylabs[cut] )
		
		gname = '%s_%03d_%07d%s' % (cut, loc, step, gtype)				
	
		if os.path.exists(fname) :
			ftext = 'Time: %5.2f s -- Step: %d' % (times[i][1],step)
			F.text( 0.5, 0.95, ftext, horizontalalignment='center' )
			F.savefig( gpath + gname, dpi=dpi )
			F.clear()
			gc.collect()