import gpkimgclass
import tsops
import Num as N
import weighted_percentile
import die

# import sys
# sys.path.append('/home/gpk/oxi/lib')
import voicing_marks
from oxigen_exceptions import NotEnoughData
from gpkimgclass import CannotReadDataFile, NoSuchColumn

SILENCE_FRAC = 0.05
WT_EXP = 1.0

def max0(d):
	return N.maximum(0, d)


def SetWtExp(x):
	print "SETTING WT_EXP to", x
	global WT_EXP
	WT_EXP = x

# This sets the scale for ignoring single wild points.
FJUMP = 10.0 # Hz

def compute_wt1(pv, f, delta):
	assert pv.shape[0] == f.shape[0]
	wt1 = N.ones((pv.shape[0],), N.Float)
	deltaffac = delta**2/((f[1:]-f[:-1])**2 * pv[1:] * pv[:-1] + delta**2)
	N.multiply(wt1[:-1], deltaffac, wt1[:-1])
	N.multiply(wt1[1:], deltaffac, wt1[1:])
	deltavfac = 1/((pv[1:]-pv[:-1])**2 + 1.0)
	N.multiply(wt1[:-1], deltavfac, wt1[:-1])
	N.multiply(wt1[1:], deltavfac, wt1[1:])
	return wt1


def read_f0_correct(fbase, adjustf0=1, conservative_v=1, f0_is='pitch'):
	d = gpkimgclass.read('%s.f0.dat' % fbase)
	# print "DTYPES1=", d.types
	assert d.column('prob_voice').shape == (d.n[2],)
	wt1 = compute_wt1(d.column('prob_voice'), d.column('F0'), FJUMP)
	assert wt1.shape == (d.n[2],)
	f0, valid, irr = voicing_marks.f0correct(
				'%s.weirdf0' % fbase,
				d.time(),
				d.column('F0'),
				d.column('prob_voice'),
				required=0,
				conservative_v = conservative_v,
				f0_is = f0_is,
				adjustf0=adjustf0
				)
	o = gpkimgclass.gpk_img(d.hdr,
			N.transpose([f0, valid, irr, wt1])
			)
	o.register('F0', 0)
	o.register('valid', 1)
	o.register('irr', 2)
	o.register('wt1', 3)
	return o


def read_f0_weight(fbase, adjustf0=1,
			conservative_v=1, f0_is='pitch',
			silence_frac=None):
	d = read_f0_correct(fbase, adjustf0, conservative_v, f0_is)
	wt = d.column('valid') * d.column('wt1') / (1 + d.column('irr'))
	l = read_loudness(fbase, silence_frac=silence_frac)
	li = tsops.interp_fill(l.select_cols(['loudness']), d.time(), 0.0)
	irr = tsops.apply(max0, gpkimgclass.read('%s.irr.dat' % fbase))
	ii = tsops.interp_fill(irr.select_cols(['aperiodicity']), d.time(), 0.0)
	lci = li * N.maximum(1-ii, 0.0)
	irr = None	# Just recycle memory.
	l = None

	# print "swt1", N.sum(wt > 0.00001)
	# print 'tsoifs', tsops.interp_fill(lc, d.time(), 0.0).shape
	# print "slci", N.sum(lci > 0.00001)
	# print 'lc time range', lc.start(), lc.end(), 'd range', d.start(), d.end(), d.time()[0], d.time()[-1]
	assert lci.shape[1] == 1
	# print 'types=', type(wt), type(lci)
	N.multiply(wt, lci[:,0], wt)
	d.d[:,d.types['wt1']] = wt**(2*WT_EXP)
	d.register('Weight', d.types['wt1'])
	# print "DTYPES=", d.types

	return d



def read_bvoicing_weight(fbase, conservative_v=1, silence_frac=None):
	d = read_f0_correct(fbase, 0, conservative_v, 'f0')
	d.register('bvoicing', d.types['valid'])
	d.d[:,d.types['wt1']] = N.ones((d.n[2],), N.Float)
	d.register('Weight', d.types['wt1'])
	return d




