#! python

"""This was used for Speech Prosody 2010."""

import os
import sys
import math
import random

import numpy

import gpkimgclass

from gmisclib import die
from gmisclib import mcmc
from gmisclib import avio
from gmisclib import stats
from gmisclib import tsops
from gmisclib import fiatio
from gmisclib import gpkmisc
from gmisclib import dictops
from gmisclib import g2_select
from gmisclib import g_localfit
from gmisclib import g_datetime as GD
from gmisclib import Numeric_gpk as NG
from gmisclib import wavesurfer_lab as WL
from gmisclib import weighted_percentile as WP
from gmisclib import blue_data_attributes as BD

from gpk_voicing import voice_misc

# from newstem2 import indexclass as IC

# from sloppy_oscillator import delay4 as O
from gpk_voicing import zero as Z


SEGMENTATION = 'recVapE_HTKmfcp.lab'
# MAXFIT = 12000
MAXFIT = 5000


def f_value(ER, EF, dfR, dfF):
    """Returns an F-statistic given the following:
        ER  = error associated with the null hypothesis (the Restricted model)
        EF  = error associated with the alternate hypothesis (the Full model)
        dfR = degrees of freedom the Restricted model
        dfF = degrees of freedom associated with the Restricted model
    """
    return ((ER-EF)/float(dfR-dfF) / (EF/float(dfF)))



class DubiousData(ValueError):
	def __init__(self, *s):
		ValueError.__init__(self, *s)


Map = {'sil': 'S', 'c': 'C', 'v': 'V', 'm': 'V'}

def remap(lbls):
	"""Maps from a sequence of labels produced by the ASR system
	to a sequence of consonantal and vocalic regions.
	"""
	rv = []
	last = ''
	ss = None
	ee = None
	for (s, e, lbl) in lbls:
		lbl = Map[lbl]
		if lbl != last:
			if last:
				rv.append( (ss, ee, last) )
			ss = s
			ee = e
			last = lbl
	if last:
		rv.append( (ss, ee, last) )
	return rv


class utterance_c(object):
	"""Information about an utterance."""

	def __init__(self, dbdir, info, segmentation):
		"""
		@param segmentation: last part of the pathname for the label files
			that define the consonantal and vocalic regions.
		@type segmentation: str
		"""
		self.info = info
		self.dbdir = dbdir
		self.segmentation = segmentation
		self.read_data(dbdir, info, segmentation)

	def read_data(self, dbdir, info, segmentation):
		"""
		@param info: metadata for an utterance.
		@param dbdir: folder containing the metadata file.
		@param segmentation: last part of the pathname for the label files
			that define the consonantal and vocalic regions.
		@type segmentation: str
		@except Z.ZeroProblem: if the background noise level cannot be determined.
		"""
		self.path = os.path.join(dbdir, info['d'], info['f'])
		die.info('Loading %s' % self.path)
		ld = gpkimgclass.read(os.path.join(self.path, 'loud.dat'))
		ird = gpkimgclass.read(os.path.join(self.path, 'irr.dat'))
		pdd = gpkimgclass.read(os.path.join(self.path, 'pdur.dat'))
		fvd = gpkimgclass.read(os.path.join(self.path, 'specfv.dat'))
		self.taxis, loud, irr, pdur, fv = tsops.common([ld, ird, pdd, fvd])
		self.pdur = pdur[:,pdd.types['pseudoduration']]

		assert fv.shape[0]==loud.shape[0]
		fvc = numpy.sum(fv*loud, axis=0)/numpy.sum(loud)
		fvs = numpy.std(fv, axis=0)
		assert fvs.shape[0] == fv.shape[1], "shape mismatch %s vs %s" % (fvs.shape, fv.shape)
		self.centering = numpy.sum(((fv-fvc)/fvs)**2, axis=1)
		assert self.centering.shape == (loud.shape[0],)
		
		loud = Z.zero_sub(loud[:,ld.types['loudness']])
		avloud = numpy.sum(loud**2)/numpy.sum(loud)
		self.lnorm = loud/avloud
		irr = irr[:,ird.types['aperiodicity']]
		irr_avg = numpy.sum(irr*loud)/numpy.sum(loud)
		self.fric = self.lnorm * (irr - irr_avg) * 2.0
		# The extra factor of 2.0 is so that the variance(fric) is similar to var(lnorm)
		
		h, lbls = WL.read(os.path.join(self.path, segmentation))
		lbls = remap(lbls)
		self.data = []
		if len(lbls) <= 3:
			die.warn('lbls too short: %d for %s / %s' % (len(lbls), self.path, segmentation))
			return
		# So, here we loop through all the syllables, collecting statistics
		# and stowing them away in self.data.    The syllables are marked
		# by the immediate context of the sound under investigation.
		for i in range(0, len(lbls)):
			tmp = self.compute_statistics(lbls[i][0], lbls[i][1],
							is_edge=(i==0 or i==len(lbls)-1))
			self.data.append(syldatum(lbls[i][2], tmp))


	def __repr__(self):
		return "<utterance_c: %s//%s//%s : %s >" % (
			self.dbdir, self.info['d'], self.info['f'],
			','.join([q.lbl for q in self.data])
			)



	def compute_statistics(self, vs, ve, is_edge):
		"""
		@param vs: the start of this phoneme
		@param ve: the end of this phoneme
		@type vs: float (time, seconds)
		@type ve: float (time, seconds)
		"""
		si = int(round(self.taxis.index(vs, limit=is_edge)))
		ei = int(round(self.taxis.index(ve, limit=is_edge)))
		if ei - si < 2:
			raise DubiousData("Too short")
		TYPDUR = 0.10
		dur = (ve-vs)/TYPDUR
		lbar = numpy.mean(self.lnorm[si:ei])
		fric = self.fric[si:ei+1]
		lnorm = self.lnorm[si:ei]
		fr = numpy.mean(fric)
		speechrate = math.log(numpy.sum(TYPDUR/self.pdur[si:ei])/dur)
		skew = (numpy.arange(ei-si) - 0.5*(ei-1-si)) * (self.taxis.dt()/TYPDUR)
		lskew = numpy.average(lnorm*skew)/(1.0+dur)
		lstd = numpy.std(lnorm)
		fstd = numpy.std(fric)
		centering = numpy.average(self.centering[si:ei])
		return [math.log(dur), lbar, speechrate, lskew, fr,
				lstd, fstd, centering]
	
	idx = {'duration':0, 'loud':1, 'speechrate':2, 'lskew':3, 'frication': 4,
			'loud_std': 5, 'fric_std': 6,
				'centering': 7
			}

	ridx = dictops.rev1to1(idx)

