#!/usr/bin/env python3

#Copyright      2017  Ashish Arora

""" This module will be used by scripts for open vocabulary setup.
 If the hypothesis transcription contains <unk>, then it will replace the 
 <unk> with the word predicted by <unk> model by concatenating phones decoded 
 from the unk-model. It is currently supported only for triphone setup.
 Args:
  phones: File name of a file that contains the phones.txt, (symbol-table for phones).
          phone and phoneID, Eg. a 217, phoneID of 'a' is 217. 
  words: File name of a file that contains the words.txt, (symbol-table for words). 
         word and wordID. Eg. ACCOUNTANCY 234, wordID of 'ACCOUNTANCY' is 234.
  unk: ID of <unk>. Eg. 231.
  one-best-arc-post: A file in arc-post format, which is a list of timing info and posterior
               of arcs along the one-best path from the lattice.
               E.g. 506_m01-049-00 8 12  1 7722  282 272 288 231
                    <utterance-id> <start-frame> <num-frames> <posterior> <word> [<ali>] 
                    [<phone1> <phone2>...]
  output-text: File containing hypothesis transcription with <unk> recognized by the
               unk-model.
               E.g. A move to stop mr. gaitskell.
  
  Eg. local/unk_arc_post_to_transcription.py lang/phones.txt lang/words.txt 
      data/lang/oov.int
"""
import argparse
import io
import os
import sys
parser = argparse.ArgumentParser(description="""uses phones to convert unk to word""")
parser.add_argument('phones', type=str, help='File name of a file that contains the'
                    'symbol-table for phones. Each line must be: <phone> <phoneID>')
parser.add_argument('words', type=str, help='File name of a file that contains the'
                    'symbol-table for words. Each line must be: <word> <word-id>')
parser.add_argument('unk', type=str, default='-', help='File name of a file that'
                    'contains the ID of <unk>. The content must be: <oov-id>, e.g. 231')
parser.add_argument('--one-best-arc-post', type=str, default='-', help='A file in arc-post'
                    'format, which is a list of timing info and posterior of arcs'
                    'along the one-best path from the lattice')
parser.add_argument('--output-text', type=str, default='-', help='File containing'
                    'hypothesis transcription with <unk> recognized by the unk-model')
args = parser.parse_args()

### main ###
phone_handle = open(args.phones, 'r', encoding='utf8') # Create file handles 
word_handle = open(args.words, 'r', encoding='utf8')
unk_handle = open(args.unk,'r', encoding='utf8')
if args.one_best_arc_post == '-':
    arc_post_handle = io.TextIOWrapper(sys.stdin.buffer, encoding='utf8')
else:
    arc_post_handle = open(args.one_best_arc_post, 'r', encoding='utf8')
if args.output_text == '-':
    output_text_handle = io.TextIOWrapper(sys.stdout.buffer, encoding='utf8')
else:
    output_text_handle = open(args.output_text, 'w', encoding='utf8')

id2phone = dict() # Stores the mapping from phone_id (int) to phone (char)
phones_data = phone_handle.read().strip().split("\n")

for key_val in phones_data:
  key_val = key_val.split(" ")
  id2phone[key_val[1]] = key_val[0]

word_dict = dict()
word_data_vect = word_handle.read().strip().split("\n")

for key_val in word_data_vect:
  key_val = key_val.split(" ")
  word_dict[key_val[1]] = key_val[0]
unk_val = unk_handle.read().strip().split(" ")[0]

utt_word_dict = dict() # Dict of list, stores mapping from utteranceID(int) to words(str)
for line in arc_post_handle:
  line_vect = line.strip().split("\t")
  if len(line_vect) < 6: # Check for 1best-arc-post output
    print("Error: Bad line: '{}'   Expecting 6 fields. Skipping...".format(line),
          file=sys.stderr)
    continue
  utt_id = line_vect[0]
  word = line_vect[4]
  phones = line_vect[5]
  if utt_id not in list(utt_word_dict.keys()):
    utt_word_dict[utt_id] = list()

  if word == unk_val: # Get the 1best phone sequence given by the unk-model
    phone_id_seq = phones.split(" ")
    phone_seq = list()
    for pkey in phone_id_seq:
      phone_seq.append(id2phone[pkey]) # Convert the phone-id sequence to a phone sequence.
    phone_2_word = list()
    for phone_val in phone_seq:
      phone_2_word.append(phone_val.split('_')[0]) # Removing the world-position markers(e.g. _B)
    phone_2_word = ''.join(phone_2_word) # Concatnate phone sequence
    utt_word_dict[utt_id].append(phone_2_word) # Store word from unk-model
  else:
    if word == '0': # Store space/silence
      word_val = ' '
    else:
      word_val = word_dict[word]
    utt_word_dict[utt_id].append(word_val) # Store word from 1best-arc-post

transcription = "" # Output transcription
for utt_key in sorted(utt_word_dict.keys()):
  transcription = utt_key
  for word in utt_word_dict[utt_key]:
    transcription = transcription + " " + word
  output_text_handle.write(transcription + '\n')
