"""This module is a python interface to the gpkio library.
That library can read in a variety of 2-dimensional image
types and multichannel audio data.

When used with a current version of the libgpkio library,
* it reads and writes in NASA FITS files that contain 2-D images
	(these are used for astronomical data, typically
	images from telescopes),
* it reads HTK HParm format files
	(these are used for speech data and feature vectors
	for the HTK Hidden Markov Model Toolkit,
	for speech recognition systems),
* it reads and writes the NIST SPHERE speech data format
	(though not all variants),
* it reads a header + data format called GPK ASCII Image
	that I frequently use.   (This format can
	give you a completely ASCII data file, or have
	a human-readable ASCII header followed by binary
	data.  It's designed for 2-D images and adapts well
	to multichannel audio data.)
* it reads and writes the obsolete J. A. Tyson data format called FOCAS,
	which is a binary format for 2-D astronomical images,
* it reads and writes an old AT&T speech data format called SIG
	(used for speech data),
* it reads an old AT&T speech data format called ATTSSW
	(used for speech data),
* it reads a data format called "EST" from the "Edinburgh
	Speech Tools".

The code is designed to share an ontology, so that you can
read in one format and write the data in another format,
without the user needing to transform the header or data.

The code is designed to support data formats that have a
header that can store key=value strings, followed by a
1- or 2-dimensional data array.

Some of the data formats have fewer restrictions than FITS
as to the content of file headers: longer keys or attributes
may be allowed, and/or a wider variety of characters.
The module does its best not to lose header information,
but sometimes it must be re-keyed or truncated.
"""

import types
import gpkimg
import numpy


class CannotReadDataFile(IOError):
	"""Thrown by read()."""
	def __init__(self, *x):
		IOError.__init__(self, x)


class NoSuchColumn(KeyError):
	"""Thown by column()."""
	def __init__(self, *d):
		KeyError.__init__(self, d)

def _pull_types(d, nc):
	"""Finds all the column definitions in the data file,
	that is, attributes in the form 'TTYPE#'.   These
	attributes tell you what kind of data is in a specific
	column.  It then builds an output dictionary that
	tells you what column contains any desired type of data."""

	PREFIX = 'TTYPE'
	Pl = len(PREFIX)
	o = {}
	for k in range(nc):
		o[k] = k	# Columns can also be indexed by integer numbers.
		sk = '%d' % k
		o[sk] = k
	for k in d.keys():
		assert isinstance(k, str), 'Key must be string'
		# print "PULLTYPES(", k, ")", 'type(k)=', type(k)
		if k.startswith(PREFIX):
			col = int(k[Pl:])
			if col>0 and col<=nc:
				o[d[k]] = col-1
	return o




def _del(d, k):
	if d.has_key(k):
		del d[k]
		return 1
	return 0


def read(f):
	"""Reads in a data file and returns its contents.
	File names may end in .Z or .gz if the data is compressed.
	@param f: name of the data file
	@type f: str
	@rtype: L{gpk_img}
	@return: the image
	@raise CannotReadDataFile: If the file format is unrecognized or invalid.
	"""
	try:
		hdr, data = gpkimg.read(f, 0)
	except IOError, x:
		raise CannotReadDataFile(x.args[0])

	hdr['_NAME'] = f
	return gpk_img(hdr, data)


def read_hdr(f):
	"""Reads in the header of a data file and returns its contents.
	File names may end in .Z or .gz if the data is compressed.
	@param f: name of the data file
	@type f: str
	@rtype: L{dict}
	@return: the image's header information
	@raise CannotReadDataFile: If the file format is unrecognized or invalid.
	"""
	try:
		hdr = gpkimg.read(f, 1)
	except IOError, x:
		raise CannotReadDataFile(x.args[0])
	hdr['_NAME'] = f
	return hdr


# Note: this isn't a member function because it can be usefully applied to the
# output of read_header()
def extract_prefix(prefix, hdr):
	"""This extracts a vector of float values from a header dictionary.
	It takes values from keys in the form prefix1 ... prefixN; it takes
	however many there are and returns them in a numpy vector.
	@type prefix: str
	@type hdr: dict(str:str)
	@return: a vector of the extracted values.
	@rtype: numpy.ndarray(float)
	"""
	v = []
	while True:
		try:
			v.append( float(hdr["%s%d" % (prefix, len(v)+1)]) )
		except KeyError:
			break
	if not v:
		raise KeyError("No key in the form '%s%%d'" % prefix)
	return numpy.array(v, numpy.float)


