#! python

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'
MINPAUSE = 0.15
MAXPAUSE = 0.50
MAXFIT = 12000


def asinh(x):
	if x > 0:
		sign = 1
	elif x == 0:
		sign = 0
	else:
		sign = -1
	return sign*numpy.log(sign*x + numpy.sqrt(x*x + 1))


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.lpdur = numpy.log(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 = []
		assert len(lbls) > 2
		for (s, e, lbl) in lbls:
			try:
				self.data.append([lbl] + self.compute_statistics(s, e))
			except IndexError, x:
				die.warn('OutOfRange: %s for (%.3f,%.3f,%s) in %s' % (x, s, e, lbl, self.path))


	def compute_statistics(self, s, e):
		TYPDUR = 0.10
		si = int(round(self.taxis.index(s, limit=True)))
		ei = int(round(self.taxis.index(e, limit=True))) + 1
		dur = (e-s)/TYPDUR
		lbar = numpy.average(self.lnorm[si:ei])
		ldur = lbar*dur/(1.0+dur)
		fbar = numpy.average(self.fric[si:ei])
		fdur = fbar*dur/(1.0+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)
		pdur = numpy.exp(self.lpdur)/TYPDUR
		complex_dur = math.log(numpy.sum(1.0/pdur[si:ei])/dur)
		return [math.log(dur), ldur, fdur, complex_dur, lskew]
	
	idx = {'lbl':0, 'duration':1, 'ldur':2, 'fdur':3, 'complex_dur':4, 'lskew': 5}

	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]


RSCALE = 0.001


class context(object):
	def __init__(self, udata, i):
		if i > 1:
			self.lleft = udata[i-2][0]
		else:
			self.lleft = 'S'
		if i > 0:
			self.left = udata[i-1][0]
		else:
			self.left = 'S'
		try:
			self.right = udata[i+1][0]
		except IndexError:
			self.right = 'S'
		try:
			self.rright = udata[i+2][0]
		except IndexError:
			self.rright = 'S'
		self.center = udata[i][0]


	def amatch(self, lbl, ttm):
		if lbl == '*':
			return True
		if lbl.startswith('^'):
			return lbl[1:] != ttm
		return lbl == ttm


	def matches(self, lbl):
		a = lbl.split(',')
		if len(a) == 1:
			return self.center == a[0]
		if len(a) == 3:
			return (self.amatch(a[0], self.left) and
				self.amatch(a[1], self.center) and
				self.amatch(a[2], self.right)
				)
		assert len(a) == 5
		return (self.amatch(a[0], self.lleft) and
			self.amatch(a[1], self.left) and
			self.amatch(a[2], self.center) and
			self.amatch(a[3], self.right) and
			self.amatch(a[4], self.rright)
			)


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_with_replacement(x, n):
	rv = []
	for i in range(n):
		rv.append(random.choice(x))
	return rv


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
	np = N*npper+constant
	tofit = []
	for u in ulist:
		itgt = u.idx[target]
		for i in range(N, len(u.data)):
			if context(u.data, i).matches(lbl):
				oc = (u.data[i][itgt],)
				ic =  []
				cS = cross_S
				for j in range(1, N+1):
					if u.data[i-j][0]=='S':
						if cS:
							cS -= 1
						else:
							ic = None
							break
					ic.extend(u.data[i-j][1:])
				if ic is not None:
					tofit.append( (ic,oc) )
	info = {'dof': len(tofit)-1-np, 'ddof': np,
		'N': N, 'lbl': lbl, 'target': target,
		'significant': 0
		}
	if len(tofit) < 2*np+11:
		return info
	tofit = sample_with_replacement(tofit, min(MAXFIT, len(tofit)))
	info['sampled_frac'] = float(min(MAXFIT, len(tofit)))/float(len(tofit))
	# for (i,o) in tofit:
		# print 'I', i
		# print 'O', o
	info.update( get_sigmas(tofit, npper) )
	perrs = g_localfit.err_before_fit(tofit, constant=constant)[0]
	regstr = numpy.identity(np)
	for i in range(npper+constant):
		regstr[i,i] = 0.0
	for i in range(npper, N*npper):
		regstr[constant+i,constant+i] = info[name('s', npper, i)]
	A, B, errs, sv, rank = g_localfit.reg_localfit(tofit, constant=constant, regstr=1.0, rscale=RSCALE)
	if errs[0] > perrs:
		die.warn("Fit gets worse!")
	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])
			# info['C:%d:%s' % (1+(i//npper), utterance_c.ridx[1+(i%npper)])] = float(A[0,i])
		return info
	return info


def name(prefix, npper, i):
	return '%s:%d:%s' % (prefix, 1+(i//npper), utterance_c.ridx[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:
			for i in range(reps):
				print avio.concoct(F(ulist, n, lbl, target, cross_S=int(lbl[0]=='S')))
				sys.stdout.flush()


if __name__ == '__main__':
	dblist = [q for q in sys.argv[1:] if 'ASR' not in q]
	LTlist = []
	for (lbl,target) in [('V,S,^S','duration'), ('C,S,^S','duration'),
				('S,C,V', 'duration'), ('V,C,V', 'duration'), ('V,C,S', 'duration'),
				('S,C,V', 'complex_dur'), ('V,C,V', 'complex_dur'), ('V,C,S', 'complex_dur'),
				('S,C,V', 'ldur'), ('V,C,V', 'ldur'), ('V,C,S', 'ldur'),
				('S,V,C', 'ldur'), ('S,C,V,C,V', 'ldur'), ('V,C,V,C,V', 'ldur'),
				('V,C,V,C,S', 'ldur'), ('C,V,S', 'ldur'),
				('S,V,C', 'lskew'), ('S,C,V,C,V', 'lskew'), ('V,C,V,C,V', 'lskew'),
				('V,C,V,C,S', 'lskew'), ('C,V,S', 'lskew'),
				('S,V,C', 'duration'), ('S,C,V,C,V', 'duration'), ('V,C,V,C,V', 'duration'),
				('V,C,V,C,S', 'duration'), ('C,V,S', 'duration'),
				('S,V,C', 'complex_dur'), ('S,C,V,C,V', 'complex_dur'), ('V,C,V,C,V', 'complex_dur'),
				('V,C,V,C,S', 'complex_dur'), ('C,V,S', 'complex_dur'),
				]:
			LTlist.append((lbl, target))
	Selector = 'True'
	Reps = 1
	Nlist = range(1,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 == '-rscale':
			RSCALE = float(arglist.pop(0))
			assert 0 < RSCALE
		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)
