#! python

"""This code looks at the output of a HERest run and reports
suspiciously good matches.

It takes log(P) for the utterances and models it with
a linear regression of log(P) on the per-utterance phoneme frequencies.
(I.e. it assumes that some phonemes tend to increase log(P) for
an utterance and that others systematically reduce it.)
It prints these regression coefficients, which can tell you which
phonemes are able to represent the data well and which cannot.

Then, it looks at the residuals.  It prints warning messages about
utterances with unexpectedly small values of log(P) - the few
that are the worst fits.   These utterances perhaps are bad data,
or have bad transcriptions, or some other problem.

Finally, it returns a list of utterances that can be used in a future
run of HERest.   This list excludes a few of the bad fits.


Usage:
HERest ...... |
python watch_HERest.py labeldir traindir > new_list_of_files_to_train

Where labeldir tells it where to look for .lab files (it needs these to
find out how many of each phone are in each utterance).   Note that this
code does *not* know how to read MLF files.

Traindir just becomes the beginning of each filename that's written to
the standard output.
"""

import os
import re
import sys
import math
import numpy
from gmisclib import die
from gmisclib import gpkmisc
from gmisclib import avio
from gmisclib import dictops
from gmisclib import gpk_lsq
from gmisclib import weighted_percentile

#: DROPFRAC helps to set how many utterances are dropped because of bad log(P).
#: This is the maximum fraction of utterances that will be dropped.
DROPFRAC = 0.002

#: Zth helps to set how many utterances are dropped because of bad log(P)
#: We look at the residuals as a Gaussian distribution; Zth is the z-score
#: for dropping utterances.  (Though subject to limitation by DROPFRAC.)
Zth = 3.0

#: This is a pattern for picking file names from the HERest output
_cpn = re.compile('\s*Processing Data: (.+); Label:? (.+)[.]lab$')

#: This is a pattern for picking log(P) from the HERest output
_upf = re.compile('\s*Utterance prob per frame = ([0-9eE.+-]+)')


def collect(fd):
	"""This function reads output from HERest and extracts the useful
	bits of information.
	@type fd: file
	@param fd: file descriptor to read
	@return: Yields name and log(P)
	@rtype: sequence(tuple(str, float))
	"""
	name = None
	names = set()
	for (lnum, l) in enumerate(fd):
		# Normal operation of HERest:
		m = _cpn.match(l)
		if m:
			if name is not None:
				die.warn("Whoops! name is not None: %s %d" % (name, lnum+1))
			name = m.group(2)
			if name in names:
				die.warn("Duplicate name: %s %d" % (name, lnum+1))
				continue
			names.add(name)
		m = _upf.match(l)
		if m:
			assert name is not None, "Line %d" % (lnum+1)
			yield (name, float(m.group(1)))
			name = None

		# HERest is in trouble:
		if l.strip().startswith('WARNING'):
			if name is not None:
				die.info("Dropping %s because of %s" % (name, l.strip()))
			name = None
		if l.strip().startswith('ERROR'):
			if name is not None:
				die.die("ERROR in HERest %s: %s" % (name, l.strip()))
			name = None

		# This message is generated by run_several.py:
		if l.startswith('# stdout'):
			name = None

def read_labels(fn):
	"""Read the phoneme sequence from a label file.
	@param fn: name of the label file
	@type fn: str
	@return: yields a sequence of phonemes
	@rtype: sequence(str)
	"""
	for l in open(fn, 'r'):
		l = l.strip()
		if l:
			yield l


def analyze(fd, labdir, out=None):
	"""This is the main routine.   It prints warning messages on the standard
	error for dubious utterances.   It returns a list of "good" utterances.
	@param fd: where the data comes from
	@type fd: L{file}
	@param labdir: where to find the label files.
	@type labdir: str
	@rtype: str
	@return: yields a list of the names of good utterances.
	"""
	# Read stuff from HERest.   Get a first estimate of the
	# mean of log(P) and how much it varies from utterance to utterance.
	np = []
	plist = []
	for (name, prob) in collect(fd):
		plist.append(prob)
		np.append( (name, prob) )
	die.info("%d collected" % len(plist))
	median = gpkmisc.median(plist)
	stdev = gpkmisc.median_ad(plist) *  1.4826
	del plist

	# Read in label files and produce some warning messages.
	pset = set()
	npl = []
	for (name, prob) in np:
		if prob > median + 10*stdev + 10:
			die.warn("Wildly good probability: %g vs %g for %s" %
					(prob, median, name))
		if prob > median + 6*stdev:
			# Suspiciously good: don't use it in the linear
			# regression, but there's no obvious reason to drop
			# this file:
			yield name
			continue
		labels = read_labels(os.path.join(labdir, name) + '.lab')
		pc = dictops.dict_of_accums()
		for lbl in labels:
			pc.add(lbl, 1)
			pset.add(lbl)
		npl.append( (name, prob, pc) )
	del np
	# None will turn (magically/klugically) into the constant term
	# of the linear regression.
	pset.add(None)

	# Create a map from phoneme labels to integers.
	pimap = {}
	for (i,p) in enumerate(sorted(pset)):
		pimap[p] = i

	# Set up the linear regression:
	a = numpy.zeros((len(npl), len(pimap)))
	b = numpy.zeros((len(npl),1))
	die.info('Sizes: %s %s' % (a.shape, b.shape))
	for (i,(name, prob, nlab)) in enumerate(npl):
		b[i,0] = prob
		nlabtot = 0
		for n in nlab.values():
			nlabtot += n
		fac = math.log(nlabtot)/nlabtot
		for (p, n) in nlab.items():
			a[i,pimap[p]] = n * fac
		a[i,pimap[None]] = 1.0
	# Do the regression:
	soln = gpk_lsq.linear_least_squares(a, b, minsvr=1e-5)

	# Print the coefficients:
	if out is not None:
		out.comment("Produced by %s" % sys.argv[0])
		out.comment("Coefficients from linear regression of log(P) on")
		out.comment("log(utterance length) * fraction of phoneme X.")
	sx = soln.x()
	for p in sorted(pset):
		die.info("Coefficient of %s = %.2f" % (p, sx[pimap[p],0]))
		if out is not None:
			out.datum({'ltype': 'coefficient', 'phone': p, 'coefficient': sx[pimap[p],0]})

	# Set a threshold for rejecting utterances:
	r = numpy.absolute(soln.residual())
	assert r.shape[1]==1
	th1 = weighted_percentile.wp(r[:,0], None, [1.0-DROPFRAC])[0]
	threshold = max(th1, Zth*math.sqrt(math.pi/2.0)*numpy.average(r))
	die.info("median=%g threshold=%g" % (median, threshold))

	# Reject the bad ones:
	p_pred = soln.fit()
	for ((name, prob, pc), pp) in zip(npl, p_pred):
		# print "# pp:", prob, pp, pc
		if prob > pp - threshold:
			yield name
		else:
			die.info("Dropping utterance: bad prob=%g: %s" % (prob, name))


if __name__ == '__main__':
	arglist = sys.argv[1:]
	if arglist[0] == '-coef':
		coeffd = avio.writer(open(arglist[1], 'w'))
		arglist = arglist[2:]
	labdir = arglist[0]
	traindir = arglist[1]
	for n in analyze(sys.stdin, labdir, out=coeffd):
		print os.path.basename(os.path.join(traindir, n))