def inspector(info):
	return (info['gender'], info['textfile'], info['subjectID'], info['language'])



def read_and_select(dblist, selector, ofd):
	"""Given a list of metadata files, return a selected subset of the utterances.
	@param dblist: metadata file names
	@type dblist: list(str)
	@param selector: A little snippet of python code.
	@type selector: L{str}
	@return: a selected list of utterances
	@rtype: list(dict(str:str))
	"""
	assert len(dblist) > 0
	dur = []
	syl = []
	tmp = []
	for dbpath in dblist:
		dbd, dbf = os.path.split(dbpath)
		die.info("Reading metadata from %s" % dbpath)
		h, d, c = fiatio.read(open(os.path.join(dbd, dbf), 'r'))
		for info in g2_select.filterlist(selector, d):
			info['__dbdir'] = dbd
			ns = int(info['son_nas_vowel_exp'])
			syl.append(ns)
			info['__nsyls'] = ns
			d = GD.timedelta2float(GD.ISO2datetime(info['RecordEndTime'])
						- GD.ISO2datetime(info['RecordStartTime']))
			info['__duration'] = d
			dur.append(d)
			tmp.append(info)
	assert len(tmp) > 0
	dl, dh = WP.wp(dur, None, [0.1, 0.9])
	snl, snh = WP.wp(syl, None, [0.1, 0.9])
	ofd.write( '# Range of durations accepted: %s %s\n' % (dl, dh))
	infolist = []
	for info in tmp:
		if dl <= info['__duration'] <= dh :
			infolist.append(info)
	return infolist




def read_in_utterances(dblist, selector, segmentation, ofd):
	"""
	@return: a list of utterances
	@rtype: list(utterance_c)
	"""
	ulist = []
	for info in read_and_select(dblist, selector, ofd):
		try:
			ulist.append(utterance_c(info['__dbdir'], info, segmentation))
		except Z.ZeroProblem, x:
			die.warn("Dropping %s/%s because of ZeroProblem.%s" % (info['d'], info['f'], str(x)))
		except DubiousData, x:
			die.warn("Dropping %s/%s because of DubiousData.%s" % (info['d'], info['f'], str(x)))
		except WL.NoSuchFileError, x:
			die.warn("Dropping %s/%s because of NoSuchFile.%s" % (info['d'], info['f'], str(x)))
	return ulist


def allbut(a, *i):
	return [a for (a,j) in enumerate(a) if j not in i]


MINSV = None
MINSVR = 1e-5

class syldatum(object):
	__slots__ = ['lbl', 'data']

	def __init__(self, lbl, data):
		self.lbl = lbl
		self.data = data

	def ok(self):
		assert self.data and self.lbl




