#!/usr/bin/env python

"""Feature vectors for optimized DTW alignment.
Exemplar project with Ladan Baghai-Ravary.
"""

import math
import numpy
from gmisclib import die
from gmisclib import cache as CC
from gmisclib import gpkmisc
from gmisclib import erb_scale
# from gpk_voicing import percep_spec as PS
from gpk_voicing import cached_ps as PS
from gpk_voicing import fv_misc as M
from gpk_voicing import zero as Z
SillyWidthException = M.SillyWidthException

DB = 0.95

Opt_text = """
# summarize_logs -uid UID -best -fromstart advdIrL,-fv201003opt_0.av
# len(currentlist)= 190 tail= 0.0 nsamp= 190
# max(advdIrL,-fv201003opt_0.av) = 84.47
#  samples used = 190 filename= advdIrL,-fv201003opt_0.av
84.47 logp
# n= 41
3.02401 Scale,fricative%Sedge:.*:11.9-14.8
14.2481 Scale,fricative%Sedge:.*:15.7-17.6
16.3675 Scale,fricative%Sedge:.*:18.6-23.3
0.925029 Scale,fricative%Sedge:.*:24.3-29.0
6.3387 Scale,fricative%Sedge:.*:4.3-7.2
22.348 Scale,fricative%Sedge:.*:8.1-11.0
0.798545 Scale,haspitch1:.*
3.12358 Scale,vowel:.*:10.0
1.17288 Scale,vowel:.*:11.0
1.01182 Scale,vowel:.*:11.9
2.37575 Scale,vowel:.*:12.9
14.409 Scale,vowel:.*:13.8
1.01614 Scale,vowel:.*:14.8
3.41617 Scale,vowel:.*:15.7
2.49094 Scale,vowel:.*:16.7
7.17409 Scale,vowel:.*:17.6
2.301 Scale,vowel:.*:18.6
3.45347 Scale,vowel:.*:19.5
5.5532 Scale,vowel:.*:20.5
1.77577 Scale,vowel:.*:21.4
0.365305 Scale,vowel:.*:22.4
5.62289 Scale,vowel:.*:23.3
2.07074 Scale,vowel:.*:24.3
0.757549 Scale,vowel:.*:25.2
2.91112 Scale,vowel:.*:26.2
3.47592 Scale,vowel:.*:27.1
1.79999 Scale,vowel:.*:28.1
10.2112 Scale,vowel:.*:29.0
1.37166 Scale,vowel:.*:4.3
1.8372 Scale,vowel:.*:5.3
0.646851 Scale,vowel:.*:6.2
0.480479 Scale,vowel:.*:7.2
12.4389 Scale,vowel:.*:8.1
6.54689 Scale,vowel:.*:9.1
0.0053143 border_shift,beginning
2.98189 dist,alpha
0.619647 dist,beta
1.35488 len_exp
0.468854 lf,elf
0.00879217 lf,lf
0.00784977 nsv
"""

Scale = M.scale_xform(Opt_text,
			name="2010-06-05 optimization in cayenne:/home/gpk/ItakuraSaitoDistance/IS-DTW/dvdIrL,-fv201003opt_0.av"
			)


