#!/usr/bin/env python


import Num
import mcmc
import mcmc_helper
import gpkimgclass
import gpkmisc

EXTRASMOOTH = 0.040
W = 3
HARD_THRESHOLD = 0
MINW = 0.050	# THE SHORTEST POSSIBLE UTTERANCE, DEFAULT.
MINFRAC = 0.25	# The shortest utterance is 25% of the length of the data.

Fend = 0.70	# Height of loudness near the end, relative to near the start.
BACKGROUND = 0.05
Minsilence = 0.050

HUGE = 1e30
E = 1.0/0.31

class Whoops(Exception):
	def __init__(self, s):
		Exception.__init__(self, s)


def fit_fun(p, n, dt):
	"""Designed for loudness or cube-rooted RMS.
	"""

	Efrac = 0.0
	MAXRISE = 0.060
	MAXFALL = 0.150

	assert len(p) == 7
	assert n > 20
	istart = p[0]
	iend = p[1]
	if istart < 0 or iend>n-1:
		raise Whoops('istart<0 or iend>n-1: %g %g' % (istart, iend))
	if iend < istart:
		raise Whoops('iend < istart: %g %g' % (istart, iend))
	f_end = p[2]
	if f_end<0.2 or f_end>1.3:
		raise Whoops('f_end bad')
	i_rise = p[3]
	if i_rise <= 0.0 or i_rise>MAXRISE/dt:
		raise Whoops('i_rise bad')
	i_fall = p[4]
	if i_fall <= 0.0 or i_fall>MAXFALL/dt:
		raise Whoops('i_fall bad')
	if i_rise+i_fall > 0.5*(iend-istart):
		raise Whoops('rise and fall time long compared to utterance')
	bkgnd = p[5]
	if bkgnd <= 0.0 or bkgnd > 0.1:
		raise Whoops('background bad')
	if f_end <= bkgnd:
		raise Whoops('Fend is ridiculously small vs. background')
	maxval = p[6]
	if maxval <= 0.0:
		raise Whoops('maxval bad')
	i = Num.arrayrange(n)
	top = 1.0 - (i-istart)*((1.0-f_end)/(iend-istart))
	le = (i-istart)/i_rise + Efrac
	re = f_end*((iend-i)/i_fall + Efrac)
	speech = Num.maximum(0.0, Num.minimum(top, Num.minimum(re, le)))**E
	bkg = bkgnd**E
	return maxval*(bkg + speech)**(1.0/E)

def eval_err(p, (loud, dt, fac, minw)):
	"""This compares the observed above_thresh to a model
	which is one between istart and iend, and zero outside.
	That corresponds to a single region of voicing.
	"""
	istart = p[0]
	iend = p[1]
	i_rise = p[3]
	i_fall = p[4]
	fudges = 0.0
	assert len(loud.shape) == 1
	assert loud.shape[0] > 20
	c = 0.05*loud.shape[0]
	fudges += c*(minw/dt + 0.25*loud.shape[0])/(iend-istart)	# length too short
	fudges += c*(Minsilence/dt)/(istart+(loud.shape[0]-iend))	# too little silence
	try:
		fit = fit_fun(p, loud.shape[0], dt)
	except Whoops:
		return -HUGE
	residual = loud - fit
	rv = -Num.sum( Num.absolute(residual) )
	# print 'rv=', rv, fac, rv*fac
	return rv*fac - fudges


def maxval(x):
	return x[Num.argmax(x)]

