#!/usr/bin/env python

"""When run as a script:
python ~/lib/wavio.py [-g gain] -wavin|-wavout infile outfile
Reads or writes .wav files from any format supported by
gpkimgclass.py .

As a library, allows reading of class gpk_img data from a .WAV
file, and writing class gpk_img to a .WAV file.
"""
import wave
import gpkimgclass
import Num

#: Cause L{write} to raise an error on overflow.
OV_ERROR = 0
#: Cause L{write} to truncate overflows.
OV_LIMIT = 1
#: Cause L{write} to silently ignore overflows.
OV_IGNORE = 2

class Overflow(ValueError):
	def __init__(self, *s):
		ValueError.__init__(self, *s)

class BadFileFormatError(Exception):
	def __init__(self, *s):
		Exception.__init__(self, *s)



_sizes = {
	2: (Num.int16, 16),
	4: (Num.int32, 32)
	}


def read_hdr(fn):
	"""Read header information for a .WAV file.   The header contains
	information formatted as per L{gpkimgclass.gpk_img}, i.e. FITS standard information.
	@param fn: filename
	@type fn: str
	@rtype: dict(str: str or float or int)
	"""
	try:
		w = wave.open(fn, "r")
	except wave.Error, x:
		raise BadFileFormatError, "%s: %s" % (x, fn)
	assert w.getcomptype() == 'NONE', "Can't handle %s compression" % w.getcompname()
	nc = w.getnchannels()
	dt = 1.0/w.getframerate()
	numtype, bitpix = _sizes[w.getsampwidth()]
	nf = w.getnframes()
	return {
		'NAXIS1': nc, 'NAXIS2': nf,
		'CDELT2': dt, 'CRPIX2': 1, 'CRVAL2': 0.0,
		'BITPIX': bitpix, '_NAME': fn, '_FILETYPE': 'WAV'
		}


def read(fn, tstart=None, t_end=None):
	"""Reads in a .WAV file and returns a L{gpkimgclass.gpk_img<gpkimgclass>} instance
	that contains the header and data information.   If tstart and/or
	t_end are specified, it seeks and only reads in the necessary data
	between tstart and t_end.
	@param fn: filename
	@type fn: str
	@param tstart: time at which to start reading audio data
	@param t_end: time at which to stop reading audio data
	@type tstart: float
	@type t_end: float
	@return: audio data
	@rtype: L{gpkimgclass} instance
	"""
	try:
		w = wave.open(fn, "r")
	except wave.Error, x:
		raise BadFileFormatError, "%s: %s" % (x, fn)
	assert w.getcomptype() == 'NONE', "Can't handle %s compression" % w.getcompname()
	nc = w.getnchannels()
	dt = 1.0/w.getframerate()
	numtype, bitpix = _sizes[w.getsampwidth()]
	if tstart is not None:
		istart = int(round(tstart/dt))
		if istart < 0:
			raise ValueError, "tstart=%.3fs < 0" % tstart
		w.setpos(istart)
	else:
		istart = 0
		tstart = 0.0
	if t_end is not None:
		nf = int(round(t_end/dt))
		if nf > w.getnframes():
			raise ValueError, "t_end=%.3fs which is beyond the end of the wave file (%.3fs)" % (t_end, w.getnframes()*dt/nc)
		nf -= istart
	else:
		nf = w.getnframes() - istart
	if nf < 0:
		raise ValueError, "t_end=%.4f < tstart=%.4f" % (t_end, tstart)
	data = Num.fromstring(w.readframes(nf), numtype)
	w.close()
	assert data.shape[0]==nf*nc, "Reshape from %d to %dx%d" % (data.shape[0], nf, nc)
	data = Num.reshape(data, (nf, nc))
	hdr = {
		'NAXIS1': nc, 'NAXIS2': nf,
		'CDELT2': dt, 'CRPIX2': 1, 'CRVAL2': istart*dt,
		'BITPIX': bitpix, '_NAME': fn, '_FILETYPE': 'WAV'
		}
	return gpkimgclass.gpk_img(hdr, data)


