# -*- coding: utf-8 -*-
#! python
from __future__ import with_statement

"""This script reads in the fiat file with data, finds all manually segmented files and creates symlinks
to mfcc and lab files for training data set for HTK. Can be run with mfcc.bash
Arguments: name_of_mfcc_file (no extension) metadata_file
The metadata file is in fiat format; each line describes one utterance.
"""

import math

import fpconst as FP

from gmisclib import die
from gmisclib import fiatio
from gmisclib import dictops
from gmisclib import gpkmisc
# from gmisclib import threaded_io
from gmisclib import nicknames
from gmisclib import wavesurfer_lab as WS

from select_paragraphs import misc

import phones as PH
import generate_misc as GM

RESP = '/proj/aesop/proc'
MFCC = 'mfcc'
LAB = 'lab'
GRAM = 'gram'
# Drop English ASR sentences and also group!=1 in English
# This helps balance the corpus
Vfile = '/proj/aesop/select_paragraphs/List_of_phones.txt'
TONES = ('_0', '_1', '_2', '_3', '_4')



def isV(phonology, lbl):
	properties = phonology[lbl]
	if properties.get('VorC', '') != '':
		return int(properties['VorC'])==1
	return False


def isC(phonology, lbl):
	properties = phonology[lbl]
	if properties.get('VorC', '') != '':
		return int(properties['VorC'])==2
	return False




# PP = {'sonority': prep_p_transitions.prep_phones,
# 	'VorC': prep_p_transitions.prep_phones,
# 	'I': prep_p_simple.prep_phones,
# 	'fa': prep_p_fa.prep_phones
# 	}


def replace_X_with_XX(phones):
	"""Map transitions between and among rare garble phones (X)
	into a single generic garble transition.
	"""
	for ph in phones:
		if ph.lbl.startswith('IPP_') and 'X' in ph.lbl:
			ph.lbl = 'IPP_XX'
		if ph.lbl.startswith('IWP_') and 'X' in ph.lbl:
			ph.lbl = 'IWP_XX'



def fill_NaNs_in_phone_sequence(phones, start, end):
	if FP.isNaN(phones[0].s):
		phones[0].s = start
	if FP.isNaN(phones[-1].e):
		phones[-1].e = end
	changed = True
	while changed:
		changed = False
		for i in range(0, len(phones)):
			if phones[i].dur is not None:
				if(i>0 and i+1<len(phones)
						and FP.isNaN(phones[i].s) and FP.isNaN(phones[i].e)
						and not FP.isNaN(phones[i-1].e)
						and not FP.isNaN(phones[i+1].s)):
					t = 0.5*(phones[i+1].s+phones[i-1].e)
					phones[i-1].e = t-phones[i].dur/2.0
					phones[i].s = phones[i-1].e
					phones[i].e = t+phones[i].dur/2.0
					phones[i+1].s = phones[i].e
					changed = True
				elif FP.isNaN(phones[i].e) and not FP.isNaN(phones[i].s):
					phones[i].e = phones[i].s + phones[i].dur
					changed = True
				elif FP.isNaN(phones[i].s) and not FP.isNaN(phones[i].e):
					phones[i].s = phones[i].e - phones[i].dur
					changed = True
			if (i+1<len(phones)
					and FP.isNaN(phones[i].e)
					and not FP.isNaN(phones[i+1].s)
					):
				phones[i].e = phones[i+1].s
				changed = True
			elif i>0 and FP.isNaN(phones[i].s) and not FP.isNaN(phones[i-1].e):
				phones[i].s = phones[i-1].e
				changed = True







def phone_name_distance(p, q):
	d = 0.0
	for (i, (pp, qq)) in enumerate(zip(p, q)):
		if pp != qq:
			d += 1.0/math.sqrt(1.0 + i)
	for i in range(min(len(p), len(q)), max(len(p), len(q))):
		d += 0.5/math.sqrt(1.0 + i)
	return d



