import Num
import gpkimgclass
import gpkavg
from oxigen_exceptions import NotEnoughData, UnNormalizable
import die
import oxigen_signals
import math
import os.path

ADJUSTF0 = 1
SLOPE_DT = 0.5


def set_silence_frac(f):
	"""Call before read_X_weight()."""
	oxigen_signals.SILENCE_FRAC = f


def datacol_base(dcol):
	return dcol.split(':')[0]

def datacol_flags(dcol):
	a = dcol.split(':')
	if len(a) < 2:
		return []
	return a[1].split(',')


NORMTYPE = {'F0': '/-',
		'ssl': '-', 'lpdur': '-',
		'pkpk': '^2/', 'rms': '^2/', 'loudness': '^2/',
		'pdur': '/', 'irregularity': '/',
		'bvoicing': '',
		'voicing': '-', 'spectral_slope': '-',
		'peakiness': '-'
		}


def per_phrase_weight(ws, dflags):
	for x in dflags:
		if x.startswith('twt'):
			TWTF = float(x[3:])/ws.shape[0]
			print 'TWTF=', TWTF
			assert len(ws.shape) == 1
			assert ws.shape[0] > 10
			ws = ws * (1.0+TWTF*Num.arrayrange(ws.shape[0]))
			break
	return ws


def per_phrase_norm(ds, ws, datacol):
	CLIP = 0.05

	dflags = datacol_flags(datacol)
	ws = per_phrase_weight(ws, dflags)
	if not 'ppn' in dflags:
		return (ds, ws)

	datacol = datacol_base(datacol)

	if NORMTYPE[datacol] == '/-':
		try:
			dsa, dss = gpkavg.avg(ds, ws, CLIP)
			# We assume it's already been normalized
			# over the entire corpus, so 1 has been
			# subtracted.
			dsa += 1
		except ValueError, x:
			if 'Nonpositive weight sum' in str(x):
				raise NotEnoughData, x
		print 'NORMALIZING %s by %g' % (datacol, dsa)
		if dsa <= 0.0:
			raise UnNormalizable, "Normalizing %s; dsa=%g" % (datacol, dsa)
		ds = ds / dsa
			# We assume it's already been normalized
			# over the entire corpus, so 1 has been
			# subtracted; we don't need to do it again.

	elif NORMTYPE[datacol] == '-':
		try:
			dsa, dss = gpkavg.avg(ds, ws, CLIP)
		except ValueError, x:
			if 'Nonpositive weight sum' in str(x):
				raise NotEnoughData, x
		print 'NORMALIZING %s by %g' % (datacol, dsa)
		ds = ds - dsa

	elif NORMTYPE[datacol] == '^2/':
		try:
			dsa, dss = gpkavg.avg(ds**2, ws, CLIP)
		except ValueError, x:
			if 'Nonpositive weight sum' in str(x):
				raise NotEnoughData, x
		print 'NORMALIZING %s by %g' % (datacol, dsa)
		if dsa <= 0.0:
			raise UnNormalizable, "Normalizing %s; dsa=%g" % (datacol, math.sqrt(dsa))
		ds = ds / math.sqrt(dsa)

	elif NORMTYPE[datacol] == '/':
		try:
			dsa, dss = gpkavg.avg(ds, ws, CLIP)
		except ValueError, x:
			if 'Nonpositive weight sum' in str(x):
				raise NotEnoughData, x
		print 'NORMALIZING %s by %g' % (datacol, dsa)
		if dsa <= 0.0:
			raise UnNormalizable, "Normalizing %s; dsa=%g" % (datacol, dsa)
		ds = ds / dsa
	elif NORMTYPE[datacol] == '':
		pass
	else:
		die.die('Unrecognized normtype: %s from %s' % NORMTYPE[datacol], datacol)

	return (ds, per_phrase_weight(ws, dflags))



def pun_norm(ds, ws, dflags, dt):
	"""This supplies an estimate of a normalization function
	for the data.   It assumes you're listening to a loop,
	and that the normalization function is a weighted average
	of your recent history.     Specifically, the weight
	is ws (the intrinsic weight of each datum) times
	exp(-time_delay/tau).
	"""

	assert ds.shape == ws.shape
	n = ds.shape[0]
	assert(ws.shape == (n,))

	Ww = None
	for fl in dflags:
		if fl.startswith('pun'):
			Ww = float(fl[3:])
			assert Ww > 0.0
	if Ww is None:
		# Take this exit if there is no "pun" flag.
		return (None, None)

	# This is what you'd get for a circular overlap
	# of an exponential function (neglecting a constant factor).
	tw = Num.exp(-(n-Num.arrayrange(n))*(dt/Ww))

	ftw = Num.FFT.real_fft(tw)
	w_convolved = Num.FFT.inverse_real_fft(Num.FFT.real_fft(ws)*ftw).real
	assert Num.alltrue(Num.greater(w_convolved, 0.0))
	wdconvolved = Num.FFT.inverse_real_fft(Num.FFT.real_fft(ws*ds)*ftw).real
	assert w_convolved.shape == ws.shape
	d_o = wdconvolved/w_convolved
	wbar = Num.sum(wd_convolved)/Num.sum(w_convolved)
	return (d_o, w_convolved/wbar)



