# -*- coding: utf-8 -*-
import math
import numpy
from gmisclib import HTK_HMM_io as HMM


class probseq(object):
	def __init__(self, p):
		self.p = p
		self.a = 0.0
	
	def __call__(self):
		self.a += self.p
		if self.a > 1:
			self.a -= 1
			return True
		return False
	
	

def makeHMM(length, width, nf, nsideways=0.0, nskip=0.0, nbranch=0.0,
		nmix=1, name=None, sself=0.0, nback=0.0):
	"""This makes a cylinder of states.    Transitions are allowed as follows:
		* forwards (always)
		* sideways to the left or right (controlled by nsideways)
		* Diagonally forwards (controlled by nbranch)
		* Skipping directly forwards (controlled by nskip)
		* Looping transitions (controlled by sself).
		* Backwards transitions (nback)
	
	@note: If C{sself==0} and C{nsideways < width-1} then you can enforce
		both a minimum and maximum duration for the HMM.    The minimum duration
		is about C{length} (if C{nskip==0}) or C{length/2} with skips.
		The maximum duration is about C{nsideways*length}, with no skips,
		and about half of that with skips turned on.
	@param length: the length of the cylinder (in states)
	@param width: the circumference of the cylinder (in states)
	@param nf: the length of each acoustic feature vector
	@param nsideways: the number of sideways transitions in each ring.
		If C{nsideways <= width-1}, all the transitions are leftwards;
		for larger values (up to C{2*(width-1)}) there are also rightwards transitions.
	@param nskip: the number of skip transitions per ring.
	@param nback: the number of backwards transitions per ring.   Backwards transitions
		are equally likely to be to any of the previous rings.    Strictly speaking,
		this number applies to the final ring: earlier rings have proportionally
		lower probabilities.
	@param nbranch: the number of forward branching transitions per ring.
		This can be as large as C{width-1}, which connects each state with
		all the states in the next ring.
	@param sself: the strength of the looping transition.
	@param nmix: How many mixtures per state?
	@param name: the name of the state.
	"""
	SFWD = 0.3
	SSKIP = 0.1
	SBRANCH = 0.1
	SSIDE = 0.1
	SBACK = 0.05
	assert nsideways <= 2*(width-1)
	assert nskip <= width
	assert nbranch <= width-1
	assert sself >= 0.0
	hmm = HMM.hmm(numstates=length*width+2, name=name)
	transitions = []
	
	def index(i, j):
		if i < 0:
			return 1
		elif i >= length:
			return 2 + width*length
		return 2 + width*i + (j+width)%width
	
	has_lsideways = probseq(min(width-1, nsideways)/float(width))
	has_rsideways = probseq(max(0.0, nsideways-width+1)/float(width))
	has_skip = probseq(nskip/float(width))
	has_back = probseq(nback/float(width*length))
	has_branch = probseq( nbranch - math.floor(nbranch))
	
	for j in range(width):
		ij = index(-1, 0)
		transitions.append( (ij, index(0, j), SFWD) )
		if has_skip():
			transitions.append( (ij, index(1, j), SSKIP) )
	for i in range(length):
		for j in range(width):
			s = HMM.state(name="%s%d,%d" % (name, i, j))
			for k in range(1, nmix+1):
				s.create_mx(k)
				s.mx[k].m = HMM.vec('u', numpy.zeros((nf,)))
				s.mx[k].v = HMM.vec('v', numpy.ones((nf,)))
				s.mx[k].check()
			s.check()
			ij = index(i, j)
			if has_lsideways():
				transitions.append( (ij, index(i,(j+1)%width), SSIDE) )
			if has_rsideways():
				transitions.append( (ij, index(i,(j-1)%width), SSIDE) )
			if has_skip():
				transitions.append( (ij, index(i+2, j), SSKIP) )
			for k in range(0, i):
				if has_back():
					transitions.append( (ij, index(k, j), SBACK) )
			transitions.append( (ij, index(i+1, j), SFWD) )
			transitions.append( (ij, ij, sself) )
			inbr = int(math.floor(nbranch))
			for m in range(inbr):
				transitions.append( (ij, index(i+1, j+m), SBRANCH) )
			if has_branch():
				transitions.append( (ij, index(i+1, j+inbr), SBRANCH) )
			hmm.states.append(s)
	ns = index(length+1, 0) + 1
	hmm.t = numpy.zeros((ns,ns))
	for (s, e, w) in transitions:
		hmm.t[s, e] = w
	numpy.divide(transitions[1:], transitions[1:].sum(axis=1), transitions[1:])
	hmm.check()
	return hmm
