# -*- 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 os
import re
import sys
import math

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

from select_paragraphs import misc

import generate_misc as GM
import phones as PH

sys.path.insert(0, '/proj/OLD/MRImodel/text/lib')
sys.path.insert(0, '/proj/OLD/MRImodel/text/bin')
import unisyn
NotInLexicon = unisyn.NotInLexicon
import asr_greedy

Unisyn_Cache = '/tmp/ucache.gdbm'

ROOT = '/proj/aesop'
RESP = os.path.join(ROOT, '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 read_phones_nt(pfn):
	"""Used for the Rhythm2010 main corpus; /proj/aesop/data_files/metadata"""
	pat = re.compile('[.]phones$')
	pfn2 = pat.sub('.phones2', pfn)
	try:
		fd = open(pfn2, 'r')
		priority = 1
	except IOError:
		die.info("No phone2 file: using %s instead" % pfn)
		fd = open(pfn, 'r')
		priority = 2
	pa = []
	for line in fd:
		if line.startswith('# '):
			continue
		for p in line.strip().split():
			pa.append(PH.phone_c(p))
	return (pa, priority)


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



class find_manual_segmentations_from_name(object):
	"""Find manual segmentation if you have a database of
	metadata and if the filename of the manual segmentation
	is stored in the metadata file.
	"""

	def __init__(self, manDBlist, man_seg_key='man_lab'):
		self.index = dictops.dict_of_lists()
		for manDB in manDBlist:
			dbd = os.path.dirname(manDB)
			dbf = manDB
			h, d, c = fiatio.read(open(dbf, 'r'))
			for x in d:
				df = os.path.join(x['d'], x['f'])
				self.index.add(df, os.path.join(dbd, x[man_seg_key]))

	def __call__(self, info):
		return self.index[os.path.join(info['d'], info['f'])]


class find_manual_segmentations_from_pat(object):
	"""Find manual segmentation if you have a database of
	metadata and if you know regular expressions that might match
	the manual segmentations, and if they are in the d/f directory.
	"""

	def __init__(self, manDBlist, accept, deny=[]):
		self.index = dictops.dict_of_lists()
		for manDB in manDBlist:
			dbd = os.path.dirname(manDB)
			dbf = manDB
			h, d, c = fiatio.read(open(dbf, 'r'))
			for x in d:
				df = os.path.join(x['d'], x['f'])
				ddf = os.path.join(dbd, df)
				for x in os.listdir(ddf):
					for pat in accept:
						if pat.search(x):
							denied = False
							for pat in deny:
								if pat.search(x):
									denied = True
									break
							if not denied:
								self.index.add(df, os.path.join(ddf, x))

	def __call__(self, info):
		return self.index[os.path.join(info['d'], info['f'])]





def get_manual_transcription(phonefile):
	return [PH.phone_c(lbl, s=s, e=e) for (s, e, lbl) in WS.read(phonefile)[1]]


def get_transcription(info):
	phonefile = os.path.join(info['dbdir'], info.get('trans_real', info['trans_exp']))
	return read_phones_nt(phonefile) + (phonefile,)






def add_to_dict(data, **kv):
	"""Add some keys to each dictionary in a list of dictionaries."""
	o = []
	for datum in data:
		dd = datum.copy()
		dd.update(kv)
		o.append(dd)
	return o



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
		



#: Maps from IViE dialect names to roughly correct dialect names for UNISYN
Dialect_map = {'Belfast': 'ccl1', 'Cambridge': 'rpx', 'Dublin': 'ccl1',
		'Leeds': 'lds', 'Jamaicain': 'rpx',
		'Punjabi': 'lds', 'Newcastle': 'lds'
		}


def IViE_rewrite_transcription_name(info):
	if info['textfile'].startswith('fragments'):
		return info['textfile']
	elif 'Cinderella' in info['textfile']:
		part = 'Cinderella'
	else:
		raise ValueError, "Cannot rewrite filename: no keyword: %s" % info['textfile']

	return os.path.join(ROOT, 'IViE/proc/transcriptions',
				'%s_%s' % (info['dialect'], info['speaker']),
				'%s_%s_%s.txt' % (part, info['dialect'], info['speaker'])
				)


Map_to_XSAMPA = dictops.read2c(open('/proj/aesop/IViE/Unisyn_to_xsampa_map_MRI.txt', 'r'))[0]
Map_to_XSAMPA.update({'<BOS>': '<BOS>', '<EOS>': '<EOS>', '<W>': '<W>',
			'<B1>': '<B1>', '<B2>': '<B2>', '<B3>': '<B3>'
			})


_punct = re.compile("""[,.;":!?]+""")
_unisyn_transcriber = None
def get_IViE_transcription(info):
	global _unisyn_transcriber
	try:
		textfile = open(os.path.join(info['dbdir'], info['textfile']), 'r')
	except IOError:
		textfn = IViE_rewrite_transcription_name(info)
		textfile = open(textfn, 'r')
	text = ' '.join([q.strip() for q in textfile.readlines()])
	text = _punct.sub(',', text.strip())

	if len(text)==0:
		raise GM.BadPhonesError, "No phones in data files"
	while text and text[-1] in "!?.,;:":
		text = text[:-1]
	text = text + '.'
	unisyn_dialect = Dialect_map[info['dialect']]
	if _unisyn_transcriber is None:
		_unisyn_transcriber = unisyn.unisyn_split(unisyn_dialect, cache=Unisyn_Cache)
	_unisyn_transcriber.set_dialect(unisyn_dialect)
	s = asr_greedy.Sentence(text, 0, _unisyn_transcriber)
	plist = [PH.phone_c(Map_to_XSAMPA[q]) for q in s.symbols]
	return (plist, 1, "Unisyn from %s" % textfile.name)



def read_IViE_VY(metadatafile):
	pat = re.compile('/([a-zA-Z]+)_([a-z][mf][0-9])_([0-9])/')
	h, infolist, c = fiatio.read(metadatafile)
	mdfn = metadatafile.name
	for info in infolist:
		dbdir = os.path.dirname(mdfn)
		info['metadatafile'] = mdfn
		info['dbdir'] = dbdir
		m = pat.search(info['d'])
		if m:
			info['sex'] = m.group(2)[1]
			info['speaker'] = m.group(2)
			info['dialect'] = m.group(1)
	return infolist


def NotEnoughData(Exception):
	def __init__(self, *s):
		Exception.__init__(self, *s)


def read_Rhythm_metadata(metadatafile):
	"""Useful for the Rhythm project metadata files."""
	h, infolist, c = fiatio.read(metadatafile)
	if not infolist:
		die.warn("Odd: no metadata in %s" % metadatafile)
	mdfn = metadatafile.name
	dbdir = os.path.dirname(mdfn)
	for info in infolist:
		info['metadatafile'] = mdfn
		info['dbdir'] = dbdir
	return infolist


def get_start_end(info):
	h = wavio.read_hdr(os.path.join(info['dbdir'], info['d'], info['f'], 'raw.wav'))
	return (h['CRVAL2'] + (1-h['CRPIX2'])*h['CDELT2'],
		h['CRVAL2'] + (h['NAXIS2']-h['CRPIX2'])*h['CDELT2']
		)
