#!/usr/bin/env python

"""This program takes an ESPS label file containing marks
(and a data file of f0 vs. time),
and returns a two-column ASCII file of f0 and voicing.

The labels are described in README.tex, under '*.weirdf0'.

You can merge region end and start marks like this: '][=123'.

This script ignores voicing information in the f0_file from the pitch
tracker, except to check your marks for consistency and sanity.
You need to mark all voiced regions by hand.
"""

import xwaves_lab
import xwaves_errs
import Num
import die
import gpkavg
import math

LOWF0 = 30	# The lowest plausible F0 that people produce
HIGHF0 = 2400	# The highest plausible F0 that people produce

class VoicingFileError(xwaves_errs.BadFileFormatError):
	def __init__(self, x):
		xwaves_errs.BadFileFormatError.__init__(self, x)

class BadF0(ValueError):
	def __init__(self, x):
		ValueError.__init__(self, x)

class SillyCorrection(ValueError):
	def __init__(self, x):
		ValueError.__init__(self, x)

class NothingToCorrect(ValueError):
	def __init__(self, x):
		ValueError.__init__(self, x)

def _check_set(valid, between_vec, of0, target):
	if target<LOWF0 or target>HIGHF0:
		raise SillyCorrection, "Silly target: %g" % target
	CLIP = 0.1
	NICETOL = 0.10
	nice = [0.3333, 0.5, 0.6667, 1.0, 1.5, 2.0, 3.0]
	assert valid.shape == of0.shape
	assert between_vec.shape == valid.shape
	ivb = Num.nonzero(valid * between_vec)
	if len(ivb) == 0:
		return (0, None)
	raw, sd = gpkavg.avg(Num.take(of0, ivb), None, CLIP)
	for ni in nice:
		if abs(1 - (ni*raw/target)) < NICETOL:
			die.info('_do_set: correcting by a factor of %g r=%g t=%g'
					% (ni, raw, target))
			return (0, ni)
	die.info('_do_set: no correction: marking %d invalid' % len(ivb))
	return (1, None)	# Should be marked as invalid.


def _log2(x):
	return math.log(x)/math.log(2.0)


def _act(time, (t0, t), flags, correction, irr, of0, valid,
				(adjustf0, conservative_v, f0_is)):
	between_vec = (time >= t0) * (time < t)
	between_idx = Num.nonzero( between_vec )
	if '*' in flags:
		if correction<0.2 or correction>5.0:
			raise SillyCorrection, 'Silly correction factor: %g at %.3f' % (correction, t0)
	invalidate = 0
	if '=' in flags:
		invalidate, correction = _check_set(valid, between_vec,
							of0, correction)

	if f0_is == 'pitch':
		if 'c' in flags:
			correction = 1.0
		elif 'u' in flags:
			if correction is None or correction <= 1.0:
				correction = 2.0
		elif 'd' in flags:
			if correction is None or correction >= 1.0:
				correction = 0.5
		else:
			invalidate = 1
	else:
		invalidate = not correction

	if 's' in flags and not conservative_v:
		invalidate = 1

	if 'i' in flags:
		Num.put(irr, between_idx, 1.0)

	if correction is not None and adjustf0 and abs(correction-1.0)>0.001:
		Num.put(of0, between_idx,
			Num.take(of0, between_idx)*correction)

	if invalidate:	# Mark the range of time as having invalid f0
		Num.put(valid, between_idx, 0)




def f0correct(markfile, time, f0, voicing=None, required=0, loose=0, adjustf0=1, conservative_v=1, f0_is='pitch'):
	assert time.shape == f0.shape
	die.note('file', markfile)
	of0 = Num.array(f0, Num.Float)
	if voicing is not None:
		assert  voicing.shape == time.shape
		valid = (voicing > 0.5)*(f0 > LOWF0)
	else:
		valid = f0 > LOWF0
	nv0 = Num.sum(valid)
	irr = Num.zeros(f0.shape, Num.Float)

	assert of0.shape == valid.shape

	try:
		voicinghdr, voicingdata = xwaves_lab.read(markfile)
	except xwaves_errs.NoSuchFileError:
		if required:
			raise
		else:
			return (of0, valid, irr)

	die.info('Correcting f0, using %s' % markfile)
	state = 'out'
	flags = []
	correction = None
	t0 = None
	for (t, lab) in voicingdata:
		recognized = False
		if state == 'in' and lab.startswith(']'):
			_act(time, (t0, t), flags, correction, irr, of0, valid,
				(adjustf0, conservative_v, f0_is))
			recognized = True
			state = 'out'
			flags = []
			correction = None
			lab = lab[1:]	# strip off beginning of '][' labels.

		if state == 'out' and lab.startswith('['):
			state = 'in'
			t0 = t
			i = 1
			while i < len(lab):
				if lab[i] in 'iscud':
					flags.append(lab[i])
				elif lab[i] in '=*':
					flags.append( lab[i] )
					correction = float(lab[i+1:])
					break
				else:
					raise VoicingFileError, 'Bad flag: %s in %s' % (lab[i], lab)
				i += 1
			state = 'in'
			recognized = True

		if state == 'out' and not recognized:
			raise VoicingFileError, 'Something other than a beginning mark in "out" state: %.3f' % t

		if state == 'in' and not recognized:
			raise VoicingFileError, 'Something other than an end mark in "in" state: %.3f' % t

	if state == 'in':
		raise VoicingFileError, "No terminating ']'"

	ival = Num.nonzero(valid)
	vf0 = Num.take(of0, ival)
	if Num.sometrue( (vf0<LOWF0) + (vf0>HIGHF0) ):
		raise BadF0, 'F0 is out of (%g, %g) range somewhere' % (LOWF0, HIGHF0)

	die.info('f0correct: number of valid points %d -> %d' % (nv0, Num.sum(valid)))
	return (of0, valid, irr)



