#! 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




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)
		return [math.log(dur), ldur]
	
	idx = {'lbl':0, 'duration':1, 'ldur':2}

	ridx = dictops.rev1to1(idx)

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



def select_data_files(dblist, selector, nselections):
	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)

	if nselections >= 0:
		bd = BD.blue_attributes(inspector, infolist)
		infolist = bd.pick(nselections)
		# infolist = random.sample(infolist, nselections)
	return infolist




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


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['s:%d:%s' % (1+(i//npper), utterance_c.ridx[1+(i%npper)])] = float(sigma[i])
	return info


def quantize(x):
	return '%d' % (int(x>1.0))

def F(ulist, N, lbl, target):
	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 =  []
				for j in range(1, N+1):
					if u.data[i-j][0] == 'S':
						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
	if len(tofit) > MAXFIT:
		tofit = random.sample(tofit, MAXFIT)
	tmp = []
	itmp = dictops.dict_of_averages()
	for d in tofit:
		tmp.append(d[1][0])
		for (i,x) in enumerate(d[0]):
			itmp.add(i, x)
	obar = gpkmisc.median(tmp)
	ibar = itmp.get_avgs()
	for d in tofit:
		qq = []
		for (i,x) in enumerate(d[0]):
			qq.append(quantize(x/ibar[i]))
		print '\t'.join([quantize(d[1][0]/obar)] + qq) 
	

NSELECTIONS = 100


def process(dblist, selector, seglist):
	ulist = []
	for s in seglist:
		ulist.extend( read_in_utterances(dblist, selector, NSELECTIONS, s) )
	print '\t'.join(['D0', 'D1', 'L1', 'D2', 'L2', 'D3', 'L3', 'D4', 'L4'])
	for (lbl,target) in [('^S,S,^S','duration'), 
				# ('V,C,S', 'duration'),
				# ('V,C,V,C,V', 'duration'),
				]:
		F(ulist, 4, lbl, target)
		sys.stdout.flush()


if __name__ == '__main__':
	dblist = [q for q in sys.argv[1:] if 'ASR' not in q]
	process(dblist, 'True', [SEGMENTATION])
