#! python

"""This is a script for generating MMF files (Master Model Files) for the
HTK speech recognition toolkit.
It takes a simple description (e.g. L{TM}) and expands it in all the
messy glory that HTK needs for an initial MMF.

The description is taken in on the standard input, and the models are
written out to files named by their phoneme.   The shared macros are written
to a file called "macros".   You specify the directory for the output.

In the input language, there are five kinds of lines: int, float, macro, phone, and comment.

	- Comment lines are blank or begin with '#'.  They are ignored.

	- Int and float lines set named parameters.

		- int ndim sets the dimension of the feature vector

		- float eps sets the probability of stepping forward to the next state.
			(Eps defaults to 0.05.)

	- Macro lines name a collection of states so that they can be used several times.
		Note that re-use of a macro does not imply that state parameters are shared.

	- Phone lines define a phoneme.

General properties:

	- All lines end with a semicolon.

	- Spaces are separators.

	- Macros must be defined before they are used.   You cannot start using a
		name as a state, then later define it to be a macro.

	- States are named, for convenience.  However the name of a state makes
		no difference at all.   Even two states with identical names within
		the same macro or phoneme are treated as separate unless some
		flag is used to share properties.

	- After each state name, you can have a colon (C{:}), and then flags
		that set properties of the state.    Available flags are:

		- "+" which makes the state loop

		- "v=name" which sets up a shared/tied diagonal variance.

		- "V=name" which sets up a shared/tied full (InvCov) variance.

		- "m=name" which sets up a shared/tied mean.

		- "s=name" which sets up a shared/tied state.

		The names in the above flags can contain C{$s} that expands to the name of the state,
		C{$p} that expands to the name of the phoneme, or C{$m}
		that expands to the names of the innermost macro.
		When the "N*" form is used to repeat a state description, the C{$i} and C{$I}
		forms are also available.  These evaluate to an integer C{0...N-1} or
		C{1...N}, respectively.

		Note that only the "+" flag can be attached to uses of macros.  Sharing flags
		are nonsensical there.

Sharing example: These two phonemes start identically, but then diverge::

phone A = x:s=fred y;
phone B = x:s=fred z;

Or, equivalently::

phone A = x:s=$s y;
phone B = x:s=$s z;

Even after training, the first state will be identical between phones C{A} and C{B}.
If you train C{A}
data, you will also update the model for the initial state of C{B}, and vice versa.
In this example, all properties of the first state are shared: means, variances,
and even mixtures.

Sharing example: These two phonemes share the same variance everywhere, but
each state can have a different mean (this would still all be true even if the
states all had the same names)::

phone A = x1:v=fred x2:v=fred;
phone B = y1:v=fred y2:v=fred;



"""


#: Sample input
TM = """
# Set the dimension of the feature vector to 4:
int ndim = 4;

phone test1 = state;
phone test2 = state1 state2;
phone test3 = stateA | stateB ;
phone test4 = state:+ ;
phone test5 = stateA1 stateA2:+ | stateB:+ ;
macro testM = stateAm | stateBm ;
phone testMp = testM:+ ;

# The "silence" macro is defined here.  It branches into four paths:
# one for coughs, one for breathing, one for lip smacks, and one for true silence.
# (Of course, you cannot really assign particular meanings to these paths!
# HTK will make whatever use of them that it wants to.)
# Each path has a few states in it.   All those states are individually looped.
#
# The "v=cough" sets up a shared (tied) variance amongst all the states that
# refer to it.
macro silence = cough0:v=cough:+ cough1:v=cough:+ cough2:v=cough:+ cough4:v=cough:+
	| breath0:v=breath:+ breath1:v=breath:+ breath2:v=breath:+
	| smack0:v=smack:+ smack1:v=smack:+ smack2:v=smack:+
	| quiet:+ ;

# Start defining a silence phone for the end of an utterance.
# It consists of a single transition state, and then whatever states are
# in the "silence" macro.    The state transitions are set up to loop
# from the end of the "silence" macro back to the beginning, in addition
# to going onward.
phone _sil = tr_in silence:+ ;

phone sil_ = silence:+ tr_out ;

# This is a inter-word space.   It consists of either a single "glue" state
# called IWS, or a longer sequence which has a transition to silence,
# the silence macro (looped), and a transition out of silence.
phone _IWS_ = IWS | tr_in silence:+ tr_out ;


macro stop = tr_in closure:+ burst:+ tr_out ;
macro fricative = tr_in | frication:+ | tr_out ;
macro other = tr_in  stuff:+  tr_out ;
macro c = stop | fricative | other ;
phone C = c:+ ;
phone V = tr_in vocalic:+ tr_out ;
"""