def _itemsize(z):
	"""This exists for Numeric/NumPy compatibility"""
	x = z.itemsize
	if isinstance(x, int):
		return x
	if callable(x):
		return x()
	raise RuntimeError, 'Cannot deal with %s' % str(type(x))

_typecodes = {
	_itemsize(Num.zeros((1,), Num.Int32))*8:
			(Num.Int32, 2.0**31-1.0, 1.0-2.0**31),
	_itemsize(Num.zeros((1,), Num.Int16))*8:
			(Num.Int16, 2.0**15-1.0, 1.0-2.0**15),
	_itemsize(Num.zeros((1,), Num.Int8))*8:
			(Num.Int8, 2.0**7-1.0, 1.0-2.0**7)
	}

def write(data, fname, scalefac=1, allow_overflow=OV_ERROR):
	"""@param data: is a class gpk_img object containing
		data to be written (note that the header information is
		ignored except for the sampling rate, bits per pixel,
		and number of channels.
	@param fname: is the name of a file to write it to (or a file object),
	@type fname: str or file,
	@param scalefac: is a factor to multiply the data
	@param allow_overflow: can be either
		- L{OV_ERROR} (default, means raise a ValueError exception
			if the data*scalefac overflows),
		- L{OV_LIMIT} (means limit the data*scalefac to prevent
			overflows -- this clips the audio), or
		- L{OV_IGNORE} (means let the overflows happen and don't worry.)
	@except ValueError: Missing information in C{data}.
	@except Overflow: The output format is clipping the data.   Strictly
		speaking, on 16-bit data, 32767 is considered clipping even
		though it possibly might be OK.   However such extreme values
		are very likely to be generated by clipping at an earlier
		stage of the processing, so it's probably an error even if
		the error is not being made here.
	"""
	if not (data.dt() > 0.0):
		raise ValueError, "Cannot set sampling rate: dt=%g\n" % data.dt()

	hdr = data.hdr

	if not hdr.has_key('BITPIX') or int(hdr['BITPIX'])==0:
		bitpix = 16
	else:
		bitpix = abs(int(hdr['BITPIX']))
	assert bitpix % 8 == 0
	if bitpix > 32:
		bitpix = 32

	tc, vmax, vmin = _typecodes[bitpix]
	dds = data.d * scalefac
	if allow_overflow == OV_ERROR:
		if not Num.alltrue(Num.greater_equal(Num.ravel(dds), vmin)):
			raise Overflow, "Scaled data overflows negative"
		if not Num.alltrue(Num.less_equal(Num.ravel(dds), vmax)):
			raise Overflow, "Scaled data overflows positive"
	elif allow_overflow == OV_LIMIT:
		dds = Num.clip(dds, vmin, vmax)
	else:
		assert allow_overflow == OV_IGNORE

	d = Num.around(dds).astype(tc).tostring()
	w = wave.open(fname, "w")
	w.setnchannels(data.d.shape[1])
	w.setnframes(data.d.shape[0])
	w.setsampwidth(bitpix/8)
	w.setframerate( int(round( 1.0/float(data.dt()) )) )
	w.writeframesraw(d)
	w.close()


if __name__ == '__main__':
	import sys
	arglist = sys.argv[1:]
	gain = None
	if len(arglist)==0:
		print __doc__
		sys.exit(1)
	if arglist[0] == '-g':
		arglist.pop(0)
		gain = float(arglist.pop(0))
	if arglist[0] == '-wavin':
		x = read(arglist[1])
		if gain is not None:
			Num.multiply(x.d, gain, x.d)
		x.write(arglist[2])
	elif arglist[0] == '-wavout':
		x = gpkimgclass.read(arglist[1])
		rxd = Num.ravel(x.d)
		mxv = max( rxd[Num.argmax(rxd)], -rxd[Num.argmin(rxd)] )
		print "# max=", mxv
		if gain is None:
			gain = 32000.0/mxv
		write(x, arglist[2], gain)
	else:
		print __doc__
		sys.exit(1)
