#! python

"""Usage: make_mfcc_files_for_rec.py [flags] list-of-metadata-files
flags: -Fsimple2
	-Fsimple4

This computes MFCC files containing feature vectors for HTK.   The overall
intent is to do a simple segmentation of the speech based on just
loudness and a vowel/consonant distinction.

RUN THIS BEFORE running rec*.sh rec.sh or rec_all.sh or rec_all.py.
It generates necessary .mfcc files.
"""

import os
import sys
import math
import tempfile
import subprocess
import numpy
import gpkimgclass as GI
from gmisclib import die
from gmisclib import avio
from gmisclib import tsops
from gmisclib import fiatio
from gmisclib import gpkmisc
from gmisclib import threaded_io
from gpk_voicing import voice_misc
from gpk_voicing import zero as Z
from gmisclib import Numeric_gpk as NG
from gmisclib import weighted_percentile as WP


IRR_AVG = 0.43	# Derived from simple segmentation .../proc/bin/do_simple_segmentation.sh

ThisCode = gpkmisc.get_mtime(sys.argv[0])

def simple2(path, oname, lazy=False):
	"""This produces a time series of sonority and frication for HTK's enjoyment.
	"""
	if lazy:
		mtime = gpkmisc.prereq_mtime(gpkmisc.get_mtime(os.path.join(path, 'loud.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'irr.dat'))
					)
		if not gpkmisc.need_to_recompute(oname, mtime, size=0):
			return

	ld = GI.read(os.path.join(path, 'loud.dat'))
	ird = GI.read(os.path.join(path, 'irr.dat'))
	taxis, loud, irr = tsops.common([ld, ird])
	loud = loud[:,ld.types['loudness']]
	loud -= Z.loud_zero(loud)
	avloud = numpy.sum(loud**2)/numpy.sum(loud)
	lnorm = loud/avloud
	irr = irr[:,ird.types['aperiodicity']]
	irr_avg = numpy.sum(irr*loud)/numpy.sum(loud)
	fric = lnorm * (irr - irr_avg)
	h = {'CDELT2': taxis.dt()}
	o = GI.gpk_img(h, numpy.transpose([lnorm, fric]))
	o.write(oname)



def simple3(path, oname, lazy=False, norm_irr=True):
	"""This produces a time series of sonority and frication for HTK's enjoyment.
	"""
	if lazy:
		mtime = gpkmisc.prereq_mtime(gpkmisc.get_mtime(os.path.join(path, 'loud.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'irr.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'pdur.dat'))
					)
		if not gpkmisc.need_to_recompute(oname, mtime, size=0):
			return

	ld = GI.read(os.path.join(path, 'loud.dat'))
	ird = GI.read(os.path.join(path, 'irr.dat'))
	pd = GI.read(os.path.join(path, 'pdur.dat'))
	taxis, loud, irr, pdur = tsops.common([ld, ird, pd])
	loud = loud[:,ld.types['loudness']]
	loud = numpy.maximum(0.0, loud-Z.loud_zero(loud))
	avloud = numpy.sum(loud**2)/numpy.sum(loud)
	lnorm = loud/avloud
	irr = irr[:,ird.types['aperiodicity']]
	if norm_irr:
		irr_avg = numpy.sum(irr*loud)/numpy.sum(loud)
	else:
		irr_avg = IRR_AVG
	lpdur = numpy.log(pdur[:,pd.types['pseudoduration']])
	lpdur -= WP.wtd_median(lpdur, loud)
	fric = lnorm * (irr - irr_avg)
	h = {'CDELT2': taxis.dt()}
	o = GI.gpk_img(h, numpy.transpose([lnorm, fric, lpdur]))
	o.write(oname)


def simple3u(path, oname, lazy):
	return simple3(path, oname, lazy, norm_irr=False)


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

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

class DataProblem(Exception):
	def __init__(self, *s):
		Exception.__init__(self, *s)
	
	def annotate(self, *s):
		self.args = self.args + tuple(s)
		return self


def minimum_near1(t1, x):
	"""Smooth x over t1, then look for the smallest value of the
	result over a range of t2.    Basically, this is a pause
	detector.   It's small if you are in a region of silence
	of length greater than t1.
	"""
	HUGE = 1e10
	nx = x.shape[0]
	np1 = int(round(t1))
	if nx < np1:
		raise NotEnoughData
	pk1 = voice_misc.window(2*np1+1, 0, norm=1)
	t = numpy.arange(pk1.shape[0]) - np1
	t2 = math.sqrt(numpy.sum(pk1*t**2)/numpy.sum(pk1))
	np2 = int(round(t2))
	tmp = numpy.convolve(x, pk1, 1)
	rv = numpy.zeros((nx+2*np2,)) + HUGE
	for i in range(-np2, np2+1):
		numpy.minimum(rv[np2+i:np2+i+nx], tmp, rv[np2+i:np2+i+nx])
	return rv[np2:nx+np2]


def median_before(t1, x):
	nx = x.shape[0]
	np1 = int(round(t1*2))
	if nx < np1:
		raise NotEnoughData
	xx = numpy.zeros((nx+np1,))
	xx[np1:np1+nx] = x
	rv = numpy.zeros((nx,))
	for i in range(nx):
		rv[i] = numpy.median( xx[i-np1:i] )
	return rv


def simpleP1(path, oname, lazy=False):
	"""This produces a time series of sonority and frication for HTK's enjoyment.
	"""
	Pause1 = 0.075	# Seconds
	if lazy:
		mtime = gpkmisc.prereq_mtime(gpkmisc.get_mtime(os.path.join(path, 'loud.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'irr.dat'))
					)
		if not gpkmisc.need_to_recompute(oname, mtime, size=0):
			return

	ld = GI.read(os.path.join(path, 'loud.dat'))
	ird = GI.read(os.path.join(path, 'irr.dat'))
	taxis, loud, irr = tsops.common([ld, ird])
	loud = loud[:,ld.types['loudness']]
	loud -= Z.loud_zero(loud)
	avloud = numpy.sum(loud**2)/numpy.sum(loud)
	lnorm = loud/avloud
	irr = irr[:,ird.types['aperiodicity']]
	irr_avg = numpy.sum(irr*loud)/numpy.sum(loud)
	fric = lnorm * (irr - irr_avg)
	lnear1 = minimum_near1(Pause1/taxis.dt(), lnorm)
	h = {'CDELT2': taxis.dt()}
	o = GI.gpk_img(h, numpy.transpose([lnorm, lnear1, fric]))
	o.write(oname)


def simpleP2(path, oname, lazy=False):
	"""This produces a time series of sonority and frication for HTK's enjoyment.
	"""
	Pause1 = 0.075	# Seconds
	if lazy:
		mtime = gpkmisc.prereq_mtime(gpkmisc.get_mtime(os.path.join(path, 'loud.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'irr.dat'))
					)
		if not gpkmisc.need_to_recompute(oname, mtime, size=0):
			return

	ld = GI.read(os.path.join(path, 'loud.dat'))
	ird = GI.read(os.path.join(path, 'irr.dat'))
	taxis, loud, irr = tsops.common([ld, ird])
	loud = loud[:,ld.types['loudness']]
	loud -= Z.loud_zero(loud)
	avloud = numpy.sum(loud**2)/numpy.sum(loud)
	lnorm = loud/avloud
	irr = irr[:,ird.types['aperiodicity']]
	irr_avg = numpy.sum(irr*loud)/numpy.sum(loud)
	fric = lnorm * (irr - irr_avg)
	lnear1 = median_before(Pause1/taxis.dt(), lnorm)
	h = {'CDELT2': taxis.dt()}
	o = GI.gpk_img(h, numpy.transpose([lnorm, lnear1, fric]))
	o.write(oname)


def simple4(path, oname, lazy=False):
	PAST = 0.450	# Seconds
	return simple4_guts(path, oname, PAST, lazy=lazy)


def simple4s(path, oname, lazy=False):
	PAST = 0.250	# Seconds
	return simple4_guts(path, oname, PAST, lazy=lazy)


def simple4_guts(path, oname, past, lazy=False):
	"""This produces sonority, frication, and (presumably for comparison purposes),
	averages of the recent past sonority and frication.
	"""
	if lazy:
		mtime = gpkmisc.prereq_mtime(gpkmisc.get_mtime(os.path.join(path, 'loud.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'irr.dat'))
					)
		if not gpkmisc.need_to_recompute(oname, mtime, size=0):
			return

	ld = GI.read(os.path.join(path, 'loud.dat'))
	ird = GI.read(os.path.join(path, 'irr.dat'))
	taxis, loud, irr = tsops.common([ld, ird])
	loud = loud[:,ld.types['loudness']]
	loud -= Z.loud_zero(loud)
	avloud = numpy.sum(loud**2)/numpy.sum(loud)
	lnorm = loud/avloud
	irr = irr[:,ird.types['aperiodicity']]
	irr_avg = numpy.sum(irr*loud)/numpy.sum(loud)
	fric = lnorm * (irr - irr_avg)
	NP = int(round(past/taxis.dt()))
	recentKernel = numpy.zeros((2*NP+1,))
	recentKernel[1+NP:] = NP - numpy.arange(NP)
	recentKernel /= numpy.sum(recentKernel)
	recent_lnorm = NG.convolve(lnorm, recentKernel)
	recent_fric = NG.convolve(fric, recentKernel)
	h = {'CDELT2': taxis.dt()}
	o = GI.gpk_img(h, numpy.transpose([lnorm, recent_lnorm, fric, recent_fric]))
	o.write(oname)



def mfcc(path, oname, lazy=False):
	if lazy:
		wav = os.path.join(path, 'raw.wav')
		try:
			mtime = gpkmisc.prereq_mtime(gpkmisc.get_mtime(wav),
					gpkmisc.get_mtime('training/coding')
					)
		except gpkmisc.PrereqError, x:
			x.repl((wav, 'training/coding')[x[1]])
			raise x
		if not gpkmisc.need_to_recompute(oname, mtime, size=0):
			return
	subprocess.check_call(['HCopy', '-C', 'training/coding', os.path.join(path, 'raw.wav'), oname])


def check_lengths_agree(*x):
	"""Do a set of L{gpkimgclass.gpk_img} objects contain about the same amount of data?
	"""
	for xx in x:
		for yy in x:
			if id(xx) < id(yy):
				if abs(xx.n[2]-yy.n[2])/float(xx.n[2]+yy.n[2]+30) > 0.1:
					raise "Lengths disagree in samples: %d vs %d" % (xx.n[2], yy.n[2])
				xxdur = xx.end() - xx.start()
				yydur = yy.end() - yy.start()
				if abs(xxdur-yydur)/(xxdur+yydur+0.3) > 0.1:
					raise "Lengths disagree in time: %.3f-%.3f=%.3f vs %.3f-%.3f=%.3f" % (
						xx.end(), xx.start(), xxdur, yy.end(), yy.start(), yydur
						)


def mfccplus(path, oname, lazy=False):
	if lazy:
		wav = os.path.join(path, 'raw.wav')
		try:
			mtime = gpkmisc.prereq_mtime(gpkmisc.get_mtime(wav),
					gpkmisc.get_mtime('training/coding'),
					gpkmisc.get_mtime(os.path.join(path, 'loud.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'irr.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'pdur.dat')),
					ThisCode
					)
		except gpkmisc.PrereqError, x:
			x.repl((wav, 'training/coding')[x[1]])
			raise x
		if not gpkmisc.need_to_recompute(oname, mtime, size=0):
			return
	fd, tmpname = tempfile.mkstemp(suffix='.mfcc')
	subprocess.check_call(['HCopy', '-C', 'training/coding', os.path.join(path, 'raw.wav'), tmpname])

	mfccd = GI.read(tmpname)
	os.close(fd)
	os.remove(tmpname)
	ld = GI.read(os.path.join(path, 'loud.dat'))
	ird = GI.read(os.path.join(path, 'irr.dat'))
	pd = GI.read(os.path.join(path, 'pdur.dat'))
	try:
		check_lengths_agree(mfccd, ld, ird, pd)
	except DataProblem, x:
		raise x.annotate(path)
	taxis, loud, irr, pdur, mfcc = tsops.common([ld, ird, pd, mfccd])
	loud = loud[:,ld.types['loudness']]
	if loud.shape[0] < 2:
		raise NotEnoughData, "Cannot normalize loudness: only one value: utterance too short."
	lmin, lmax = WP.wp(loud, None, [0.01, 0.99])
	if not (lmax > lmin and lmin>=0.0):
		raise CannotNormalize, "Weird loudnesses: lmax=%s lmin=%s, length=%s" % (lmax, lmin, loud.shape)
	loud = (loud-lmin)/(lmax-lmin)

	weight = numpy.maximum(0.0, loud)
	assert numpy.greater(weight, 0.0).any()
	irr = irr[:,ird.types['aperiodicity']]
	irr_avg = numpy.sum(irr*weight)/numpy.sum(weight)
	lpdur = numpy.log(pdur[:,pd.types['pseudoduration']])
	lpdur -= WP.wtd_median(lpdur, weight)
	fric = loud * (irr - irr_avg)

	PHONE = 0.125
	NP1 = int(round(PHONE/taxis.dt()))
	NP2 = int(round(1.414*PHONE/taxis.dt()))
	NP3 = int(round(2*PHONE/taxis.dt()))
	kernel1 = voice_misc.window(NP1, 0, norm=1)
	kernel2 = voice_misc.window(NP2, 1, norm=1)
	kernel3 = voice_misc.window(NP3, 2, norm=1)
	loud3 = NG.convolve(loud, kernel3)
	fric1 = NG.convolve(fric, kernel1)
	fric2 = NG.convolve(fric, kernel2)
	fric3 = NG.convolve(fric, kernel3)

	h = {'CDELT2': taxis.dt()}
	x = numpy.transpose([loud3, fric1, fric2, fric3, lpdur])
	y = numpy.concatenate((mfcc, x), axis=1)
	assert x.shape[0]==y.shape[0], "%s %s" % (x.shape, y.shape)
	GI.gpk_img(h, y).write(oname)


def gaussianize(fv, fvhdr):
	vowels = []
	i = 1
	while 'F_INFO%d' % i in fvhdr:
		x = avio.parse(fvhdr['F_INFO%d' % i])
		if x['type'] == 'vowel':
			vowels.append( (float(x['erb']), i, x) )
		i += 1
	vowels.sort()
	out = numpy.array(fv)
	N = 3
	assert len(vowels) > 3
	for (j,(e, i, x)) in enumerate(vowels):
		e2, i2, x2 = vowels[(j+N)%len(vowels)]
		out[:,i] = fv[:,i] - 0.7*fv[:,i2]
	return out


def fv(path, oname, lazy=False):
	"""This format-converts the full Acoustic Description Vector."""
	if lazy:
		mtime = gpkmisc.prereq_mtime(gpkmisc.get_mtime(os.path.join(path, 'loud.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'irr.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'pdur.dat')),
					gpkmisc.get_mtime(os.path.join(path, 'fv.dat'))
					)
		if not gpkmisc.need_to_recompute(oname, mtime, size=0):
			return

	fvd = GI.read(os.path.join(path, 'fv.dat'))
	ld = GI.read(os.path.join(path, 'loud.dat'))
	ird = GI.read(os.path.join(path, 'irr.dat'))
	pd = GI.read(os.path.join(path, 'pdur.dat'))
	taxis, loud, irr, pdur, fv = tsops.common([ld, ird, pd, fvd])
	loud = loud[:,ld.types['loudness']]

	# Here, we are normalizing the loudness, so we darn well
	# ought to have some speech in the utterance as a comparison.
	lmin, lmax = WP.wp(loud, None, [0.01, 0.99])
	loud -= lmin
	weight = numpy.maximum(0.0, loud)
	irr = irr[:,ird.types['aperiodicity']]
	irr_avg = numpy.sum(irr*weight)/numpy.sum(weight)
	lpdur = numpy.log(pdur[:,pd.types['pseudoduration']])
	lpdur -= WP.wtd_median(lpdur, weight)
	fric = loud * (irr - irr_avg)

	PHONE = 0.125
	NP1 = int(round(PHONE/taxis.dt()))
	NP2 = int(round(1.414*PHONE/taxis.dt()))
	NP3 = int(round(2*PHONE/taxis.dt()))
	kernel1 = voice_misc.window(NP1, 0, norm=1)
	kernel2 = voice_misc.window(NP2, 1, norm=1)
	kernel3 = voice_misc.window(NP3, 2, norm=1)
	loud3 = NG.convolve(loud, kernel3)
	fric1 = NG.convolve(fric, kernel1)
	fric2 = NG.convolve(fric, kernel2)
	fric3 = NG.convolve(fric, kernel3)

	h = {'CDELT2': taxis.dt()}
	x = numpy.transpose([loud3, fric1, fric2, fric3, lpdur])
	assert x.shape[0]==fv.shape[0], "%s %s" % (x.shape, fv.shape)
	y = numpy.concatenate((x, gaussianize(fv, fvd.hdr)), axis=1)
	assert x.shape[0]==y.shape[0], "%s %s" % (x.shape, y.shape)
	GI.gpk_img(h, y).write(oname)



def process_one_db(db, fcn, oname, lazy=False, skiperr=False):
	dbdir, dbf = os.path.split(db)
	h, d, c = fiatio.read(open(db, 'r'))
	for info in d:
		# die.info('\tprocessing %s / %s / %s' %(dbdir, info['d'], info['f']))
		path = os.path.join(dbdir, info['d'], info['f'])
		pon = os.path.join(path, oname)
		try:
			fcn(path, pon, lazy=lazy)
		except (CannotNormalize, NotEnoughData, Z.ZeroProblem, GI.CannotReadDataFile), x:
			die.warn("Failure on %s/%s: %s" % (info['d'], info['f'], str(x)))
			try:
				os.remove(pon)
			except OSError, x:
				pass
		except gpkmisc.PrereqError, x:
			if skiperr:
				die.warn("Prerequisite Error: %s" % x)
			else:
				raise



def process_all(dblist, fcn, oname, lazy=False, skiperr=False):
	tp = threaded_io.Thread_poolW()
	for db in dblist:
		die.info('Processing %s' % db)
		tp.start( process_one_db, args=(db, fcn, oname), kwargs=dict(lazy=lazy, skiperr=skiperr) )



if __name__ == '__main__':
	Lazy = False
	SkipErr = False
	Oname = None
	Oname2 = None
	Fcn = None
	open("LOG", "a").writelines( [ ' '.join(sys.argv), '\n'] )
	arglist = sys.argv[1:]
	while arglist and arglist[0].startswith('-'):
		arg = arglist.pop(0)
		if arg == '--':
			break
		elif arg == '-lazy':
			Lazy = True
		elif arg == '-skiperr':
			SkipErr = True
		elif arg == '-o':
			Oname = arglist.pop(0)
		elif arg == '-Fsimple2':
			Fcn = simple2
			Oname2 = 'HTKsimple2.mfcc'
		elif arg == '-Fsimple3':
			Fcn = simple3
			Oname2 = 'HTKsimple3.mfcc'
		elif arg == '-Fsimple3u':
			Fcn = simple3u
			Oname2 = 'HTKsimple3u.mfcc'
		elif arg == '-FsimpleP1':
			Fcn = simpleP1
			Oname2 = 'HTKsimpleP1.mfcc'
		elif arg == '-FsimpleP2':
			Fcn = simpleP2
			Oname2 = 'HTKsimpleP2.mfcc'
		elif arg == '-Fsimple4':
			Fcn = simple4
			Oname2 = 'HTKsimple4.mfcc'
		elif arg == '-Fsimple4s':
			Fcn = simple4s
			Oname2 = 'HTKsimple4s.mfcc'
		elif arg == '-Ffv':
			Fcn = fv
			Oname2 = 'HTKfv.mfcc'
		elif arg == '-Fmfcc':
			Fcn = mfcc
			Oname2 = 'HTKmfcc.mfcc'
		elif arg == '-Fmfcp':
			Fcn = mfccplus
			Oname2 = 'HTKmfcp.mfcc'
		else:
			die.die("Unrecognized flag: %s" % arg)
	if not Fcn:
		die.die("Need to specify -Fsomething flag")
	if not Oname:
		Oname = Oname2
	if not Oname:
		die.die("Need to specify -o name flag")
	process_all(arglist, Fcn, Oname, lazy=Lazy, skiperr=SkipErr)