import os
import re
import sys
import numpy
from gmisclib import die
from gmisclib import HTK_HMM_io


#: Helps in error detection and reporting:
RESERVED_WORDS = ['int', 'string', 'float', 'phone', 'macro']

#: Indents for printing.
I = '  '
_namepat = re.compile('^[a-zA-Z][a-zA-Z0-9_]*$')

def _or(a, b):
	if a is not None:
		return a
	elif b is not None:
		return b
	return None


class select_info(object):
	def __init__(self, phone=None, macro=None, state=None, istate=None, imix=None):
		self.phone = phone
		self.state = state
		self.macro = macro
		self.istate = istate
		assert imix is None or imix > 0, "Imix=%s" % str(imix)
		self.imix = imix

	def s(self, phone=None, macro=None, state=None, istate=None, imix=None):
		return select_info(_or(phone, self.phone), _or(macro, self.macro),
					_or(state, self.state), _or(istate, self.istate),
					_or(imix, self.imix)
				)

	def __repr__(self):
		tmp = []
		if self.phone is not None:
			tmp.append('ph=%s' % self.phone)
		if self.state is not None:
			tmp.append('st=%s' % self.state)
		if self.macro is not None:
			tmp.append('mac=%s' % self.macro)
		if self.istate is not None:
			tmp.append('i=%s' % self.istate)
		if self.imix is not None:
			tmp.append('mx=%s' % self.imix)
		return "[Selector %s]" % ( ' '.join(tmp))


def search_previous_macros(previous, selector, need):
	if selector.macro is None:
		return None
	rv = None
	if need == 's':
		rv = previous.s_macros.get(selector.macro, None)
	if need == 'u':
		rv = previous.u_macros.get(selector.macro, None)
	if need == 'v':
		rv = previous.v_macros.get(selector.macro, None)
	if need == 'i':
		rv = previous.i_macros.get(selector.macro, None)
	if rv is not None:
		die.info("Found matching macro from previous: %s %s" % (need, selector.macro))
	return rv


def search_previous_tree(previous, selector, need):
	bestscore = 0
	beststate = None
	besthmm = None
	for hmm in previous.phones:
		score = 0
		for (i, state) in enumerate(hmm.states):
			if state is None:
				continue
			if selector.phone==hmm.name:
				score += 8
			if(selector.state==state.fullname or selector.state==state.name):
				score += 3
			elif(state.name is not None and selector.state is not None
					and selector.state.startswith(state.name)):
				score += 2
			elif i==selector.istate:
				score += 1
			if score > bestscore:
				bestscore = score
				beststate = state
				besthmm = hmm
	if beststate is None:
		return None
	if need is not None:
		die.info("Found match to previous tree %s %s" % (need, selector))
	if need == 's':
		return beststate
	if need == 't':
		return besthmm.t
	tmp = beststate.mx[(_or(selector.imix, 1)-1) % (len(beststate.mx)-1) + 1]
	if need == 'u':
		# return tmp.m + numpy.random.normal(0.0, 0.01, tmp.m.shape)*numpy.sqrt(tmp.v)
		return tmp.m
	if need == 'v':
		return tmp.v
	if need == 'i':
		die.die("Sorry. Search for 'i' is not yet implemented.")
	return None


def search_previous(previous, selector, need):
	if previous is None or selector is None:
		return None
	tmp = search_previous_macros(previous, selector, need)
	if tmp is not None:
		return tmp
	return search_previous_tree(previous, selector, need)


