#! 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 mcmc_idxr
# from gmisclib import spread_jobs
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 RM_ts_lib import zero as Z


SEGMENTATION = 'recVx_HTKfv.lab'
MAXFIT = 12000

LTlist = [('S','duration'), ('S','ldur'), ('S', 'ccluster'),
		('initial', 'ldur'), ('medial', 'ldur'), ('final', 'ldur'),
		('initial', 'duration'), ('medial', 'duration'), ('final', 'duration'),
		('initial', 'speechrate'), ('medial', 'speechrate'), ('final', 'speechrate'),
		('initial', 'voice1'), ('medial', 'voice1'), ('final', 'voice1'),
		('initial', 'ccluster'), ('medial', 'ccluster'), ('final', 'ccluster'),
		('initial', 'vowel'), ('medial', 'vowel'), ('final', 'vowel'),
		('initial', 'vowel'), ('medial', 'vowel'), ('final', 'vowel'),
		('initial', 'lskew'), ('medial', 'lskew'), ('final', 'lskew')
		]


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)


class utterance_c(object):
	def __init__(self, dbdir, info, segmentation):
		self.info = info
		self.dbdir = dbdir
		self.read_data(dbdir, info, segmentation)

	def read_data(self, dbdir, info, segmentation):
		self.segmentation = segmentation
		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'))
		self.taxis, loud, irr, pdur = tsops.common([ld, ird, pdd])
		self.pdur = pdur[:,pdd.types['pseudoduration']]
		
		loud = loud[:,ld.types['loudness']]
		loud -= Z.loud_zero(loud)
		if not numpy.sum(numpy.maximum(0.0, loud)) > 30.0*numpy.sum(numpy.maximum(0.0, -loud)):
			raise DubiousData, "Loudness is not dominantly positive."
		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))
		self.data = []
		if len(lbls) <= 3:
			die.warn('lbls too short: %d for %s / %s' % (len(lbls), self.path, self.segmentation))
			return
		for i in range(1, len(lbls)-1):
			if(lbls[i][2]=='V' and lbls[i-1][2]=='C' and lbls[i+1][2]=='C'):
				tmp = self.compute_statistics(lbls[i-1][0],
								lbls[i][0],
								lbls[i][1],
								lbls[i+1][1])
				self.data.append(['CVC'] + tmp)
			elif(lbls[i][2]=='V' and lbls[i-1][2]=='S' and lbls[i+1][2]=='C'):
				tmp = self.compute_statistics(lbls[i][0],
								lbls[i][0],
								lbls[i][1],
								lbls[i+1][1])
				self.data.append(['SVC'] + tmp)
			elif(lbls[i][2]=='V' and lbls[i-1][2]=='C' and lbls[i+1][2]=='S'):
				tmp = self.compute_statistics(lbls[i-1][0],
								lbls[i][0],
								lbls[i][1],
								lbls[i][1])
				self.data.append(['CVS'] + tmp)
			elif(lbls[i][2]=='S' and lbls[i-1][2]=='C' and lbls[i+1][2]=='C'):
				tmp = self.compute_statistics(lbls[i-1][0],
								lbls[i][0],
								lbls[i][1],
								lbls[i+1][1])
				self.data.append(['CSC'] + tmp)
			elif(lbls[i][2]=='S' and lbls[i-1][2]=='V' and lbls[i+1][2]=='C'):
				tmp = self.compute_statistics(lbls[i][0],
								lbls[i][0],
								lbls[i][1],
								lbls[i+1][1])
				self.data.append(['VSC'] + tmp)
			elif(lbls[i][2]=='S' and lbls[i-1][2]=='C' and lbls[i+1][2]=='V'):
				tmp = self.compute_statistics(lbls[i-1][0],
								lbls[i][0],
								lbls[i][1],
								lbls[i][1])
				self.data.append(['CSV'] + tmp)
			elif(lbls[i][2]=='S' and lbls[i-1][2]=='V' and lbls[i+1][2]=='V'):
				tmp = self.compute_statistics(lbls[i][0],
								lbls[i][0],
								lbls[i][1],
								lbls[i][1])
				self.data.append(['VSV'] + tmp)
			else:
				# print 'Must Be C: %s lbls[%d]=%s %s' % (lbls[i-1][2], i, lbls[i], lbls[i+1][2])
				assert lbls[i][2]=='C'


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


	def compute_statistics(self, c1s, vs, ve, c2e):
		si = int(round(self.taxis.index(0.25*vs+0.75*c1s)))
		ei = int(round(self.taxis.index(0.75*ve+0.25*c2e)))
		TYPDUR = 0.10
		dur = (0.75*ve+0.25*c2e-0.25*vs-0.75*c1s)/TYPDUR
		lbar = numpy.average(self.lnorm[si:ei])
		ldur = lbar*dur/(1.0+dur)
		fric = self.fric[si:ei+1]
		v1 = numpy.mean(fric)
		cc = numpy.sum(fric * numpy.greater(fric,0.0))*self.taxis.dt()/(1.0+dur)
		vowel = numpy.sum(fric * numpy.less(fric,0.0))*self.taxis.dt()/(1.0+dur)
		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(self.lnorm[si:ei]*skew)/(1.0+dur)
		return [math.log(dur), ldur, v1, speechrate, cc, lskew, vowel]
	
	idx = {'lbl':0, 'duration':1, 'ldur':2, 'voice1':3, 'speechrate':4, 'ccluster': 5,
			'lskew': 6, 'vowel': 7}

	ridx = dictops.rev1to1(idx)

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