def feature_vec_guts(data, dt, DT,
		LF=1.0, Nsv=M.NSV, ELF=1.0,
		do_voicing=0, do_dissonance=False,
		PlompBouman=False, do_pdur=False,
		zerosub=0.0, do_irx=False, E=None, cache_info=None):
	assert not do_irx, "Do_irx=%s" % do_irx
	assert not do_pdur and not do_dissonance
	FORMANT_LOW = erb_scale.f_to_erb(130.0)
	FORMANT_HIGH = erb_scale.f_to_erb(5000.0)
	assert float(DT)>0.0
	assert float(LF)>0.0
	bmin = erb_scale.f_to_erb(100.0)
	bmax = erb_scale.f_to_erb(5500.0)
	ectrs, neural, t0 = PS.perceptual_spec(data, dt, DT,
							bmin, bmax, DB,
							do_mod=do_voicing,
							do_dissonance=do_dissonance,
							do_peakalign=False,
							PlompBouman=PlompBouman,
							e=E,
							cache_info = cache_info.copy() if cache_info is not None else None
							)

	assert 1 < neural.shape[0] < 200
	band_indices = [i for (i,ec) in enumerate(ectrs) if ec['type']=='band']
	neural_b = neural.take(band_indices, axis=0)
	assert neural_b.shape[1]==neural.shape[1]
	if zerosub != 0.0:
		neural_b = neural_b.copy()
		try:
			z = Z.PS(neural_b)
		except Z.ZeroProblem:
			die.warn("Bad utterance: zero does not converge.")
			z = numpy.zeros((neural_b.shape[0],))
		assert neural_b.shape[0]==z.shape[0]
		numpy.subtract(neural_b, z[:,numpy.newaxis], neural_b)
	assert neural_b.shape[1]==neural.shape[1]
	assert neural_b.shape[0]<=neural.shape[0]

	neural_now = numpy.average(neural_b, axis=0)	# Average over frequency.
	assert neural_now.shape[0] == neural.shape[1]
	neural_avg = numpy.square(neural_now).sum()/numpy.sum(neural_now)	# Average over time.
		# neural_avg is a scalar, grand average.
	numpy.divide(neural, neural_avg, neural)
		# Now, we've normalized by an over-all average loudness.
	numpy.divide(neural_now, neural_avg, neural_now)
		# Now, we've normalized by an over-all average loudness.

	for (i,e) in enumerate(ectrs):
		if e['type'] == 'haspitch':
			numpy.divide(neural[i,:], neural_avg, neural[i,:])
		if e['type'] == 'dissonance':
			numpy.divide(neural[i,:], neural_avg, neural[i,:])
		if e['type'] == 'peakalign':
			numpy.divide(neural[i,:], numpy.square(neural_avg), neural[i,:])

	o = []
	descr = []
	w = 0.04*LF/DT
	tmpo, tmpd = M.vowel(w, ectrs, neural, neural_now, Nsv,
						formant_low=FORMANT_LOW,
						formant_high=FORMANT_HIGH
						)
	o.extend(tmpo)
	descr.extend(tmpd)
	assert len(descr)==len(o), "Descriptor mismatch"

	w = 0.04*ELF/DT
	tmpo, tmpd = M.fricative_edge(w, ectrs, neural, neural_now, Nsv,
						do_abs=False, version=2
						)
	o.extend(tmpo)
	descr.extend(tmpd)
	assert len(descr)==len(o), "Descriptor mismatch"

	if do_voicing:
		w = 0.02*math.sqrt(LF)/DT
		tmpo, tmpd = M.haspitch(w, ectrs, neural, neural_now, Nsv)
		o.extend(tmpo)
		descr.extend(tmpd)
	assert len(descr)==len(o), "Descriptor mismatch"

	w = 0.03*ELF/DT
	tmpo, tmpd = M.peakiness(w, ectrs, neural, neural_now, Nsv)
	o.extend(tmpo)
	descr.extend(tmpd)
	assert len(descr)==len(o), "Descriptor mismatch"

	N = neural[0].shape[0]
	for (i, (tmp, dsc)) in enumerate(zip(o, descr)):
		assert tmp.shape == (N,), "Wrong size: %d, %s = %d vs. %d" % (i, str(dsc), tmp.shape[0], N)

	return (o, descr, DT, t0)


def feature_vec(data, dt, DT,
		LF=1.0, Nsv=M.NSV, ELF=1.0,
		do_voicing=0, do_dissonance=False,
		PlompBouman=False, do_pdur=False,
		zerosub=0.0, do_irx=False, E=None, cache_info=None):
	o = None
	ci = None
	if cache_info is not None:
		assert isinstance(cache_info, CC.cache_info)
		ci = cache_info.addinfo(dt, DT, LF, Nsv, ELF, do_voicing, do_dissonance, PlompBouman,
				do_pdur, zerosub, do_irx, E, 'fv201003opt')
		try:
			o = ci.load()
		except ci.Errors:
			pass
	if o is None:
		o = feature_vec_guts(data, dt, DT, LF=LF, Nsv=Nsv, ELF=ELF, do_voicing=do_voicing,
					do_dissonance=do_dissonance, PlompBouman=PlompBouman,
					do_pdur=do_pdur, zerosub=zerosub, do_irx=do_irx, E=E,
					cache_info=cache_info
					)
		if ci is not None:
			ci.dump(o)
	return o
