# -*- coding: utf-8 -*-

# Copyright 2015 Language Technology, Technische Universitaet Darmstadt (author: Benjamin Milde)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function

import argparse
import pickle
import io

#generate string for one entry of the dictionary
def generateEntry(word,entry,sphinx_format=False):
    freqMult = 0.0
    #find most frequent entry
    freqs = [pron['freq'] for pron in entry]
    freqMult = 1.0 / float(max(freqs))
    freqs = [float(pron['freq'])*freqMult for pron in entry]
    prons = [elem['pron'] for elem in entry]    

    txt = ''
    for i,(freq,pron) in enumerate(sorted(zip(freqs,prons), reverse=True, key=lambda x:x[0])):
        if sphinx_format:
            txt += word + ('('+str(i)+')' if i>0 else '') + '  ' + ' '.join(pron) + '\n'
        else:
            #Kaldi probabilty format <word> <freq> <pronounciation>
            txt += word + ' ' + str(freq) + ' ' + ' '.join(pron) + '\n'
    return txt

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Prepares the files from the TUDA corpus (XML) into text transcriptions for KALDI')
    parser.add_argument('-f', '--file', dest='file', help='process this (python pickle) lexicon file', type=str)
    parser.add_argument('-o', '--outfile', dest='outfile', help='lexicon out file', type=str, default='lexiconp.txt')
    parser.add_argument('-sph', '--sphinx-format', dest='sphinx_format', help='export lexicon in sphinx format', action='store_true', default=False)

    args = parser.parse_args()
    
    if args.sphinx_format:
        print('Note: will use the Sphinx lexicon for this export.')

    print('Load ', args.file)
    combinedDict = pickle.load( open( args.file , 'rb' ) )

    print('Succesfully opened pickle file, now exporting to:', args.outfile) 
    with io.open(args.outfile,'w',encoding='utf-8') as outfile:
        if '%' in combinedDict:
            combinedDict['<UNK>'] = combinedDict['%']
            print('<UNK> is:', combinedDict['<UNK>'])
            #del combinedDict['%']
        else:
            print('Warning!: No % entry found! Will add <UNK> -> usb mapping manually.')
            combinedDict['<UNK>'] = [{'pron': ['usb'], 'freq': 100, 'manual': True}]

        if '<Lachen>' not in combinedDict:
            print('Warning!: No <Lachen> entry found! Will add <Lachen> -> lau mapping manually.')
            combinedDict['<Lachen>'] = [{'pron': ['lau'], 'freq': 100, 'manual': True}]

        for key in sorted(list(combinedDict.keys())):
            txt = generateEntry(key,combinedDict[key],args.sphinx_format)
            outfile.write(txt)

