#!/usr/bin/env python

'''Plot rectangular grid of values'''

# Sky Coyote, June 2008

import sys, signal
from numpy import *
import wx

class ImagePlot(wx.Frame):
	'''Frame containing 2d array image'''
	def __init__(self, data, mask, title='Image plot', min=0, max=0, size=(512, 512), pos=(-1, -1), mode=1):
		'''Initialize frame'''
		# create bitmap
		if mode == 1:
			# in color
			self.bmp = self.array2bitmap2(data, mask, min, max)
		else:
			# in b & w
			self.bmp = self.array2bitmap(data, mask, min, max)
		# create frame
		wx.Frame.__init__(self, parent=None, id=-1, title=title, size=size, pos=pos)
		self.offset = array((0.0, 0.0))
		self.ctrlKeyDown = False
		self.leftDown = False
		# bind event handlers
		self.Bind(wx.EVT_KEY_DOWN, self.OnKeyDown)
		self.Bind(wx.EVT_KEY_UP, self.OnKeyUp)
		self.Bind(wx.EVT_PAINT, self.OnPaint)
		self.Bind(wx.EVT_LEFT_DOWN, self.OnLeftDown)
		self.Bind(wx.EVT_LEFT_UP, self.OnLeftUp)
		self.Bind(wx.EVT_MOTION, self.OnMotion)
		  
	def OnPaint(self, event):
		'''Draw image plot'''
		dc = wx.PaintDC(self)
		dc.Clear()		
		dc.DrawBitmap(self.bmp, self.offset[0], self.offset[1])
			
	def OnKeyDown(self, evt):
		'''Handle key down'''
		keycode = evt.GetKeyCode()
		if keycode == 308:	# ctrl
			self.ctrlKeyDown = True
		elif keycode == 67:	# c
			if self.ctrlKeyDown: sys.exit(0)
			
	def OnKeyUp(self, evt):
		'''Handle key up'''
		keycode = evt.GetKeyCode()
		if keycode == 308:	# ctrl
			self.ctrlKeyDown = False
			
	def OnLeftDown(self, evt):
		'''Handle mouse down'''
		self.mousePosition0 = array(evt.GetPositionTuple())
		self.offset0 = self.offset
		self.leftDown = True
		self.CaptureMouse()
		  
	def OnLeftUp(self, evt):
		'''Handle mouse up'''
		if self.HasCapture():
			self.ReleaseMouse()
		self.leftDown = False
		  
	def OnMotion(self, evt):
		'''Handle mouse movement'''
		if evt.Dragging() and evt.LeftIsDown() and self.leftDown:
			# move bitmap
			self.mousePosition = array(evt.GetPositionTuple())
			self.offset = self.offset0 + (self.mousePosition - self.mousePosition0)
			self.OnPaint(None)
		evt.Skip()
		  
	def array2bitmap(self, data, mask, min=0, max=0):
		'''Convert numpy ndarray to wx.Bitmap'''
		# clip image
		if min == 0 and max == 0: 
			min = data[mask].min()
			max = data[mask].max()
		data = data.clip(min=min, max=max)
		if max == min:
			max += 1.0
			min -= 1.0
		data[mask != True] = min

		# Create color planes
		red = 1.0 * data
		green = 1.0 * data
		blue = 1.0 * data
			
		# Flip arrays
		red = red[::-1, :]
		green = green[::-1, :]
		blue = blue[::-1, :]
		
		# Scale and convert to bytes
		if max == min: 
			redBytes = (127.0 * ones(data.shape)).astype('b')
			greenBytes = (127.0 * ones(data.shape)).astype('b')
			blueBytes = (127.0 * ones(data.shape)).astype('b')
		else: 
			redBytes = (255.0 * (red - min) / (max - min)).astype('b')
			greenBytes = (255.0 * (green - min) / (max - min)).astype('b')
			blueBytes = (255.0 * (blue - min) / (max - min)).astype('b')

		# Create composite planes
		composite = concatenate((redBytes, greenBytes, blueBytes)).reshape(3, \
			data.shape[0], data.shape[1])
		# Shuffle and get raw interleaved data
		byteString = composite.swapaxes(0, 2).swapaxes(0, 1).tostring()

		# create wx.Bitmap from byte buffer
		return wx.BitmapFromBuffer(data.shape[1], data.shape[0], byteString)
		
	def hsv2rgb(self, h, s, v):
		'''Convert hue/saturation/value to red/green/blue'''
		# http://en.wikipedia.org/wiki/HSV_color_space
		hi = int(h / 60.0) % 6
		f = h / 60.0 - int(h / 60.0)
		p = v * (1.0 - s)
		q = v * (1.0 - f * s)
		t = v * (1.0 - (1.0 - f) * s)
		if hi == 0:
			r, g, b = v, t, p
		elif hi == 1:
			r, g, b = q, v, p
		elif hi == 2:
			r, g, b = p, v, t
		elif hi == 3:
			r, g, b = p, q, v
		elif hi == 4:
			r, g, b = t, p, v
		elif hi == 5:
			r, g, b = v, p, q
		return (r * 255, g * 255, b * 255)
	
	def array2bitmap2(self, data, mask, min=0, max=0):
		'''Convert numpy ndarray to wx.Bitmap'''
		# clip image
		if min == 0 and max == 0: 
			min = data[mask].min()
			max = data[mask].max()
		data = data.clip(min=min, max=max)
		if max == min:
			max += 1.0
			min -= 1.0
		data[mask != True] = min

		# Create color planes
		red = zeros(data.shape)
		green = zeros(data.shape)
		blue = zeros(data.shape)
		
		# convert value to hue
		# this is the slow part, as it is a nested loop
		dz = max - min
		for j in range(data.shape[0]):
			for i in range(data.shape[1]):
				if mask[j, i]:
					x = 270.0 - 360.0 * (data[j, i] - min) / dz
					x = x % 360.0
					red[j, i], green[j, i], blue[j, i] = self.hsv2rgb(x, 1.0, 1.0)
			
		# Flip arrays
		redBytes = red[::-1, :].astype('b')
		greenBytes = green[::-1, :].astype('b')
		blueBytes = blue[::-1, :].astype('b')
		
		# Create composite planes
		composite = concatenate((redBytes, greenBytes, blueBytes)).reshape(3, \
			data.shape[0], data.shape[1])
		# Shuffle and get raw interleaved data
		byteString = composite.swapaxes(0, 2).swapaxes(0, 1).tostring()

		# create wx.Bitmap from byte buffer
		return wx.BitmapFromBuffer(data.shape[1], data.shape[0], byteString)
		