def write_mean(ndim, previous, prev_select):
	o = []
	p = search_previous(previous, prev_select, 'u')
	if p is None:
		# p = HTK_HMM_io.vec('u', numpy.random.lognormal(0.0, 0.2, (ndim,)))
		p = HTK_HMM_io.vec('u', numpy.ones((ndim,)))
	tmp = p.write()
	o.append(I*5 + ' '.join(tmp))
	return o


def write_variance(ndim, isfull, previous, prev_select):
	o = []
	if not isfull:
		p = search_previous(previous, prev_select, 'v')
		if p is None:
			# p = HTK_HMM_io.vec('v', numpy.random.lognormal(0.0, 0.01, (ndim,)))
			p = HTK_HMM_io.vec('v', numpy.ones((ndim,)))
		tmp = p.write()
	else:
		p = search_previous(previous, prev_select, 'v')
		if p is None:
			# p = HTK_HMM_io.vec('i', numpy.random.lognormal(0.0, 0.01, (ndim,ndim)))
			p = HTK_HMM_io.vec('i', numpy.ones((ndim,ndim)))
		tmp = p.write()
	o.append(' '.join(tmp))
	return o



class context_c(object):
	def __init__(self):
		self.macros = {}
		self.states = {}
		self.ints = {}
		self.strings = {'SOURCEKIND': 'USER'}
		self.floats = {'eps': 0.05}
		self.shared_v = set()
		self.shared_m = set()
		self.shared_s = set()
		self.globalstatemap = []

	def prep_globalstatemap(self):
		for (phone, i, state) in self.globalstatemap:
			yield '%s %d %s %s' % (phone, i, state.name, state.fullname)



def state_factory(s, context, pname=None, mname=None, irep=None):
	a = s.strip().split(':')
	name = a[0].strip()
	if name in RESERVED_WORDS:
		die.die('Illegal name for a state: "%s".  Forgot a semicolon?' % name)
	if not name[0].isalpha():
		die.die('State name must start with letter: "%s".' % name)
	flags = a[1:]
	if name in context.macros:
		for flag in flags:
			if flag.startswith('P='):
				pname = flag[2:]
		if '+' in flags:
			tmp = looper_c(context.macros[name].realize(pname=pname), context)
		else:
			tmp = context.macros[name].realize(pname=pname)
	else:
		tmp = state_c(name, flags, context, pname=pname, mname=mname, irep=irep)
		if 'v' in tmp.shared:
			context.shared_v.add(tmp.shared['v'])
		if 'm' in tmp.shared:
			context.shared_m.add(tmp.shared['m'])
		if tmp.s_macro:
			context.shared_s.add(tmp.s_macro)
		if '+' in flags:
			tmp = looper_c(tmp, context)
	return tmp
	

_re = re.compile('[$][psmiI]')
def sharename(s, name, pname=None, mname=None, irep=None):
	repl = {'$s': name, '$p': '', '$m': '', '$i': '', '$I': ''}
	if pname is not None:
		repl['$p'] = pname
	if mname is not None:
		repl['$m'] = mname
	if irep is not None:
		repl['$i'] = '%d' % irep
		repl['$I'] = '%d' % (irep+1)
	a, nm = s.split('=', 1)
	return _re.sub(lambda x: repl[x.group()], nm)


class V_c(object):
	def __init__(self, name, fullV):
		self.name = name
		self.fullV = fullV
	
	def symbol(self):
		return 'vi'[self.fullV]

	def __hash__(self):
		return hash( (self.name, self.fullV) )

	def __eq__(self, other):
		return self.name==other.name and self.fullV==other.fullV

	def __repr__(self):
		return "[V %s(%s)]" % (self.name, ['Diag', 'InvCov'][self.fullV])


def get_int(s, context):
	s = s.split('=')[1].strip()
	if s.startswith('$'):
		return context.ints[s[1:]]
	return int(s)