def select_data_files(dblist, selector):
	assert len(dblist) > 0
	dur = []
	syl = []
	tmp = []
	for dbpath in dblist:
		dbd, dbf = os.path.split(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])
	print '# Range of durations counts accepted:', dl, dh
	print '# Range of syllable counts accepted:', snl, snh
	infolist = []
	for info in tmp:
		if(dl <= info['__duration'] <= dh and
				snl <= info['__nsyls'] <= snh):
			infolist.append(info)

	return infolist




def read_in_utterances(dblist, selector, segmentation):
	ulist = []
	for info in select_data_files(dblist, selector):
		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


def context(udata, i):
	if udata[i][0][1] == 'S':
		return 'S'
	if udata[i][0][1] != 'V':
		return 'dontmatch'
	initial = i==0 or udata[i-1][0][1]=='S'
	final = i==len(udata)-1 or udata[i+1][0][1]=='S'
	if initial and final:
		return 'tooshort'
	if initial and udata[i+1][0][1]=='V':
		return 'initial'
	if final and udata[i-1][0][1]=='V':
		return 'final'
	if udata[i+1][0][1]=='V' and udata[i-1][0][1]=='V':
		return 'medial'
	return 'leftovers'


def get_sigmas(tofit, npper):
	info = {}
	s = None
	ss = None
	ld = len(tofit)
	for d in tofit:
		assert len(d)==2
		i = d[0]
		if s is None:
			s = numpy.zeros((len(i),))
		if ss is None:
			ss = numpy.zeros((len(i),))
		i = numpy.asarray(i)
		numpy.add(s, i, s)
		numpy.add(ss, i*i, ss)
	sigma = numpy.sqrt((ss - s*s/ld)/(ld-1))
	for i in range(sigma.shape[0]):
		info[name('s', npper, i)] = float(sigma[i])
		# info['s:%d:%s' % (1+(i//npper), utterance_c.ridx[1+(i%npper)])] = float(sigma[i])
	return info


def sample_people_with_replacement(ulist):
	by_subject = dictops.dict_of_lists()
	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


def center_silence(s):
	assert len(s)%2==1
	return len(s)==3 and s[1]=='S'