def sample_people_with_replacement(ulist):
	"""Sample data by subject-sized blocks.   The final output list contains
	either all of a given subject's data or none.

	@type ulist: list(utterance_c)
	@rtype: list(utterance_c)
	"""
	by_subject = dictops.dict_of_lists()	#: utterances indexed by subjectID
	for u in ulist:
		by_subject.add(u.info['subjectID'], u)
	slist = list(by_subject.keys())
	sl2 = []
	for i in range(len(slist)):
		sl2.append(random.choice(slist))
	rv = []
	for s in sl2:
		rv.extend(by_subject[s])
	return rv


class fitdata(object):
	def __init__(self, u, ival, desc):
		self.u = u
		self.ival = ival
		self.dvals = []
		self.desc = desc


def package(ulist, N, target, direction):
	assert N > 0
	assert ulist
	tofit = dictops.dict_of_lists()
	for u in ulist:
		no_s = 0
		while u.data[0].lbl != 'S':
			u.data.pop(0)
			no_s += 1
		while u.data[-1].lbl != 'S':
			u.data.pop(-1)
			no_s += 1
		if no_s > len(u.data)//10:
			die.warn("Too much trimming to find end silences: %s" % str(u))
			continue
		elif no_s:
			die.info("U didn't start and end with 'S': %s" % str(u))
		if len(u.data) < 2+N:
			continue
		# print 'u=', u
		itgt = u.idx[target]
		for jtgt in range(1, len(u.data)-1):
			if u.data[jtgt].lbl == 'S':
				ttype = 'S'
			elif u.data[jtgt+1].lbl=='S' and u.data[jtgt-1].lbl=='S':
				continue
			elif u.data[jtgt+1].lbl == 'S':
				ttype = 'final'
			elif u.data[jtgt-1].lbl == 'S':
				ttype = 'initial'
			else:
				ttype = 'medial'

			s_count = 0
			if ttype=='final' and direction<0:
				s_count = 1
			elif ttype=='initial' and direction>0:
				s_count = 1
			elif ttype=='S':
				s_count = 1
			
			key = (ttype, u.data[jtgt].lbl, target, direction)
			tmp = fitdata(u, u.data[jtgt].data[itgt], key)
			failed = False
			for k in range(1, N+1):
				if direction != 0:
					j = jtgt - direction*k
				else:
					x = (1+k)//2
					if k%2 == 0:
						j = jtgt - x
					else:
						j = jtgt + x
				if j<1 or j>=len(u.data)-1:
					failed = True
					break
				if u.data[j].lbl == 'S':
					s_count -= 1
				if s_count < 0:
					failed = True
					break
				tmp.dvals.extend( u.data[j].data )
			if not failed:
				tofit.add(tmp.desc, tmp)
	return tofit



def multiF(ulist, n, direction):
	assert n > 0
	"""Compute a linear predictor for a set of data.
	@type n: int
	@param n: how many syllables should be used to predict?
	@return: a dictionary containing a lot of information about the
		linear regression.
	@rtype: dict(str:something)
	"""
	ulist = sample_people_with_replacement(ulist)
	o = []
	for target in utterance_c.idx.keys():
		tmpfit = package(ulist, n, target, direction)
		for (key, flist) in tmpfit.items():
			tmp = oneF(key, flist)
			tmp['N'] = n
			o.append(tmp)
			# return o	# TEST
	return o



def prep_for_localfit(tfi, subjs, acoustics=False, constant=False, dosub=False):
	def mka(s, sidx):
		tmp = [-1.0/len(sidx)]*(len(sidx)-1)
		i = sidx[s]
		if i < len(sidx)-1:
			tmp[i] = 1
		return tmp
	sidx = dict([(s,i) for (i,s) in enumerate(subjs)])
	tofit = []
	np = None
	for t in tfi:
		dvars = []
		if acoustics:
			dvars.extend(t.dvals)
		if dosub:
			dvars.extend( mka(t.u.info['subjectID'], sidx) )
		if constant:
			dvars.append(1)
		if np is None:
			np = len(dvars)
		else:
			assert np == len(dvars)
		tofit.append((dvars, (t.ival,)))
	return (tofit, np)