class state_c(object):
	def __init__(self, name, flags, context, pname=None, mname=None, irep=None):
		self.fullV = False
		if not _namepat.match(name):
			die.die("Bad name for state: '%s'." % name)
		self.name = name
		self.s_macro = None
		self.shared = {}
		self.nmix = 1
		for flag in flags:
			if flag.startswith('v='):
				self.shared['v'] = V_c(sharename(flag, name=name, pname=pname, mname=mname, irep=irep), False)
			elif flag.startswith('V='):
				self.shared['v'] = V_c(sharename(flag, name=name, pname=pname, mname=mname, irep=irep), True)
			elif flag.startswith('m='):
				self.shared['m'] = sharename(flag, name=name, pname=pname, mname=mname, irep=irep)
			elif flag.startswith('s='):
				self.s_macro = sharename(flag, name=name, pname=pname, mname=mname, irep=irep)
			elif flag == 'fullV':
				self.fullV = True
			elif flag == '+':
				pass
			elif flag.startswith('nmix='):
				self.nmix = get_int(flag, context)
				assert self.nmix > 0
			else:
				die.die('Unrecognized modifier for state named "%s": "%s"' %
						(name, flag))
		self.context = context
		self.fullname = sharename('=$s$i@$m', name=name, pname=pname, mname=mname, irep=irep)


	def __repr__(self):
		tmp = [self.fullname]
		for (k, v) in self.shared.items():
			tmp.append('%s=%s' % (k, v))
		if self.s_macro:
			tmp.append('s=%s' % self.s_macro)
		return ':'.join(tmp)


	def list(self):
		yield self


	def entrypoints(self):
		return [self]

	def tr(self, prev):
		ends = set([self])
		o = set()
		f = self.context.floats['eps']/len(prev)
		for p in prev:
			o.add( (p, self, f) )
		print 'state_c tr=%d from prev=%d*1, ends=%d' % (len(o), len(prev), len(ends))
		return (o, ends)


	def write(self, istate, previous, prev_select):
		if self.s_macro:
			return [I*2 + '<State> %d' % (istate+1), I*3 + '~s "%s"' % self.s_macro]

		return write_state(self.context.ints['ndim'], self.shared,
					previous, prev_select.s(istate=istate, state=self.fullname),
					istate=istate, fullV=self.fullV, nmix=self.nmix)


	def __hash__(self):
		return id(self)

	def __eq__(self, other):
		return self is other


def write_state(ndim, shared, previous, prev_select, fullV=False, istate=None, nmix=1):
	"""This constructs the text that describes a state in the phoneme.
	@param ndim: dimension of the feature vector
	@type ndim: int
	@param shared: information about what is shared with what
	@type shared: C{dict}
	"""

	o = []
	if istate is None:
		o.append(I*2 + '<NumMixes> %d' % nmix)
	else:
		o.append(I*2 + '<State> %d <NumMixes> %d' % (istate+1, nmix))
	for imix in range(nmix):
		o.append(I*3 + '<Mixture> %d 1' % (imix+1)) 
		if 'm' in shared:
			o.append(I*4 + '~u "%s"' % shared['m'])
		else:
			o.extend(write_mean(ndim, previous, prev_select.s(istate=istate, imix=imix+1)))
		if 'v' in shared:
			sv = shared['v']
			o.append(I*4 + '~%s "%s"' % (sv.symbol(), sv.name))
		else:
			o.extend(write_variance(ndim, fullV, previous, prev_select.s(istate=istate, imix=imix+1)))
	return o




class seq_c(object):
	"""This class represents a sequence of states within a phoneme model."""

	_reppat = re.compile('([0-9]+)[*]')

	def __init__(self, s, context, pname=None, mname=None):
		self.s = []
		reps = 1
		for a in s.strip().split():
			m = self._reppat.match(a)
			if m:
				reps = int(m.group(1))
			else:
				for i in range(reps):
					self.s.append(state_factory(a, context, pname=pname, mname=mname, irep=i))
				reps = 1


	def entrypoints(self):
		return self.s[0].entrypoints()

	def list(self):
		"""List all the states in the sequence."""
		for item in self.s:
			for state in item.list():
				yield state

	def tr(self, prev):
		prev0 = prev
		o = set()
		for s in self.s:
			items, prev = s.tr(prev)
			o |= items
		print 'seq_c tr=%d from prev=%d*1, ends=%d' % (len(o), len(prev0), len(prev))
		return (o, prev)


