"""This is code that, when given a non-negative time-series
corresponding to speech loudness, computes the background noise,
which is assumed to be constant.
It uses a simple model based on a Gaussian and exponential."""

import math
import numpy
import fpconst
from gmisclib import die
from gmisclib import load_mod
from gmisclib import Numeric_gpk

# import pylab

class ZeroProblem(RuntimeError):
	def __init__(self, *s):
		RuntimeError.__init__(self, *s)


PASSLIM = 80

def exp(x):
	return numpy.exp(numpy.maximum(x, -50.0))

def find_zero(x, debug=None):
	"""Models the loudness as being drawn either from a Gaussian
	distribution (corresponding to the background noise level)
	or a positive exponential distribtion (corresponding to speech).
	It returns parameters of these distributions.

	@param x: a time-series of loudness.
	@return: four floats: (1) the estimate of the background noise level,
		(2) the width of the noise Gaussian, (3) the mean speech level,
		(4) the fraction of time with no speech (just background noise).
		return (zz, sz, e, Pz)
	@rtype: tuple(float, float, float, float)
	"""
	zz = Numeric_gpk.N_minimum(x)
	sigma2 = numpy.var(x)
	if not (sigma2 > 0.0):
		raise ZeroProblem, "Loudness zero has sigma=zero -- seemingly no signal here."
	sz2 = sigma2/64.0
	sv2 = sigma2 * 2.0
	#: fz is the estimated fraction of data that are at zero loudness, i.e. part of the
	#: 	Gaussian centered at zero.
	fz = 0.5
	ipass = 0
	q = 1000.0
	while ipass<PASSLIM and q>0.0001:
		assert 0.0 <= fz <= 1.0
		ozz = zz
		osz2 = sz2
		osv2 = sv2
		ofz = fz

		gzz = numpy.greater(x, zz)
		assert sv2 > 0.0
		assert sz2 > 0.0
		wv = gzz * exp(-numpy.square(x-zz)/(2*sv2))/math.sqrt(2*math.pi*sv2)
		wz = exp(-numpy.square(x-zz)/(2*sz2))/math.sqrt(2*math.pi*sz2)
		#: flatfudge is a fudge factor that expresses the idea that the "zero" regions
		#: are fairly flat.
		flatfudge = exp(-numpy.square(x[1:]-x[:-1])/(32*sz2))
		numpy.multiply(wz[1:], flatfudge, wz[1:])
		numpy.multiply(wz[:-1], flatfudge, wz[:-1])
		assert 0.0 < fz < 1.0
		pz = fz*wz / (fz*wz + (1-fz)*wv)

		speed = math.exp(-2.5*ipass/float(PASSLIM))
		fz = (1-speed)*fz + speed*min(0.99, max(0.01, numpy.average(pz)))
		pv = 1.0 -pz
		sumPv = numpy.sum(pv)
		sumPz = numpy.sum(pz)
		if fpconst.isNaN(sumPz) or fpconst.isNaN(sumPv):
			raise ZeroProblem, "NaN while computing zero: sumPz sumPv: ipass=%d" % ipass
		zz = (1-speed)*zz + speed * numpy.sum(pz*x) / sumPz
		sz2 = max(1e-4*sigma2, (1-speed)*sz2 + speed*numpy.sum(pz*numpy.square(x-zz))/sumPz)
		sv2 = max(sz2, (1-speed)*sv2 + speed*numpy.sum(pv*numpy.square(x-zz))/sumPv)
		if debug is not None:
			debug(x, pz, zz, fz, sz2, sv2, ipass)

		q = (ozz-zz)**2/(sz2+osz2) + math.log(sz2/osz2)**2 + math.log(sv2/osv2)**2 + math.log(fz/ofz)**2
		ipass += 1
	if debug is not None:
		debug(x, pz, zz, fz, sz2, sv2, -1)
	sz = math.sqrt(sz2)
	sv = math.sqrt(sv2)
	if ipass >= PASSLIM:
		raise ZeroProblem, "Loudness zero does not converge: %g %g %g %g" % (zz, sz, sv, fz)
	if sv<3*sz or zz<0.10*sz or zz>sv/2.0 or fz>0.9:
		raise ZeroProblem, "Strange loudness zero: %g %g %g %g" % (zz, sz, sv, fz)
	return (zz, sz, sv, fz)


def loud_zero(x, debug=None):
	"""@param x: a time-series of loudness.
	@return: An estimate of the background noise level
	@rtype: float
	"""
	zz, sz, e, Pz = find_zero(x, debug=debug)
	return zz


def zero_sub(x):
	return numpy.maximum(x**3 - loud_zero(x)**3, 0.0)**(1./3.)
	

def percep_spec_zero(x):
	assert len(x.shape)==2 and x.shape[0]<60
	loud = numpy.sum(x, axis=0)
	zz, sz, e, Pz = find_zero(loud)
	fz = Pz.sum()
	return numpy.sum(x*Pz[numpy.newaxis,:], axis=1)/fz


def diagnostic_plot(x, ex):
	pylab = load_mod.load_named('pylab')
	for i in range(x.shape[0]):
		pylab.plot(x[i,:])
	pylab.title(str(ex))
	pylab.show()


if __name__ == '__main__':
	import sys
	import pylab
	import gpkimgclass
	d = gpkimgclass.read(sys.argv[1])
	x = d.column(0)/numpy.average(d.column(0))
	def debug(x, pz, zz, fz, sz2, sv2, i):
		if debug < 0:
			return
		pylab.figure(i)
		pylab.plot(x, 'b-')
		pylab.plot(pz, 'r-')
		pylab.title('zz=%g fz=%.3f sz=%.3f sv=%.3f, i=%d' % (zz, fz, math.sqrt(sz2), math.sqrt(sv2), i))
	try:
		zz, sz, sv, pz = find_zero(x, debug)
	except ZeroProblem, x:
		print "Zero problem: %s" % x
	pylab.show()
