#! python

"""This computes a set of commands to HHed that will add or
subtract mixtures from the current recognizer.
"""

import os
import math
import numpy
from gmisclib import avio
from gmisclib import gpkmisc
from gmisclib import HTK_HMM_io

def getstats(fname):
	o = {}
	nstates  = {}
	for (i, l) in enumerate(open(fname, 'r')):
		a = l.strip().split()
		assert int(a[0]) == i+1
		assert len(a[1])>2 and a[1].startswith('"') and a[1].endswith('"')
		phone = a[1][1:-1]
		states = a[3:]
		ns = len(states)
		for istate in range(2, ns+2):
			o[(phone, istate)] = float(states[istate-2])
		nstates[phone] = ns
	return (o, nstates)


def getwatch(fname):
	hrest, comments = avio.read(open(fname, 'r'))
	rv = {}
	for d in hrest:
		if d['ltype'] == 'coefficient':
			rv[d['phone']] = float(d['coefficient'])
	return rv


def count_mixes(hmms):
	rv = {}
	for ph in hmms.phones:
		for (istate,s) in enumerate(ph.states):
			if s is not None:
				rv[(ph.name, istate)] = len(s.mx)-1
	return rv


def volume(s):
	"""@return: the log of the phase-space volume of a state, i.e. the
		discriminant of the covariance matrix.
	"""
	vols = []
	for m in s.mx:
		if m is not None:
			v = HTK_HMM_io.deref(m.v).v
			vols.append( numpy.sum(numpy.log(v)) )
	return gpkmisc.median(vols)



def volumes(hmms):
	"""Compute phase space log(volume) for all states in a HMM.
	@rtype: dict(tuple(str,str): float)
	"""
	rv = {}
	for ph in hmms.phones:
		for (istate, s) in enumerate(ph.states):
			if s is not None:
				rv[(ph.name,istate)] = volume(s)
	return rv



def go(tj):
	N_PER_MIX = 100
	# Read the stats information produced by HERest.  This tells
	# how often the HMM used each state.
	stats, nstates = getstats(os.path.join(tj, "stats.out"))

	#: What is the correlation between the number of occurrences of a
	#: phone and log(P)?  Negative correlations suggest that the phone is
	#: not well modeled and therefore needs more flexibility.
	coefs = getwatch(os.path.join(tj, "watch.av"))

	#: Read in the current HMM.
	hmms = HTK_HMM_io.read([open(os.path.join(tj, "acoustics.mmf"), 'r')])
	nmixes = count_mixes(hmms)
	vol = volumes(hmms)
	vbar = gpkmisc.median(vol.values())
	scores = []
	for ((ph, istate), occ) in stats.items():
		#: How many mixes do we currently have?
		mxc = nmixes[(ph,istate)]
		#: What's the maximum number of mixes that are tolerable?
		nmax = int(math.ceil(occ/N_PER_MIX))
		if mxc > nmax:
			print "MU %d {%s.state[%d].mix}" % (nmax, ph, istate)
			continue

		#: Score marks the states that need/can support more mixtures.
		#: It gets bigger if the state is heavily used (thus more data to support more mixtures).
		#: It gets bigger if the coefficient is negative (i.e. the phone is not well modelled).
		#: It gets bigger if the state has a large volume (thus is ripe for splitting into more mixes).
		score = math.log(occ/(N_PER_MIX*mxc)) - 0.5*coefs[ph] + 0.1*(vol[(ph,istate)]-vbar)
		if score > 0.0:
			newmixes = min(nmax, mxc + 1 + int(math.floor(0.4*mxc)))
			if newmixes > mxc:
				scores.append( (score, ph, istate, newmixes) )
		if score < 0.0 and mxc > 1:
			newmixes = mxc-1
			NegF = 0.5
			# Let's be reluctant to make a big (fractional) reduction in the number
			# of mixtures.  So, scale down the "score".   Then scale it down further
			# because we don't want a lot of backtracking.
			scores.append( (-NegF*score*float(newmixes)/float(mxc), ph, istate, newmixes) )
	scores.sort()
	for (sc, ph, istate, newmixes) in scores[len(scores)//2:]:
		print "MU %d {%s.state[%d].mix}" % (newmixes, ph, istate)


if __name__ == '__main__':
	import sys
	go(sys.argv[1])