class looper_c(object):
	def __init__(self, s, context):
		self.s = s
		self.context = context

	def list(self):
		for state in self.s.list():
			yield state

	def entrypoints(self):
		return self.s.entrypoints()

	def tr(self, prev):
		o = set()
		items, ends = self.s.tr(prev)
		o |= items
		entrypoints = self.s.entrypoints()
		f = (1-self.context.floats['eps'])/len(entrypoints)
		for p in entrypoints:
			for e in ends:
				o.add( (e, p, f) )
		return (o, ends)

	def __repr__(self):
		return '%s:+' % self.s.__repr__()


class or_c(object):
	"""This class represents a set of alternatives within a phoneme model."""

	def __init__(self, s, context, pname=None, mname=None):
		self.s = [ seq_c(a, context, pname=pname, mname=mname)
				for a in s.strip().split('|')
			]


	def list(self):
		"""List all the states within."""
		for seq in self.s:
			for state in seq.list():
				yield state

	def entrypoints(self):
		o = []
		for q in self.s:
			o.extend(q.entrypoints())
		return o


	def tr(self, prev):
		"""Compute state transition probabilities."""
		o = set()
		ee = set()
		for s in self.s:
			items, ends = s.tr(prev)
			o |= items
			ee |= ends
		return (o, ee)



class macro_c(object):
	"""This class stores a macro that can be used to build a HMM."""

	def __init__(self, s, context):
		assert '=' in s
		a = s.split('=', 1)
		if len(a) < 2:
			die.die("Macro definition line without any '='.")
			die.info("Line=%s" % s)
		name, code = a
		self.name = name.strip()
		if not _namepat.match(self.name):
			die.die("Bad name for a macro: '%s'" % self.name)
		if self.name in context.states:
			die.die('"%s" is already used as a state. It cannot be redefined as a macro.' % self.name)
		self.code = code
		self.context = context


	def realize(self, pname=None):
		return or_c(self.code, self.context, mname=self.name, pname=pname)



class hmm_c(object):
	"""This class represents a phone in a HMM."""

	def __init__(self, name, code, context):
		self.name = name.strip()
		if not _namepat.match(self.name):
			die.die("Bad name for a phone: '%s'." % self.name)
		self.code = or_c(code, context, pname=name)
		self.context = context


	def list(self):
		"""List all the states in this HMM."""

		for state in self.code.list():
			yield state


	def write(self, previous, prev_select):
		o = []
		statemap = {0:0}
		nstates = 1
		for state in self.list():
			if state not in statemap:
				statemap[state] = nstates
				self.context.globalstatemap.append( (self.name, nstates, state) )
				o.extend(state.write(nstates, previous, prev_select.s(phone=self.name)))
				nstates += 1
		nstates += 1
		tr = self._make_transition_matrix(nstates, statemap, previous, prev_select.s(phone=self.name))
		return [twiddle_o(self.context),
			'~h "%s"' % self.name, '<BeginHMM>',
			I*1 + '<NumStates> %d' % nstates] + o + tr + ['<EndHMM>']


	def _make_transition_matrix(self, nstates, statemap, previous, prev_select):
		m = search_previous(previous, prev_select, 't')
		wrongsize = m is not None and m.shape != (nstates, nstates)
		if wrongsize:
			die.warn("Transition matrix doesn't match number of states.: %d vs %s" % (nstates, m.shape))
		if m is None or wrongsize:
			m = numpy.zeros((nstates,nstates))
			items, ends = self.code.tr(set([0]))
			for (i, j, wt) in sorted(items, cmp=lambda a, b: cmp(statemap[a[0]], statemap[b[0]])):
				print 'tr %s(%d) -> %s(%d) = %.3f' % (i, statemap[i]+1, j, statemap[j]+1, wt)
				m[statemap[i],statemap[j]] = wt
			for i in ends:
				print 'tr %s(%d) -> [end](%d) = %.3f' % (i, statemap[i]+1, nstates, self.context.floats['eps'])
				m[statemap[i],nstates-1] = self.context.floats['eps']
		else:
			die.info("Found transition matrix: %s" % prev_select)
		tr = ['<TransP> %d' % nstates]
		for i in range(nstates):
			trtmp = m[i,:]
			trts = trtmp.sum()
			if trts > 0.0:
				trtmp = numpy.around(1000*trtmp/trtmp.sum())/1000.0
				trtmp = numpy.around(2000*trtmp/trtmp.sum())/2000.0
				trtmp = numpy.around(5000*trtmp/trtmp.sum())/5000.0
			tmp = []
			for j in range(nstates):
				tmp.append('%.4f' % trtmp[j])
			tr.append(I*1 + ' '.join(tmp) )
		return tr