def merge_rare_phones(odatalist, toofew=20):
	assert 1<toofew<10000
	phonelist = dictops.dict_of_accums()
	for datum in odatalist:
		for p in datum['_phones']:
			for q in p.labels():
				phonelist.add(q, 1)
	print 'phonelist=', phonelist
	to_be_merged = [ p for (p, count) in phonelist.items() if count<=toofew ]
	merge_targets = [ p for (p, count) in phonelist.items() if count>toofew ]
	assert merge_targets
	pmap = {}
	for p in to_be_merged:
		dbest = None
		qbest = None
		for q in merge_targets:
			d = phone_name_distance(p, q)
			if dbest is None or d < dbest:
				qbest = q
				dbest = d
		print 'Mapping from %s to %s' % (p, qbest)
		pmap[p] = qbest

	def phone_mapping(q):
		return pmap.get(q, q)

	for pdata in odatalist:
		for p in pdata['_phones']:
			if isinstance(p.lbl, str):
				p.lbl = phone_mapping(p.lbl)
			else:
				p.lbl = p.lbl.map( phone_mapping )
	return (odatalist, merge_targets)





# def fmt_or_list(p):
	# return '|'.join([HW.HTKformat(q) for q in p.labels()])
	





def make_pmap(phonology, mapname):
	pmap = {}
	for (phone,pinfo) in phonology.items():
		# print 'phone=', phone, 'pinfo=', pinfo
		if phone in ['#', '##']:
			out = 'sil'
		elif pinfo['non_phoneme'] != '':
			out = ''
		elif mapname == 'VorC':
			if pinfo['VorC'] == '':
				raise RuntimeError, "Missing VorC information for %s" % phone
			else:
				VorC = int(pinfo['VorC'])
			if VorC == 1:
				out = 'v'
			elif VorC == 2:
				out = 'c'
			else:
				out = 'sil'
			print 'pmap[%s] = %s' % (phone, out)
		elif mapname == 'sonority':
			if pinfo['non_phoneme']!='1' and pinfo['sonority'] == '':
				raise RuntimeError, "Missing VorC information for %s" % phone
			else:
				sonority = int(pinfo['sonority'])
			if sonority == 3:
				out = 'v'
			elif sonority==2:
				out = 'm'
			elif sonority==1:
				out = 'c'
			else:
				out = ''
		elif mapname == 'm_is_v':
			if pinfo['non_phoneme']!='1' and pinfo['sonority'] == '':
				raise RuntimeError, "Missing VorC information for %s" % phone
			else:
				sonority = int(pinfo['sonority'])
			if sonority == 3:
				out = 'v'
			elif sonority==2:
				out = 'v'
			elif sonority==1:
				out = 'c'
			else:
				out = ''
		elif mapname in ['I', 'fa']:
			if pinfo['non_phoneme']=='1':
				out = ''
			else:
				out = phone
		else:
			die.die("Whoops: mapname=%s for phone=%s" % (mapname, phone))
		pmap[phone] = out
		for tone in TONES:
			# This is a kluge to handle manual Chinese transcriptions:
			pmap['%s%s' % (phone, tone)] = out
	print 'pmap=', pmap
	return pmap


def get_most_frequent_vowel(pa, pmap, phonology):
	v_counts = dictops.dict_of_accums()
	for p in pa:
		if p.lbl in pmap and isV(phonology, p.lbl):
			v_counts.add(pmap[p.lbl], 1)
	if not v_counts:
		return None
	return dictops.argmax(v_counts)


def get_most_frequent_consonant(pa, pmap, phonology):
	c_counts = dictops.dict_of_accums()
	for p in pa:
		if p.lbl in pmap and isC(phonology, p.lbl):
			c_counts.add(pmap[p.lbl], 1)
	if not c_counts:
		return None
	return dictops.argmax(c_counts)


def add_edge_silence(phones):
	# The cutting of the data is not precise, so sometimes the silence
	# is fairly short:
	phones.append(PH.phone_c(PH.alt('sil', 'shortsil')))
	phones.insert(0, PH.phone_c(PH.alt('sil', 'shortsil')))


def junkify_edge_phones(phones):
	if phones[-1].is_all_sil():
		phones.pop(-1)
	else:
		phones[-1].lbl = PH.alt(phones[-1].likely()[0], GM.JunkPhone,
					PH.seq(phones[-1].likely()[0], GM.JunkPhone),
					)
	if phones[0].is_all_sil():
		phones.pop(0)
	else:
		phones[0].lbl = PH.alt(phones[0].likely()[-1], GM.JunkPhone,
					PH.seq(GM.JunkPhone, phones[0].likely()[-1])
					)