def read_loudness(fbase, silence_frac = None):
	"""This subtracts off an estimate for the background noise
	from the loudness measurements."""
	if silence_frac is None:
		silence_frac = SILENCE_FRAC
	assert silence_frac>=0.0 and silence_frac<1.0
	E = 1.0/3.0
	die.note('name', fbase)
	x = gpkimgclass.read('%s.loud.dat' % fbase)
	if silence_frac == 0:
		return x
	ld = x.column('loudness')
	if silence_frac >= 0.5:
		die.warn('read_loudness: percentile=%g > 0.5: t0=%g, tf=%g'
				% (silence_frac, x.start(), x.end()))
		raise NotEnoughData, 'file %s is too short' % fbase
	lpct = weighted_percentile.wp(ld, None, [ silence_frac ] )[0]
	ldc = ( N.maximum(0.0, ld**(1.0/E) - lpct**(1.0/E)) )**E
	x.d[:,x.types['loudness']] = ldc
	return x


def read_loudness_weight(fbase, silence_frac = None):
	d = read_loudness(fbase, silence_frac)
	assert d.n[1] == 1
	wt = N.ones((d.n[2],), N.Float)
	d.attach_col(wt, ttype='Weight')
	assert d.n[1] == 2
	return d


def read_ss_weight(fbase, silence_frac=None):
	d = read_f0_correct(fbase)

	s = gpkimgclass.read('%s.ss.dat' % fbase)
	ss = s.column('spectral_slope')

	vi = tsops.interp_fill(d.select_cols(['valid']), s.time(), 0.0)

	l = read_loudness(fbase, silence_frac=silence_frac)
	li = tsops.interp_fill(l.select_cols(['loudness']), s.time(), 0.0)

	lc = li * vi

	assert len(lc.shape)==2 and lc.shape[1] == 1
	# print 'types=', type(wt), type(lci)
	o = gpkimgclass.gpk_img(s.hdr, N.transpose([ss, lc[:,0]**(2*WT_EXP)]))
	o.register('spectral_slope', 0)
	o.register('Weight', 1)
	return o


def read_ssY_weight(fbase, silence_frac=None):
	"""Read spectral slope, but use a weight function that is not restricted
	to voiced regions.
	"""

	s = gpkimgclass.read('%s.ssl.dat' % fbase)
	ss = s.column('spectral_slope')

	l = read_loudness(fbase, silence_frac=silence_frac)
	lc = tsops.interp_fill(l.select_cols(['loudness']), s.time(), 0.0)

	assert len(lc.shape)==2 and lc.shape[1] == 1
	# print 'types=', type(wt), type(lci)
	o = gpkimgclass.gpk_img(s.hdr, N.transpose([ss, lc[:,0]**(2*WT_EXP)]))
	o.register('ssY', 0)
	o.register('Weight', 1)
	return o



def read_ssl_weight(fbase, silence_frac=None):
	d = read_f0_correct(fbase)

	s = gpkimgclass.read('%s.ssl.dat' % fbase)
	ss = s.column('spectral_slope')

	vi = tsops.interp_fill(d.select_cols(['valid']), s.time(), 0.0)

	l = read_loudness(fbase, silence_frac=silence_frac)
	li = tsops.interp_fill(l.select_cols(['loudness']), s.time(), 0.0)

	lc = li * vi

	assert len(lc.shape)==2 and lc.shape[1] == 1
	# print 'types=', type(wt), type(lci)
	o = gpkimgclass.gpk_img(s.hdr, N.transpose([ss, lc[:,0]**(2*WT_EXP)]))
	o.register('ssl', 0)
	o.register('Weight', 1)
	return o




def read_gpkv_weight(fbase, silence_frac=None):
	x = gpkimgclass.read('%s.gpkv.dat' % fbase)
	gpkv = N.maximum(x.column('voicing'), 0.0)

	l = read_loudness(fbase, silence_frac=silence_frac)
	li = tsops.interp_fill(l.select_cols(['loudness']), x.time(), 0.0)


	assert len(li.shape)==2 and li.shape[1] == 1
	o = gpkimgclass.gpk_img(x.hdr, N.transpose([gpkv, li[:,0]**(2*WT_EXP)]))
	o.register('voicing', 0)
	o.register('Weight', 1)
	return o