def twiddle_o(context):
	ndim = context.ints['ndim']
	sk = context.strings['SOURCEKIND']
	# return '~o <STREAMINFO> 1 %d <VECSIZE> %d <NULLD><%s><DiagC>' % (ndim, ndim, sk)
	return '~o <STREAMINFO> 1 %d <VECSIZE> %d <NULLD><%s>' % (ndim, ndim, sk)


class hmmset_c(object):
	"""This class contains a whole set of Hidden Markov Models."""

	def __init__(self, s, phonelist=[], extralines=[]):
		self.hmm = []
		self.context = context_c()
		vlines = []
		current = []
		for (i,line) in enumerate(s.split('\n')):
			if line=='' or line.startswith('#'):
				if current:
					die.die("Comment or blank inside statement line %d.  Forgot semicolon at end?" % (i+1))
				continue
			a = line.split(';')
			while a:
				apop = a.pop(0).strip()
				if apop:
					current.append(apop)
				if current and a:
					tmp = ' '.join(current)
					if ' ' not in tmp:
						die.die("Bad syntax: %s: statement %d" % (tmp, i+1))
					vlines.append(tmp)
					current = []
		stl = vlines + extralines
		if current:
			die.die("No semicolon at end of file.")
		stl2 = []
		for line in stl:
			ltype, content = line.split(None, 1)
			if ltype == 'int':
				k, v = content.strip().split('=', 1)
				k = k.strip()
				v = int(v)
				self.context.ints[k] = v
				die.info("Setting int parameter %s to %d" % (k, v))
			elif ltype == 'float':
				k, v = content.strip().split('=', 1)
				k = k.strip()
				v = float(v)
				self.context.floats[k] = v
				die.info("Setting float parameter %s to %g" % (k, v))
			elif ltype == 'string':
				k, v = content.strip().split('=', 1)
				k = k.strip()
				v = v.strip()
				self.context.strings[k] = v
				die.info("Setting string parameter %s to '%s'" % (k, v))
			else:
				stl2.append( (ltype,content) )
		plist = set(phonelist)
		for (ltype,content) in stl2:
			if ltype == 'macro':
				h = macro_c(content, self.context)
				die.info("Defining macro %s" % h.name)
				self.context.macros[h.name] = h
			elif ltype == 'phone':
				try:
					name, code = content.split('=')
				except ValueError:
					die.info("Line: %s" % content)
					die.die("A 'phone' statement must contain exactly one equals sign")
				name = name.strip()
				if name not in plist:
					die.warn("Explicitly defined phoneme %s not in phonelist." % name)
				h = hmm_c(name, code, self.context)
				self.hmm.append(h)
				plist.discard(name)
				die.info("Defining phone %s" % h.name)
			elif ltype == 'phones':
				try:
					pat, code = content.split('=', 1)
				except ValueError:
					die.info('Line: "%s"' % content)
					die.die("'phones' line must contain an '='")
				pat = pat.strip()
				cpat = re.compile('(%s)$' % pat)
				phonelist_unmatched = set()
				for name in plist:
					m = cpat.match(name)
					if m:
						h = hmm_c(name, code, self.context)
						self.hmm.append(h)
						die.info("Defining2 phone %s" % h.name)
					else:
						phonelist_unmatched.add(name)
				if len(phonelist_unmatched) == len(plist):
					die.warn("No phones in phonelist match pattern %s" % pat)
				plist = phonelist_unmatched
			else:
				die.die("Unrecognized statement: %s %s" % (ltype, content))
		if plist:
			die.die("Cannot find a definition for these phones: %s" % plist)


	def write(self, d, previous=None):
		ndim = self.context.ints['ndim']
		o = [twiddle_o(self.context)]
		for x in self.context.shared_v:
			tmp = ['~%s "%s"' % (x.symbol(), x.name)]
			tmp.extend(write_variance(ndim, x.fullV, previous, select_info(macro=x.name)))
			o.extend(tmp)
		for x in self.context.shared_m:
			tmp = ['~u "%s"' % x]
			tmp.extend(write_mean(ndim, previous, select_info(macro=x)))
			o.extend(tmp)
		for x in self.context.shared_s:
			tmp = ['~s "%s"' % x]
			tmp.extend(write_state(ndim, {}, previous, select_info(macro=x)))
			o.extend(tmp)
		print_to(d, 'macros', o)
		for h in self.hmm:
			print '# phoneme: %s' % h.name
			print_to(d, h.name, h.write(previous, select_info(phone=h.name)))
		print_to(d, 'phones', [h.name for h in self.hmm])
		print_to(d, 'statemap.txt', self.context.prep_globalstatemap(),
				hdr=['From writeHMM.py', 'Phone state_number description']
			)