def extract_n(prefix, hdr, n):
	"""This extracts values from a header dictionary.
	It takes C{n} values from keys in the form C{prefix1} ... C{prefixN}.
	@type prefix: str
	@type n: int
	@param n: now many?
	@type hdr: dict(str:str)
	@return: the extracted values.
	@rtype: list(str)
	"""
	rv = []
	for i in range(1, n+1):
		rv.append(hdr['%s%d' % (prefix, i)])
	return rv


CacheSize = 10*1024*1024	#: How many samples worth of data should we store in the cache?
_csize = 0
_cache = {}
def CachedRead(f):
	"""Reads in a data file (f is the name) and returns a
	class gpk_img object.    File names may end in .Z or .gz
	if the data is compressed.
	Raises CannotReadDataFile if it does not recognize the file
	format.
	"""
	global _csize, _cache, CacheSize
	try:
		return _cache[f].copy()
	except KeyError:
		pass
	while _cache and _csize >= CacheSize:
		x = _cache.pop()
		_csize -= x.d.shape[1]*x.d.shape[0]
	if len(_cache) == 0 or _csize < 0:
		_csize = 0
	tmp = read(f)
	_cache[f] = tmp
	_csize += tmp.d.shape[1]*tmp.d.shape[0]
	return tmp.copy()