def find_best_fit(loud, dt, minw):
	"""This finds the ends of the single best-fit region
	of speech.   The model assumes speech in the middle,
	and silence outside.
	"""
	assert len(loud.shape) == 1
	assert loud.shape[0] > 20, 'Loud.shape=%s' % str(loud.shape)
	xpeak = maxval(loud)
	x = loud/xpeak
	start = Num.array([0.1*x.shape[0], 0.9*x.shape[0],
				0.5,
				0.03/dt, 0.06/dt, 0.01, 1.0])
	V = gpkmisc.make_diag([0.3*(0.1*x.shape[0])**2, 0.3*(0.1*x.shape[0])**2,
				0.05,
				0.1*(0.1/dt)**2, 0.1*(0.1/dt)**2, 0.1*0.01**2, 0.01*1**2])
	fac = len(start)*10/Num.sum(Num.absolute(x))
	f1 = mcmc.bootstepper(eval_err, start, V, (x, dt, fac, minw))
	mcmch = mcmc_helper.threaded_stepper()
	mcmch.run_to_bottom(f1)
	fac = 3*len(start)*loud.shape[0]/Num.sum(Num.absolute(x))
	f2 = mcmc.bootstepper(eval_err, f1.current().prms(), V, (x, dt, fac, minw))
	mcmch.run_to_bottom(f2)
	p = Num.array(f2.current().prms())
	p[6] *= xpeak
	return p


def bestfit_to_start_end(bestfit, x):
	istart = bestfit[0]
	iend = bestfit[1]
	return (max(x.time(istart-0.5), x.start()),
		min(x.time(iend+0.5), x.end()))


def find_trim_times_rms(x, colname, xexp, minw=MINW):
	xx = x.column(colname)
	bestfit = find_best_fit( xx**xexp, x.dt(), minw)
	return bestfit_to_start_end(bestfit, x)




def find_ref_ampl(x, colname, tstart, tend):
	istart = int(round(x.t_index(tstart)))
	iend = int(round(x.t_index(tend)))
	return Num.sort(x.column(colname)[istart:iend])[-1]


def find_trim_times(fbase, xexp, suffix='/rms.dat', minw=MINW):
	x = gpkimgclass.read('%s%s' % (fbase, suffix))
	tstart, tend =  find_trim_times_rms(x, 'rms', xexp, minw=minw)
	ref_ampl = find_ref_ampl(x, 'rms', tstart, tend)
	return (tstart, tend, ref_ampl)


if __name__ == '__main__':
	import sys
	import die
	import xwaves_lab
	dtype = 'loud'
	dtin = 'loudness'
	E = 1.0
	DBG = False
	arglist = sys.argv[1:]
	labelfile = None
	while arglist and arglist[0].startswith('-'):
		arg = arglist.pop(0)
		if arg == '--':
			break
		elif arg == '-label':
			labelfile = arglist.pop(0)
		elif arg == '-D':
			DBG = True
		elif arg == '-RMS':
			dtype = 'rms'
			dtin = 'rms'
			E = 0.3
		else:
			die.die('Bad argument: %s' % arg)
	assert len(arglist) == 1
	# wav = gpkimgclass.read('%s/dat' % arglist[0])
	rms = gpkimgclass.read('%s/%s.dat' % (arglist[0], dtype))
	bestfit =  find_best_fit(rms.column(dtin)**E, rms.dt(), minw=MINW)
	if DBG:
		x = fit_fun(bestfit, rms.n[2], rms.dt()) * (maxval(rms.column(dtin))**E)
		rms.attach_col(x, 'Fit')
		rms.hdr['BITPIX'] = -32
		rms.write('foo.dat')
	if labelfile:
		tstart, tend = bestfit_to_start_end(bestfit, rms)
		d = [(tstart, '*'), (tend, '%')]
		hdr = {'t0': tstart, 'tf': tend,
			'Fend': bestfit[2], 't_rise': bestfit[3]*rms.dt(),
			't_fall': bestfit[4]*rms.dt(),
			'background': bestfit[5]*bestfit[6],
			'Fitted_max': bestfit[6],
			'datastart': rms.start(),
			'dataend': rms.end(),
			'startclip': tstart-rms.start(),
			'endclip': rms.end()-tend
			}
		xwaves_lab.write(open(labelfile, 'w'), hdr, d)
	else:
		print bestfit[0]*rms.dt(), bestfit[1]*rms.dt()