def package(ulist, N, target, lbl, cross_S):
	assert (lbl=='initial' and cross_S) or (lbl=='medial' and not cross_S) or (lbl=='final' and cross_S)
	assert N > 0
	assert ulist and ulist[0].data and len(ulist[0].data[0])>=1
	npper = len(ulist[0].data[0])-1
	tofit = []
	for u in ulist:
		# print 'u=', u
		itgt = u.idx[target]
		for i in range(len(u.data)):
			# print 'context=', context(u.data, i), 'lbl=', lbl
			if context(u.data, i) == lbl:
				oc = (u.data[i][itgt],)
				ic =  []
				cS = cross_S
				for ii in range(1, npper*N+1):
					j = offset(npper, ii)
					if not ( 0 <= i-j < len(u.data) ):
						ic = None
						break
					if center_silence(u.data[i-j][0]):
						if cS:
							cS -= 1
						else:
							ic = None
							break
					ic.append(u.data[i-j][item(npper, ii)])
				if ic is not None:
					assert len(oc)>0
					tofit.append( (ic,oc) )
	return (tofit, npper)




def F(ulist, N, lbl, target, cross_S=0):
	assert N > 0
	constant = True
	assert ulist and ulist[0].data and len(ulist[0].data[0])>=1
	npper = len(ulist[0].data[0])-1
	ulist = sample_people_with_replacement(ulist)
	tofit, npper = package(ulist, N, target, lbl, cross_S)
	np = N*npper+constant
	info = {'dof': len(tofit)-1-np, 'ddof': np,
		'N': N, 'lbl': lbl, 'target': target,
		'significant': 0
		}
	if len(tofit) < 2*np+11:
		return info
	if len(tofit) > MAXFIT:
		tofit = random.sample(tofit, MAXFIT)
	info['sampled_frac'] = float(min(MAXFIT, len(tofit)))/float(len(tofit))
	info.update( get_sigmas(tofit, npper) )
	perrs = g_localfit.err_before_fit(tofit, constant=constant)[0]
	A, B, errs, sv, rank = g_localfit.localfit(tofit, constant=constant, minsv=MINSV, minsvr=MINSVR)
	fstat = stats.f_value(perrs, errs[0], len(tofit)-1, len(tofit)-1-np)
	fprob = stats.fprob(np, len(tofit)-1-np, fstat)
	info.update( {'F': fstat, 'P': fprob, 'svr': sv[-1]/sv[0], 'err': errs[0], 'effect_size': (perrs-errs[0])/perrs, 'rank': rank })
	if fprob < 1e-4:
		info['significant'] = 1
	info['C:*:1'] = B[0]
	assert len(A.shape) == 2
	assert A.shape[0]==1 and A.shape[1]%npper==0
	for i in range(A.shape[1]):
		info[name('C', npper, i)] = float(A[0,i])
	return info


def name(prefix, npper, i):
	return '%s:%d:%s' % (prefix, offset(npper, i), utterance_c.ridx[1+(i%npper)])

def offset(npper, i):
	j = i//npper
	return (1 + j//2) * (-1)**(j)

def item(npper, i):
	return 1 + i%npper



def process(dblist, selector, seglist, ltlist, nlist, reps):
	ulist = []
	for s in seglist:
		ulist.extend( read_in_utterances(dblist, selector, s) )
	for n in nlist:
		for (lbl, target) in ltlist:
			if len(lbl) == 3:
				cross_S = int(lbl[0]=='S')+int(lbl[2]=='S')
			else:
				cross_S = 0
			for i in range(reps):
				print avio.concoct(F(ulist, n, lbl, target, cross_S=cross_S))
				sys.stdout.flush()


if __name__ == '__main__':
	dblist = [q for q in sys.argv[1:] if 'ASR' not in q]
	Selector = '"ASR" not in textfile'
	Reps = 1
	Nlist = [2,4,6,8]
	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 == '-n':
			Nlist = [ int(arglist.pop(0)) ]
		elif arg == '-svr':
			MINSVR = float(arglist.pop(0))
			assert 0 < MINSVR < 1.0
		elif arg == '-i':
			i = int(arglist.pop(0))
			LTlist = [ LTlist[i] ]
		elif arg == '-reps':
			Reps = int(arglist.pop(0))
			assert Reps > 0
		else:
			die.die('Unrecognized argument: %s' % arg)
	process(arglist, Selector, [SEGMENTATION], LTlist, Nlist, Reps)