def oneF(key, flist):
	constant = True
	subjs = set()
	np1 = None
	for f in flist:
		subjs.add(f.u.info['subjectID'])
	subjs = sorted(subjs)
	info = {'N0': len(flist), 'Nsubj': len(subjs),
		'significant1': 0,
		'ttype': key[0], 'lbl': key[1], 'target': key[2], 'direction': key[3]
		}
	if len(subjs) < 3:
		info['problem'] = 'too few subjects'
		return info
	if len(flist) > MAXFIT:
		info['sampled_frac'] = float(MAXFIT)/float(len(flist))
		flist = random.sample(flist, MAXFIT)
	else:
		info['sampled_frac'] = 1.0
	tofit0, np0 = prep_for_localfit(flist, subjs, acoustics=False, constant=True, dosub=False)
	if len(tofit0) < np0 + np0//3 + 11:
		info['problem'] = 'too short 0'
		return info
	A0, B0, errs0, sv0, rank0 = g_localfit.localfit(tofit0, constant=constant, minsv=MINSV, minsvr=MINSVR)
	info['dof0'] = len(tofit0)-np0
	info['np0'] = np0
	info['n'] = len(tofit0)
	del tofit0

	tofit1, np1 = prep_for_localfit(flist, subjs, acoustics=False, constant=True, dosub=True)
	if len(tofit1) < np1 + np1//3 + 11:
		info['problem'] = 'too short 1'
		return info
	A1, B1, errs1, sv1, rank1 = g_localfit.localfit(tofit1, constant=constant, minsv=MINSV, minsvr=MINSVR)
	info['dof1'] = len(tofit1)-np1
	info['np1'] = np1
	fstat1 = stats.f_value(errs0[0], errs1[0], len(tofit1)-np0, len(tofit1)-np1)
	fprob1 = stats.fprob(np1-np0, len(tofit1)-np1, fstat1)
	del tofit1
	if fprob1 < 1e-4:
		info['significant1'] = 1

	tofit2, np2 = prep_for_localfit(flist, subjs, acoustics=True, constant=True, dosub=True)
	if len(tofit2) < np2 + np2//3 + 11:
		info['problem'] = 'too short 2'
		return info
	A2, B2, errs2, sv2, rank2 = g_localfit.localfit(tofit2, constant=constant, minsv=MINSV, minsvr=MINSVR)
	info['dof2'] = len(tofit2)-np2
	info['np2'] = np2
	fstat2 = stats.f_value(errs1[0], errs2[0], len(tofit2)-np1, len(tofit2)-np2)
	fprob2 = stats.fprob(np2-np1, len(tofit2)-np2, fstat2)
	if fprob2 < 1e-4:
		info['significant2'] = 1

	info.update( {'F1': fstat1, 'P1': fprob1, 'svr1': sv1[-1]/sv1[0], 'rank1': rank1,
			'effect_size1': (errs0[0]-errs1[0])/errs0[0],
			'F2': fstat2, 'P2': fprob2, 'svr2': sv2[-1]/sv2[0], 'rank2': rank2,
			'effect_size2': (errs1[0]-errs2[0])/errs1[0],
			'err0': errs0[0]
			} )
	return info




def process(dblist, selector, seg, nlist, direction, ofd):
	"""
	@param nlist: list of how long of a predictor to use
	@type nlist: int
	"""
	ulist = read_in_utterances(dblist, selector, seg, ofd)
	for n in nlist:
		for infoplus in multiF(ulist, n, direction):
			infoplus['selector'] = selector
			infoplus['mapM'] = Map['m']
			ofd.writelines( [avio.concoct(infoplus), '\n'] )
			ofd.flush()


if __name__ == '__main__':
	dblist = [q for q in sys.argv[1:] if 'ASR' not in q]
	Selector = '"ASR" not in textfile'
	Nlist = [1,2,3,4,5,6,7,8]
	Direction = 1
	Outfile = None
	arglist = sys.argv[1:]
	while arglist and arglist[0].startswith('-'):
		arg = arglist.pop(0)
		if arg == '--':
			break
		elif arg == '-s':
			Selector = arglist.pop(0)
		elif arg == '-poetry':
			Selector = 'property=="poetry" and "ASR" not in textfile'
		elif arg == '-prose':
			Selector = 'property!="poetry" and "ASR" not in textfile'
		elif arg == '-o':
			Outfile = arglist.pop(0)
		elif arg == '-dir':
			Direction = int(arglist.pop(0))
		elif arg == '-mapM':
			Map['m'] = arglist.pop(0)
		elif arg == '-n':
			Nlist = [ int(q) for q in arglist.pop(0).split(',') ]
		elif arg == '-svr':
			MINSVR = float(arglist.pop(0))
			assert 0 < MINSVR < 1.0
		else:
			die.die('Unrecognized argument: %s' % arg)
	assert Outfile is not None
	Ofd = gpkmisc.open_nowipe(Outfile, '.av', mode='w')
	process(arglist, Selector, SEGMENTATION, Nlist, Direction, Ofd)
