#! python

import os
import sys
import math

import numpy

import gpkimgclass

from gmisclib import die
from gmisclib import mcmc
from gmisclib import avio
from gmisclib import tsops
from gmisclib import fiatio
from gmisclib import gpkmisc
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 = 'rec4fn_HTKsimple4.lab'
MINPAUSE = 0.15
MAXPAUSE = 0.50




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)
		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
		s = [ self.compute_statistics(s, e) for (s, e, lbl) in lbls ]
		s = s[-1:] + s + s[:1]
		assert len(s) == len(lbls)+2
		for i in range(1, len(lbls)-1):
			out.append( s[i-1] + s[i] + s[i+1] )
		return out


	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
		assert ei > si
		lbar = numpy.average(self.lnorm[si:ei])
		ldur = lbar*min(e-s, 2*TYPDUR)/TYPDUR
		fbar = numpy.average(self.fric[si:ei])
		lskew = numpy.sum(self.lnorm[si:ei]*numpy.arange(ei-si)) / numpy.sum(numpy.arange(ei-si)) - 0.5
		fskew = numpy.sum(self.fric[si:ei]*numpy.arange(ei-si)) / numpy.sum(numpy.arange(ei-si)) - 0.5
		complexity = math.log(numpy.sum((e-s) / numpy.exp(self.lpdur[si:ei])))
		pdur = numpy.exp(self.lpdur)/TYPDUR
		lpdur = numpy.sum(self.lnorm[si:ei]*pdur[si:ei])
		fpdur = numpy.sum(self.fric[si:ei]*pdur[si:ei])
		return [lbar, (e-s)/TYPDUR, ldur, lpdur, fbar, fpdur, lskew, fskew, complexity]
	


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



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

	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 %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]



NSELECTIONS = 100
	

def process(target, dblist, selector, seglist):
	ulist = []
	for s in seglist:
		ulist.extend( read_in_utterances(dblist, selector, NSELECTIONS, s) )
	for n in range(8):
		for lbl in ['S', 'C', 'V']:
			print lbl, n, F(ulist, n, lbl, target)


if __name__ == '__main__':
		process('lbar', sys.argv[1:], 'True', ['rec2fn_HTKsimple2.lab', SEGMENTATION])