def per_utterance_norm(d, datacol, prefix=None):
	"""Normalizes the data on a per-utterance basis.
	The details of the normalization are derived
	from flags in the datacol name."""

	if prefix is None:
		prefix = "norm_"
	dflags = datacol_flags(datacol)
	datacol = datacol_base(datacol)
	ds = d.column('%s%s' % (prefix, datacol))
	ws = d.column('Weight')
	ws = ws/Num.sum(ws)
	dsa, wsa = pun_norm(ds, ws, dflags, d.dt())
	if dsa is None:
		return d

	if NORMTYPE[datacol] == '/-':
		# We assume it's already been normalized
		# over the entire corpus, so 1 has been
		# subtracted.
		dsa += 1
		if not Num.alltrue(dsa > 0.0):
			raise UnNormalizable, "Normalizing %s" % datacol
		nds = ds / dsa
			# We assume it's already been normalized
			# over the entire corpus, so 1 has been
			# subtracted; we don't need to do it again.

	elif NORMTYPE[datacol] == '-':
		nds = ds - dsa

	elif NORMTYPE[datacol] == '^2/':
		if not Num.alltrue(dsa > 0.0):
			raise UnNormalizable, "Normalizing %s" % datacol
		nds = ds / dsa

	elif NORMTYPE[datacol] == '/':
		if not Num.alltrue(dsa > 0.0):
			raise UnNormalizable, "Normalizing %s" % datacol
		nds = ds / dsa
	elif NORMTYPE[datacol] == '':
		pass
	else:
		die.die('Unrecognized normtype: %s from %s' % NORMTYPE[datacol], datacol)
	EPS = 1e-10
	nws = ws * wsa / (ws + wsa + EPS)
	wbar = Num.sum(nws)/wsx.shape[0]
	nws /= wbar
	d.attach_col(nws, 'Weight')
	d.attach_col(nds, '%s%s' % (prefix, datacol))
	return d


def get_normalization(group, datacol, Valid):
	norminfo = []
	for info in group:
		if not Valid.data(info):
			continue
		if int(info['nspeakers']) > 1:
			# It's hard to normalize a dialog.  You'd have
			# to split it up, and normalize all the pieces.
			# Instead, we depend on every dialog speaker
			# having also done a solo performance, and
			# just use the solos in normalization.
			# This has the added benefit of making it easier
			# to compare solo and duet performances.
			continue

		try:
			norminfo.append( get_norm_f1(os.path.join(info['d'], info['f']), datacol) )
		except gpkimgclass.CannotReadDataFile, x:
			die.warn('Cannot Read Data file: %s' % x)
			continue
		except gpkimgclass.NoSuchColumn, x:
			die.warn('No such column: %s' % str(x))
			continue
		except NotEnoughData, x:
			die.warn('Not Enough Data: %s' % x)
			continue
		except:
			die.catch('get_normalization: unknown exception')
			raise
	if len(norminfo) == 0:
		raise NotEnoughData, 'When about to combine normalizations'
	return combine_normalizations(norminfo)



def get_norm_f1(fn, datacol):
	CLIP = {'F0': 0.10, 'spectral_slope': 0.10,
		'ssl': 0.20, 'ssY': 0.20,
		'voicing': 0.10, 'irregularity': 0.35,
		'rms': 0.10, 'pkpk':0.10, 'loudness':0.10,
		'pdur':0.10, 'lpdur':0.10,
		'bvoicing':0.10,
		'peakiness': 0.20
		}

	datacol = datacol_base(datacol)
	d = read_X_weight(fn, datacol)
	wt = d.column('Weight')

	if NORMTYPE[datacol] == '^2/':
		ctr, sigma = gpkavg.avg(d.column(datacol)**2, wt, CLIP[datacol])
		ctr = math.sqrt(ctr)
		sigma = 0.5*sigma/ctr
	else:
		ctr, sigma = gpkavg.avg(d.column(datacol), wt, CLIP[datacol])

	print 'File:', fn, 'ctr=', ctr, 'sigma=', sigma
	return ( Num.array([ctr, sigma]), 1.0 )