class gpk_img:
	"""This class describes a data file.  It contains header
	information and data.   The header information follows
	the NASA FITS standard closely with a few extensions.

	Header items:
	_NAME -- the filename from which the data came,
	_FILETYPE -- a string indicating the format of the data file

	Each instance has two attributes that are intended for general use:
	hdr -- a dictionary containing all the header information, and
	d -- a 2-dimensional numpy array containing the data.
	"""

	def __init__(self, hdr, data=None):
		"""Creates an gpk_img object from
		hdr (a dictionary that describes a data array)
		and the data (a Numeric 2-dimensional array).
		"""

		self.hdr = hdr.copy()
		try:
			self.d = numpy.asarray(data, numpy.float)
		except ValueError,x:
			raise ValueError, '%s (you gave it %s)' % (str(x), str(type(data)))

		if len(self.d.shape) == 1:
			self.d = numpy.transpose([self.d])
		elif len(self.d.shape) == 2:
			if(self.hdr.has_key('NAXIS')):
				assert self.hdr['NAXIS'] in ['2', 2]
				del self.hdr['NAXIS']
		else:
			raise ValueError, 'Needs 1 or 2 dimensional arrays'
	
		# self.n = (None, string.atoi(d.get('NAXIS1', '1')),
		#		string.atoi(d.get('NAXIS2', '1')))
		# self.n = (None, self.d.shape[1], self.d.shape[0])
		# _del(self.hdr, 'NAXIS1')
		# _del(self.hdr, 'NAXIS2')
		self.cdelt = (None, float(hdr.get('CDELT1', 'NaN')),
				float(hdr.get('CDELT2', 'NaN')))
		# _del(self.hdr, 'CDELT1')
		# _del(self.hdr, 'CDELT2')
		self.crpix = (None, float(hdr.get('CRPIX1', 'NaN')),
				float(hdr.get('CRPIX2', 'NaN')))
		# _del(self.hdr, 'CRPIX1')
		# _del(self.hdr, 'CRPIX2')
		self.crval = (None, float(hdr.get('CRVAL1', 'NaN')),
				float(hdr.get('CRVAL2', 'NaN')))
		# _del(self.hdr, 'CRVAL1')
		# _del(self.hdr, 'CRVAL2')
		self.types = _pull_types(self.hdr, self.d.shape[0])
	
	def __repr__(self):
		return "<gpk_img: %dx%d %g+%s* %g+%s* %s>" % (self.d.shape[0], self.d.shape[1],
				self.coord(0.0, 1), self.hdr.get('CDELT1', '???'),
				self.coord(0.0, 2), self.hdr.get('CDELT2', '???'),
				self.hdr.get('_NAME', '???')
				)

	def copy(self):
		return gpk_img(self.hdr, numpy.array(self.d, numpy.float, copy=True))

	def add_history(self, h):
		"""Adds an item to the file's history."""
		if 'HISTORY' in self.hdr:
			self.hdr['HISTORY'] += '\n%s' % h.strip()
		else:
			self.hdr['HISTORY'] = h.strip()

	def get_history(self):
		"""Returns the file's history as a list of strings."""
		return self.hdr.get('HISTORY', '').split('\n')

	def empty(self):
		"""Returns True if there is no data."""
		return self.d.shape[0]<=0 or self.d.shape[1]<=0

	def index(self, t, i, limit=False, error=True):
		"""What is the i^th index if the i^th coordinate = t?
		More plainly, this function gives you the index that
		corresponds to the coordinate, and 'i' selects which index.
		(i=1 or 2).
		"""
		n = self.d.shape[2-i]
		assert self.cdelt[i] != 0
		tmp = (self.crpix[i]-1) + (t-self.crval[i])/self.cdelt[i]
		if limit:
			if hasattr(t, 'shape'):
				numpy.maximum(tmp, 0.0, tmp)
				numpy.minimum(tmp, n-1, tmp)
			elif not (0 < tmp < n-1):
				tmp = max(0.0, min(n-1, tmp))
		elif error:
			if hasattr(t, 'shape'):
				if not (numpy.alltrue(numpy.greater_equal(tmp, 0.0))
					and numpy.alltrue(numpy.less_equal(tmp,n-1))
					):
					raise IndexError, "coordinate out of range"
			elif not (0 <= tmp <= n-1):
				raise IndexError, "coordinate out of range: %g vs. (%g, %g)" % (t, self.coord(0, i), self.coord(n-1,i))
		return tmp

	def coord(self, index, i):
		"""What is the i^th coordinate if the i^th index==index?
		More plainly, this function gives you the coordinate that
		corresponds to the index, and 'i' selects which index.
		If index=None, it returns an array of all the coordinates.
		(i=1 or 2).
		"""
		n = self.d.shape[2-i]
		if index is None:
			index = numpy.arange(n)
		return self.crval[i] + (index-(self.crpix[i]-1))*self.cdelt[i]

	def x_index(self, x, **flags):
		"""Returns the y-index for a particular x-coordinate."""
		return self.index(x, 1, **flags)

	def y_index(self, y, **flags):
		"""Returns the y-index for a particular y-coordinate.
		Also used for t-index, which returns the index
		for a particular time, if the data is a time-series.
		"""
		return self.index(y, 2, **flags)

	def x(self, index = None):
		"""Returns the x-coordinate for a particular x-index."""
		return self.coord(index, 1)

	def y(self, index = None):
		"""Returns the y-coordinate for a particular y-index."""
		return self.coord(index, 2)



	def write(self, f):
		"""Write out a data file, with the information in this class."""
		d = self.hdr.copy()
		sh = self.d.shape

		if len(sh) == 2:
			data = self.d
		else:
			raise RuntimeError, '2 dimensional arrays only'

		nc = self.d.shape[1]
		for i in (1,2):
			si = str(i)
			d['NAXIS' + si] = str(sh[i-1])
			d['CDELT' + si] = str(self.cdelt[i])
			d['CRPIX' + si] = str(self.crpix[i])
			d['CRVAL' + si] = str(self.crval[i])
		for k in d.keys():
			if k.startswith('TTYPE'):
				cs = k[len('TTYPE'):]
				col = int(cs)
				if col<=0 or col>nc:
					self._wipe(col)
		avpairs = [ (x[0], str(x[1])) for x in d.items() ]
		# print "# avpairs=", avpairs
		gpkimg.write(f, avpairs, data)


	def __len__(self):
		"""This is here to make it possible to access the class
		as a (header,data) tuple.
		"""
		return 2	# header, data


	def __getitem__(self, key):
		"""This is here to make it possible to access the class
		as a (header,data) tuple.
		"""
		return (self.hdr, self.d)[key]


	# Below is time-series stuff.   Should be a seperate class.

	def time(self, index=None):
		"""What time corresponds to a certain index?
		(This is the same as self.y(index).)"""
		return self.coord(index, 2)

	t_index = y_index


	def start(self):
		"""Starting time of the data."""
		return self.time(0)

	def end(self):
		"""The time of the last sample."""
		return self.time(self.d.shape[0]-1)


	def dt(self):
		"""Assuming time flows down a column (axis=2), what's the time interval?"""
		return self.cdelt[2]


	def column(self, s=None):
		"""Return the requested column of data.
		(This is the actual data, not a copy.)
		If s is a string, it returns the column which has TTYPEx=s.
		Otherwise, if s is an integer, it returns that column
		(zero-based index).
		As a special case, if there is only one column,
		and the function is called with s=None (the default),
		it will return the data.
		"""

		if s is None and self.d.shape[1]==1:
			return self.d[:,0]

		try:
			ic = self.types[s]
		except KeyError:
			if '_NAME' in self.hdr:
				xargs = 'file=%s; column=%s' % (
						self.hdr['_NAME'], str(s))
			else:
				xargs = 'column=%s' % str(s)
			raise NoSuchColumn, xargs
		return self.d[:,ic]


	def _idxhdrset(self, prefix, index, value):
		a = '%s%d' % (prefix, index+1)
		if value is not None:
			self.hdr[a] = value
		else:
			_del(self.hdr, a)


	def _wipe(self, col):
		nd = 0
		for x in ('TTYPE', 'TUNIT', 'TRANGE'):
			nd += _del(self.hdr, '%s%d' % (x, col+1))
		return nd


	def register(self, ttype, column, trange=None, tunit=None):
		"""This sets up the TTYPE fields so that
		the specified datatype lives in the specified
		column.  Column is zero-based.
		You can erase by specifying ttype=None.
		"""
		# print 'REGISTER', ttype, column, trange, tunit
		assert isinstance(ttype, types.StringType)
		assert isinstance(column, types.IntType)

		# First, you find out the old type that was stored
		# in the column.  That needs to be deleted from self.types.
		if self.hdr.has_key('TTYPE%d' % (column+1)):
			oldtype = self.hdr['TTYPE%d' % (column+1)]
			_del(self.types, oldtype)
		# Then, you wipe out the information in the header:
		self._wipe(column)

		# Finally, you set things:
		if ttype is not None:
			# print 'setting types', ttype, column
			self.types[ttype] = column

		# print 'SETTING HDR', column, ttype
		self._idxhdrset('TTYPE', column, ttype)
		self._idxhdrset('TUNIT', column, tunit)
		self._idxhdrset('TRANGE', column, trange)
		self.types[column] = column


	def attach_col(self, data, ttype, trange=None, tunit=None, extra_hdr={}):
		"""Attach another column to an existing class.
		This modifies the object it is called upon.
		The data may or may not be copied in the process.
		"""
		assert len(data.shape) == 1
		assert ttype is not None
		assert self.d is None or data.shape[0] == self.d.shape[0]
		self.hdr.update(extra_hdr)
		if self.d is None:
			self.d = numpy.array([data], numpy.float)
			c = 0
		elif ttype in self.types:
			c = self.types[ttype]
			self.d[:,c] = numpy.asarray(data, numpy.float)
		else:
			self.d = numpy.concatenate([self.d, numpy.transpose([data])], axis=1)
			c = self.d.shape[1]-1
		self.register(ttype, c, trange=trange, tunit=tunit)


	def select_cols(self, columns, extra_hdr = {}):
		"""Select a set of columns in an existing class,
		returning a new gpk_img class containing a subset of
		the data.   The data is copied; the new and old
		objects share nothing.
		"""
		newdata = []
		for c in columns:
			newdata.append( self.column(c) )
		n = gpk_img(self.hdr, numpy.transpose(newdata))
		n.hdr.update( extra_hdr )
		nic = 0
		for c in columns:
			ic = self.types[c]+1
			# print 'Registering %s (old col %d) into new col %d' % (c, ic, nic)
			n.register(self.hdr.get('TTYPE%d' % ic, None), nic,
					self.hdr.get('TRANGE%d' % ic, None),
					self.hdr.get('TUNIT%d' % ic, None))
			nic += 1

		while n._wipe(nic):	# Clear out misleading TTYPE info.
			nic += 1

		return n