def readData(fname):
	'''Read data file and header'''
	header = {}
	dataList = []
	dataFile = open(fname, 'r')
	
	# read each line
	for line in dataFile:
		# split line into fields
		fields = line.strip().split()
		if len(fields) < 1:
			continue
		try:
			# create row of floats
			row = map(float, fields)
			dataList.append(row)
		except:
			# create name/value pairs
			if len(fields) > 1:
				name = fields[0]
				value = fields[1]
				try:
					# try to set value as a float
					header[name] = float(value)
				except:
					header[name] = value
			
	dataFile.close()
	return (header, array(dataList))

def sigintHandler(a, b):
	'''Control-C handler'''
	sys.exit(0)

if __name__ == '__main__':
	if len(sys.argv) < 2:
		print 'Usage: PlotGrid.py data-file [min max mode]'
		print '       data-file = text file of 2d array with optional header'
		print '                   set to "demo" for demo mode'
		print '       min = data minimum'
		print '       max = data maximum (set both to 0 for auto-scaling, default)'
		print '       mode = 0 (b & w), 1 (color, default)'
		sys.exit(1)
		
	# install ctrl-c handler
	signal.signal(signal.SIGINT, sigintHandler)
	
	# get arguments
	fname = sys.argv[1]
	min = 0.0
	if len(sys.argv) > 2:
		min = float(sys.argv[2])
	max = 0.0
	if len(sys.argv) > 3:
		max = float(sys.argv[3])
	mode = 1
	if len(sys.argv) > 4:
		mode = int(sys.argv[4])
	
	if fname == 'demo':
		# create demo data
		data = fromfunction(lambda j, i: sin(j / 512.0 * 4.0 * pi) * sin(i / 512.0 * 4.0 * pi), (512, 512))
		valid = data >= 0
	else:
		# read data and header
		header, data = readData(fname)
		print 'Header = ', header
		print 'Shape = ', data.shape
		# set mask
		if 'NODATA_value' in header:
			valid = data != header['NODATA_value']
		else:
			valid = data == data
		print 'Valid data: %d entries, max = %g, min= %g, mean = %g, std = %g' % \
			(valid.sum(), data[valid].max(), data[valid].min(), data[valid].mean(), data[valid].std())
	
	# create app
	app = wx.PySimpleApp(False)
	# create window
	win = ImagePlot(data, valid, fname, min, max, mode=mode)
	win.Show()
	# handle events
	app.MainLoop()