#! python

import os
import re
import sys
import math
import random
import hashlib

import numpy

from oxigen import read_words

from gmisclib import die
from gmisclib import fiatio
from gmisclib import wavio
from gmisclib import xwaves_lab
from gmisclib import Numeric_gpk as NG

DB = '/proj/OLD/oxigen/oxi_rest/oxigen/DB.fiat'
MeanLen = 2.0
MINLENGTH = 0.5	# seconds

_punct = re.compile("""[,.;":!?]+""")

def find_chunks(db):
	hdr, d, comments = fiatio.read(open(db, 'r'))
	random.shuffle(d)
	dbdir = os.path.dirname(db)
	for info in d:
		if not (info['task']=='read' or info['task']=='retold'):
			continue
		fbase = os.path.join(dbdir, info['d'], info['f'])
		try:
			labels = read_words.read_with_beginnings(fbase, loose=0)
		except xwaves_lab.Error, x:
			die.info("No labels for %s: %s" % (fbase, x))
			continue
		die.info("Using %s" % fbase)
		if len(labels) <= 2:
			continue
		# '*' means a pause.  Convert to punctuation.
		if labels[0][1].strip() == '*':
			labels[0] = (labels[0][0], '.')
		for i in range(1,len(labels)):
			# '*' means a pause.  Convert to punctuation.
			if labels[i][1].strip() == '*':
				dur = labels[i][0] - labels[i-1][0]
				if dur >= 0.130:
					labels[i] = (labels[i][0], '.')
				else:
					labels[i] = (labels[i][0], ';')
			else:
				# Convert most other punctuation to commas.  (Not apostrophes, not '#'.)
				m = _punct.search(labels[i][1])
				if m:
					labels[i] = (labels[i][0], _punct.sub(',', labels[i][1]))

		i = 0
		while i < len(labels)-1:
			n = int(math.ceil(random.expovariate(1.0/MeanLen)))
			if i+n >= len(labels):
				i += 1
				continue
			text = ' '.join( [q[1] for q in labels[i:i+n]] )
			if '#' in text:	
				i += 1	#	"#" is the code for a noise or other problem
				continue
			info['dbdir'] = dbdir
			if i > 0:
				start = labels[i-1][0]
			else:
				start = 0.0	# Assumes that audio starts at t=0.0
			end = labels[i+n-1][0]
			if end - start > MINLENGTH:
				yield (info.copy(), start, end, text)
			i += max(n-1, 1)

# import pylab

def make_corpus_of_fragments(db, tdir, N):
	HWW = 0.010
	w = fiatio.writer(open(os.path.join(tdir, "DB.fiat"), 'w'))
	w.comment("Created by %s from %s" % (sys.argv[0], db))
	n = 0
	chunks = list(find_chunks(db))
	random.shuffle(chunks)
	for (info, start, end, text) in chunks:
		fn = os.path.join(info['dbdir'], info['d'], info['f']) + '.wav'
		try:
			d = wavio.read(fn, tstart=max(0.0, start-HWW), t_end=end+HWW)
		except (ValueError, wavio.wave.Error), x:
			die.warn("Whoops cannot read .wav file %s: %s" % (fn, x))
			continue
		ofn = hashlib.md5(repr((fn, start, end, text))).hexdigest()
		nww = int(round(HWW/d.dt()))
		win = NG.edge_window(d.d.shape[0], nww, nww, typeleft="cos", typeright="cos")
		# pylab.plot(win)
		# pylab.show()
		numpy.subtract(d.d, numpy.average(d.d), d.d)
		numpy.multiply(d.d, win[:,numpy.newaxis], d.d)
		wavio.write(d, os.path.join(tdir, ofn) + '.wav')
		ofo = {'d': '.', 'f': ofn, 'textfile': ofn + '.txt',
			'dialect': info['dialect'], 'from': fn
			}
		ofo['textfile'] = os.path.join(tdir, ofn) + '.txt'
		open(ofo['textfile'], 'w').writelines([text, '\n'])
		w.datum(ofo)
		n += 1
		if n > N:
			break
	w.close()


if __name__ == '__main__':
	tdir = sys.argv[1]
	if len(sys.argv) < 3:
		N = 2000
	else:
		N = int(sys.argv[2])
	make_corpus_of_fragments(DB, tdir, N)