def read_X_weight(fn, datacol):
	datacol = datacol_base(datacol)
	if datacol == 'F0':
		d = oxigen_signals.read_f0_weight(fn, adjustf0=ADJUSTF0,
							conservative_v=0)
	elif datacol == 'bvoicing':
		d = oxigen_signals.read_bvoicing_weight(fn, conservative_v=0)
	elif datacol == 'voicing':
		d = oxigen_signals.read_gpkv_weight(fn)
	elif datacol == 'irregularity':
		d = oxigen_signals.read_irr_weight(fn)
	elif datacol == 'spectral_slope':
		d = oxigen_signals.read_ss_weight(fn)
	elif datacol == 'ssY':
		d = oxigen_signals.read_ssY_weight(fn)
	elif datacol == 'ssl':
		d = oxigen_signals.read_ssl_weight(fn)
	elif datacol == 'peakiness':
		d = oxigen_signals.read_peakiness_weight(fn)
	elif datacol == 'pdur':
		d = oxigen_signals.read_pdur_weight(fn)
	elif datacol == 'lpdur':
		d = oxigen_signals.read_lpdur_weight(fn)
	elif datacol == 'rms':
		d = oxigen_signals.read_rms_weight(fn)
	elif datacol == 'pkpk':
		d = oxigen_signals.read_pkpk_weight(fn)
	elif datacol == 'loudness':
		d = oxigen_signals.read_loudness_weight(fn)
	else:
		die.die('Bad datacol= %s' % datacol)

	assert d.n[1] >= 2 and d.n[1]<10
	return d


def normalize(d, norm, datacol, prepend=None):
	assert d.n[1] >= 1 and d.n[1]<10
	if prepend is None:
		prepend = "norm_"

	if NORMTYPE[datacol] == '/-':
		nf = norm[0]
		d.d[:,d.types[datacol]] = d.column(datacol) / nf - 1
	elif NORMTYPE[datacol] == '':
		pass	# No normalization
	elif NORMTYPE[datacol] == '-':
		nf = norm[0]
		d.d[:,d.types[datacol]] = d.column(datacol) - nf
	elif NORMTYPE[datacol] == '/' or NORMTYPE[datacol] == '^2/':
		nf = norm[0]
		d.d[:,d.types[datacol]] = d.column(datacol) / nf
	else:
		die.die('Bad NORMTYPE=%s datacol= %s' % (NORMTYPE[datacol], datacol))

	if prepend is not None and prepend != '':
		d.register('%s%s' % (prepend, datacol), d.types[datacol])

	assert d.n[1] >= 1 and d.n[1]<10
	return d


def antinormalize(d, norm, datacol, prefix=None):
	assert d.n[1] >= 1 and d.n[1]<10

	if prefix is None:
		prefix = "norm_"
	pdc = "%s%s" % (prefix, datacol)

	if NORMTYPE[datacol] == '/-':
		d.d[:,d.types[pdc]] = (d.column(pdc)+1) * norm[0]
	elif NORMTYPE[datacol] == '':
		pass	# No normalization
	elif NORMTYPE[datacol] == '-':
		d.d[:,d.types[pdc]] = d.column(pdc) + norm[0]
	elif NORMTYPE[datacol] == '/' or NORMTYPE[datacol] == '^2/':
		d.d[:,d.types[pdc]] = d.column(pdc) * norm[0]
	else:
		die.die('Bad NORMTYPE=%s datacol= %s' % (NORMTYPE[datacol], datacol))

	if prefix is not None and prefix != '':
		d.register(datacol, d.types[pdc])
		
	assert d.n[1] >= 1 and d.n[1]<10
	d.add_history( 'fft_norm.antinormalize %s' % str(norm) )
	return d


def read_X_weight_norm1(fn, norm, datacol, root=None, prepend=None):
	assert norm is not None
	if root is not None:
		fn = os.path.join(root, fn)
	datacol = datacol_base(datacol)

	d = read_X_weight(fn, datacol)
	return normalize(d, norm, datacol, prepend=prepend)


def combine_normalizations(normlist):
	CLIP = 0.25
	vv = []
	ww = []
	for (v, wt) in normlist:
		vv.append(v)
		ww.append(wt)
	vv = Num.array(vv)
	assert vv.shape[0] == len(ww)
	o = Num.zeros((vv.shape[1],), Num.Float)
	for i in range(vv.shape[1]):
		o[i] = gpkavg.avg(Num.array(vv[:,i]), Num.array(ww), CLIP)[0]
	return o



TIME_PER_POLY = {
		'F0': 0.110,
		'loudness': 0.060,
		'spectral_slope': 0.090,
		'ssl': 0.090, 'ssY': 0.090,
		'voicing': 0.080, 'bvoicing': 0.090,
		'irregularity': 0.080,
		'rms': 0.060,
		'pkpk': 0.060,
		'pdur': 0.070,
		'lpdur': 0.070,
		'peakiness': 0.090
		}
	

def time_per_poly(datacol):
	return TIME_PER_POLY[datacol_base(datacol)]