def read_peakiness_weight(fbase, silence_frac=None):
	x = gpkimgclass.read('%s.pkns.dat' % fbase)
	pkns = x.column('peakiness')

	l = read_loudness(fbase, silence_frac=silence_frac)
	li = tsops.interp_fill(l.select_cols(['loudness']), x.time(), 0.0)


	assert len(li.shape)==2 and li.shape[1] == 1
	o = gpkimgclass.gpk_img(x.hdr, N.transpose([pkns, li[:,0]**(2*WT_EXP)]))
	o.register('peakiness', 0)
	o.register('Weight', 1)
	return o



def read_irr_weight(fbase, silence_frac=None):
	x = gpkimgclass.read('%s.irr.dat' % fbase)
	irr = N.maximum(x.column('aperiodicity'), 0.0)

	l = read_loudness(fbase, silence_frac=silence_frac)
	li = tsops.interp_fill(l.select_cols(['loudness']), x.time(), 0.0)


	assert len(li.shape)==2 and li.shape[1] == 1
	o = gpkimgclass.gpk_img(x.hdr, N.transpose([irr, li[:,0]**(2*WT_EXP)]))
	o.register('irregularity', 0)
	o.register('Weight', 1)
	return o



def read_pdur_weight(fbase, silence_frac=None):
	x = gpkimgclass.read('%s.pdur.dat' % fbase)
	pdur = x.column('pseudoduration')

	l = read_loudness(fbase, silence_frac=silence_frac)
	li = tsops.interp_fill(l.select_cols(['loudness']), x.time(), 0.0)


	assert len(li.shape)==2 and li.shape[1] == 1
	o = gpkimgclass.gpk_img(x.hdr, N.transpose([pdur, li[:,0]**(2*WT_EXP)]))
	o.register('pdur', 0)
	o.register('Weight', 1)
	return o


def read_lpdur_weight(fbase, silence_frac=None):
	x = gpkimgclass.read('%s.pdur.dat' % fbase)
	lpdur = N.log(x.column('pseudoduration'))

	l = read_loudness(fbase, silence_frac=silence_frac)
	li = tsops.interp_fill(l.select_cols(['loudness']), x.time(), 0.0)


	assert len(li.shape)==2 and li.shape[1] == 1
	o = gpkimgclass.gpk_img(x.hdr, N.transpose([lpdur, li[:,0]**(2*WT_EXP)]))
	o.register('lpdur', 0)
	o.register('Weight', 1)
	return o



def read_rms_weight(fbase, silence_frac = None):
	"""This subtracts off an estimate for the background noise
	from the RMS measurements."""
	if silence_frac is None:
		silence_frac = SILENCE_FRAC
	assert silence_frac>=0.0 and silence_frac<1.0
	E = 1.0/2.0
	die.note('name', fbase)
	x = gpkimgclass.read('%s.rms.dat' % fbase)
	ld = x.column('rms')
	if silence_frac > 0.0:
		lpct = weighted_percentile.wp(ld, None, [ silence_frac ] )[0]
		ldc = ( N.maximum(0.0, ld**(1.0/E) - lpct**(1.0/E)) )**E
	else:
		ldc = ld

	o = gpkimgclass.gpk_img(x.hdr, N.transpose([ldc, N.ones(ldc.shape, N.Float)]))
	o.register('rms', 0)
	o.register('Weight', 1)
	return o


def read_pkpk_weight(fbase, silence_frac = None):
	"""This subtracts off an estimate for the background noise
	from the RMS measurements."""
	if silence_frac is None:
		silence_frac = SILENCE_FRAC
	assert silence_frac>=0.0 and silence_frac<1.0
	E = 1.0/4.0
	die.note('name', fbase)
	x = gpkimgclass.read('%s.pkpk.dat' % fbase)
	ld = x.column('pkpk')
	if silence_frac > 0.0:
		lpct = weighted_percentile.wp(ld, None, [ silence_frac ] )[0]
		ldc = ( N.maximum(0.0, ld**(1.0/E) - lpct**(1.0/E)) )**E
	else:
		ldc = ld

	o = gpkimgclass.gpk_img(x.hdr, N.transpose([ldc, N.ones(ldc.shape, N.Float)]))
	o.register('pkpk', 0)
	o.register('Weight', 1)
	return o