def print_to(d, f, stuff, hdr=[]):
	if d is None:
		fd = sys.stdout
	else:
		fd = open(os.path.join(d, f), 'w')
	for h in hdr:
		fd.write('# %s\n' % h)
	for s in stuff:
		fd.write('%s\n' % s)

def read_statemap(fn):
	sm = open(fn, 'r')
	statemap = {}
	for line in sm:
		if not line.startswith('#'):
			ls = line.split()
			if len(ls) == 3:
				ls.append(ls[-1])
			phone, i, name, fullname = ls
			i = int(i)
			statemap[(phone,i)] = (name, fullname)
	return statemap


def read_hmm(macros, phonelist):
	fdlist = [open(macros, 'r')]
	d = os.path.dirname(macros)
	for f in phonelist:
		fdlist.append(open(os.path.join(d, f), 'r'))
	rv = HTK_HMM_io.read(fdlist)
	sm = os.path.join(d, 'statemap.txt')
	try:
		statemap = read_statemap(sm)
	except IOError:
		statemap = {}
	for hmm in rv.phones:
		for (i,state) in enumerate(hmm.states):
			if state is not None:
				try:
					state.name, state.fullname = statemap[(hmm.name, i)]
				except KeyError:
					state.name = None
					state.fullname = None
	return rv


def nmix_f(fname, fac, exp):
	nmcol = 2
	freqcol = 3
	rv = {}
	for line in open(fname, 'r'):
		cols = line.split()
		name = cols[nmcol]
		if name[0]=='"' and name[-1]==name[0]:
			name = name[1:-1]
		freq = float(cols[freqcol])
		assert freq > 0.0
		rv[name] = 1+int(fac * freq**exp)
	return rv


def go(arglist):
	extralines = []
	phonelist = set()
	fd = sys.stdin
	previous = None
	while arglist and arglist[0].startswith('-'):
		arg = arglist.pop(0)
		if arg == '--':
			break
		elif arg == '-pre':
			previous = read_hmm(arglist.pop(0), phonelist)
		elif arg == '-i':
			fd = open(arglist.pop(0), 'r')
		elif arg == '-test':
			hmmset_c(TM, extralines=extralines, phonelist=phonelist).write(None, previous=previous)
			sys.exit(0)
		elif arg == '-x':
			extralines.append(arglist.pop(0))
		elif arg == '-p':
			phonelist = set(arglist.pop(0).split())
		else:
			die.die("Unrecognized flag: %s" % arg)
	assert len(arglist) == 1
	hmmset_c(fd.read(), extralines=extralines, phonelist=phonelist).write(arglist[0], previous=previous)


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