#!/usr/bin/env python
  
  # Copyright 2016  Xiaohui Zhang
  # Apache 2.0.
  
  from __future__ import print_function
  from __future__ import division
  from collections import defaultdict
  import argparse
  import sys
  import math
  
  def GetArgs():
      parser = argparse.ArgumentParser(description = "Use a Bayesian framework to select"
                                       "pronunciation candidates from three sources: reference lexicon"
                                       ", G2P lexicon and phonetic-decoding lexicon. The inputs are a word-stats file,"
                                       "a pron-stats file, and three source lexicons (ref/G2P/phonetic-decoding)."
                                       "We assume the pronunciations for each word follow a Categorical distribution"
                                       "with Dirichlet priors. Thus, with user-specified prior counts (parameterized by"
                                       "prior-mean and prior-count-tot) and observed counts from the pron-stats file, "
                                       "we can compute posterior for each pron, and select candidates with highest"
                                       "posteriors, until we hit user-specified variants-prob-mass/counts thresholds."
                                       "The outputs are: a file specifiying posteriors of all candidate (pron_posteriors),"
                                       "a learned lexicon for words out of the ref. vocab (learned_lexicon_oov),"
                                       "and a lexicon_edits file containing suggested modifications of prons, for"
                                       "words within the ref. vocab (ref_lexicon_edits).",
                                       epilog = "See steps/dict/learn_lexicon_bayesian.sh for example.")
      parser.add_argument("--prior-mean", type = str, default = "0,0,0",
                          help = "Mean of priors (summing up to 1) assigned to three exclusive n"
                          "pronunciatio sources: reference lexicon, g2p, and phonetic decoding. We "
                          "recommend setting a larger prior mean for the reference lexicon, e.g. '0.6,0.2,0.2'")
      parser.add_argument("--prior-counts-tot", type = float, default = 15.0,
                          help = "Total amount of prior counts we add to all pronunciation candidates of"
                          "each word. By timing it with the prior mean of a source, and then dividing"
                          "by the number of candidates (for a word) from this source, we get the"
                          "prior counts we actually add to each candidate.")
      parser.add_argument("--variants-prob-mass", type = float, default = 0.7,
                          help = "For each word, we pick up candidates (from all three sources)"
                          "with highest posteriors until the total prob mass hit this amount.")
      parser.add_argument("--variants-prob-mass-ref", type = float, default = 0.9,
                          help = "For each word, after the total prob mass of selected candidates "
                          "hit variants-prob-mass, we continue to pick up reference candidates"
                          "with highest posteriors until the total prob mass hit this amount (must >= variants-prob-mass).")
      parser.add_argument("--variants-counts", type = int, default = 1,
                          help = "Generate upto this many variants of prons for each word out"
                          "of the ref. lexicon.")
      parser.add_argument("silence_file", metavar = "<silphonetic-file>", type = str,
                          help = "File containing a list of silence phones.")
      parser.add_argument("pron_stats_file", metavar = "<stats-file>", type = str,
                          help = "File containing pronunciation statistics from lattice alignment; "
                          "each line must be <count> <word> <phones>.")
      parser.add_argument("word_counts_file", metavar = "<counts-file>", type = str,
                          help = "File containing word counts in acoustic training data; "
                          "each line must be <word> <count>.")
      parser.add_argument("ref_lexicon", metavar = "<reference-lexicon>", type = str,
                          help = "The reference lexicon (most probably hand-derived)."
                          "Each line must be <word> <phones>")
      parser.add_argument("g2p_lexicon", metavar = "<g2p-expanded-lexicon>", type = str,
                          help = "Candidate ronouciations from G2P results."
                          "Each line must be <word> <phones>")
      parser.add_argument("phonetic_decoding_lexicon", metavar = "<prons-in-acoustic-evidence>", type = str,
                          help = "Candidate ronouciations from phonetic decoding results."
                          "Each line must be <word> <phones>")
      parser.add_argument("pron_posteriors", metavar = "<pron-posteriors>", type = str,
                          help = "Output file containing posteriors of all candidate prons for each word,"
                          "based on which we select prons to construct the learned lexicon."
                          "each line is <word> <pronunciation-source: one of R(ef)/G(2P)/P(hone-decoding)> <posterior> <pronunciation> ")
      parser.add_argument("learned_lexicon_oov", metavar = "<learned-lexicon-oov>", type = str,
                          help = "Output file which is the learned lexicon for words out of the ref. vocab.")
      parser.add_argument("ref_lexicon_edits", metavar = "<lexicon-edits>", type = str,
                          help = "Output file containing human-readable & editable pronounciation info (and the"
                          "accept/reject decision made by our algorithm) for those words in ref. vocab," 
                          "to which any change has been recommended. The info for each word is like:" 
                          "------------ an 4086.0 --------------"
                          "R  | Y |  2401.6 |  AH N"
                          "R  | Y |  640.8 |  AE N"
                          "P  | Y |  1035.5 |  IH N"
                          "R(ef), P(hone-decoding) represents the pronunciation source"
                          "Y/N means the recommended decision of including this pron or not"
                          "and the numbers are soft counts accumulated from lattice-align-word outputs. "
                          "See the function WriteEditsAndSummary for more details.")
  
  
      print (' '.join(sys.argv), file=sys.stderr)
  
      args = parser.parse_args()
      args = CheckArgs(args)
  
      return args
  
  def CheckArgs(args):
      args.silence_file_handle = open(args.silence_file)
      if args.pron_stats_file == "-":
          args.pron_stats_file_handle = sys.stdin
      else:
          args.pron_stats_file_handle = open(args.pron_stats_file)
      args.word_counts_file_handle = open(args.word_counts_file)
      args.ref_lexicon_handle = open(args.ref_lexicon)
      args.g2p_lexicon_handle = open(args.g2p_lexicon)
      args.phonetic_decoding_lexicon_handle = open(args.phonetic_decoding_lexicon)
      args.pron_posteriors_handle = open(args.pron_posteriors, "w")
      args.learned_lexicon_oov_handle = open(args.learned_lexicon_oov, "w")
      args.ref_lexicon_edits_handle = open(args.ref_lexicon_edits, "w")
      
      prior_mean = args.prior_mean.strip().split(',')
      if len(prior_mean) is not 3:
          raise Exception('Invalid Dirichlet prior mean ', args.prior_mean)
      for i in range(0,3):
          if float(prior_mean[i]) <= 0 or float(prior_mean[i]) >= 1:
              raise Exception('Dirichlet prior mean', prior_mean[i], 'is invalid, it must be between 0 and 1.')
      args.prior_mean = [float(prior_mean[0]), float(prior_mean[1]), float(prior_mean[2])]
  
      return args
  
  def ReadPronStats(pron_stats_file_handle):
      stats = {}
      for line in pron_stats_file_handle.readlines():
          splits = line.strip().split()
          if len(splits) == 0:
              continue
          if len(splits) < 2:
              raise Exception('Invalid format of line ' + line
                                  + ' in stats file.')
          count = float(splits[0])
          word = splits[1]
          phones = ' '.join(splits[2:])
          stats[(word, phones)] = count
      return stats
  
  def ReadWordCounts(word_counts_file_handle):
      counts = {}
      for line in word_counts_file_handle.readlines():
          splits = line.strip().split()
          if len(splits) < 2:
              raise Exception('Invalid format of line ' + line
                                  + ' in counts file.')
          word = splits[0]
          count = int(splits[1])
          counts[word] = count
      return counts
  
  def ReadLexicon(args, lexicon_file_handle, counts):
      # we're skipping any word not in counts (not seen in training data),
      # cause we're only learning prons for words who have acoustic examples.
      lexicon = defaultdict(set)
      for line in lexicon_file_handle.readlines():
          splits = line.strip().split()
          if len(splits) == 0:
              continue
          if len(splits) < 2:
              raise Exception('Invalid format of line ' + line
                                  + ' in lexicon file.')
          word = splits[0]
          if word not in counts:
              continue
          phones = ' '.join(splits[1:])
          lexicon[word].add(phones)
      return lexicon
  
  def FilterPhoneticDecodingLexicon(args, phonetic_decoding_lexicon, stats):
      # We want to remove all candidates which contains silence phones
      silphones = set()
      for line in args.silence_file_handle:
          silphones.add(line.strip())
      rejected_candidates = set()
      for word, prons in phonetic_decoding_lexicon.items():
          for pron in prons:
              for phone in pron.split():
                  if phone in silphones:
                     if (word, pron) in stats:
                         count = stats[(word, pron)]
                         del stats[(word, pron)]
                     else:
                         count = 0
                     rejected_candidates.add((word, pron))
                     print('WARNING: removing the candidate pronunciation from phonetic-decoding: {0}: '
                           '"{1}" whose soft-count from lattice-alignment is {2}, cause it contains at'
                           ' least one silence phone.'.format(word, pron, count), file=sys.stderr)
                     break
      for word, pron in rejected_candidates:
          phonetic_decoding_lexicon[word].remove(pron)
      return phonetic_decoding_lexicon, stats
  
  def ComputePriorCounts(args, counts, ref_lexicon, g2p_lexicon, phonetic_decoding_lexicon):
      prior_counts = defaultdict(list)
      # In case one source is absent for a word, we set zero prior to this source, 
      # and then re-normalize the prior mean parameters s.t. they sum up to one.
      for word in counts:
          prior_mean = [args.prior_mean[0], args.prior_mean[1], args.prior_mean[2]]
          if word not in ref_lexicon:
              prior_mean[0] = 0
          if word not in g2p_lexicon:
              prior_mean[1] = 0
          if word not in phonetic_decoding_lexicon:
              prior_mean[2] = 0
          prior_mean_sum = sum(prior_mean)
          try:
              prior_mean = [float(t) / prior_mean_sum for t in prior_mean] 
          except ZeroDivisionError:
              print('WARNING: word {} appears in train_counts but not in any lexicon.'.format(word), file=sys.stderr)
          prior_counts[word] = [t * args.prior_counts_tot for t in prior_mean] 
      return prior_counts
  
  def ComputePosteriors(args, stats, ref_lexicon, g2p_lexicon, phonetic_decoding_lexicon, prior_counts):
      posteriors = defaultdict(list) # This dict stores a list of (pronunciation, posterior)
      # pairs for each word, where the posteriors are normalized soft counts. Before normalization,
      # The soft-counts were augmented by a user-specified prior count, according the source 
      # (ref/G2P/phonetic-decoding) of this pronunciation.
  
      for word, prons in ref_lexicon.items():
          for pron in prons:
              # c is the augmented soft count (observed count + prior count)
              c = float(prior_counts[word][0]) / len(ref_lexicon[word]) + stats.get((word, pron), 0)
              posteriors[word].append((pron, c))
  
      for word, prons in g2p_lexicon.items():
          for pron in prons:
              c = float(prior_counts[word][1]) / len(g2p_lexicon[word]) + stats.get((word, pron), 0)
              posteriors[word].append((pron, c))
  
      for word, prons in phonetic_decoding_lexicon.items():
          for pron in prons:
              c = float(prior_counts[word][2]) / len(phonetic_decoding_lexicon[word]) + stats.get((word, pron), 0)
              posteriors[word].append((pron, c))
  
      num_prons_from_ref = sum(len(ref_lexicon[i]) for i in ref_lexicon)
      num_prons_from_g2p = sum(len(g2p_lexicon[i]) for i in g2p_lexicon)
      num_prons_from_phonetic_decoding = sum(len(phonetic_decoding_lexicon[i]) for i in phonetic_decoding_lexicon)
      print ("---------------------------------------------------------------------------------------------------", file=sys.stderr)
      print ('Total num. words is {}:'.format(len(posteriors)), file=sys.stderr)
      print ('{0} candidate prons came from the reference lexicon; {1} came from G2P;{2} came from'
             'phonetic_decoding'.format(num_prons_from_ref, num_prons_from_g2p, num_prons_from_phonetic_decoding), file=sys.stderr)
      print ("---------------------------------------------------------------------------------------------------", file=sys.stderr)
  
      # Normalize the augmented soft counts to get posteriors.
      count_sum = defaultdict(float) # This dict stores the pronunciation which has 
      # the sum of augmented soft counts for each word.
      
      for word in posteriors:
          # each entry is a pair: (prounciation, count)
          count_sum[word] = sum([entry[1] for entry in posteriors[word]])
      
      for word, entry in posteriors.items():
          new_entry = []
          for pron, count in entry:      
              post = float(count) / count_sum[word]
              new_entry.append((pron, post))
              source = 'R'
              if word in g2p_lexicon and pron in g2p_lexicon[word]:
                  source = 'G'
              elif word in phonetic_decoding_lexicon and pron in phonetic_decoding_lexicon[word]:
                  source = 'P'
              print(word, source, "%3.2f" % post, pron, file=args.pron_posteriors_handle)
          del entry[:]
          entry.extend(sorted(new_entry, key=lambda new_entry: new_entry[1]))
      return posteriors
  
  def SelectPronsBayesian(args, counts, posteriors, ref_lexicon, g2p_lexicon, phonetic_decoding_lexicon):
      reference_selected = 0
      g2p_selected = 0
      phonetic_decoding_selected = 0
      learned_lexicon = defaultdict(set)
  
      for word, entry in posteriors.items():
          num_variants = 0
          post_tot = 0.0
          variants_counts = args.variants_counts
          variants_prob_mass = args.variants_prob_mass
          if word in ref_lexicon:
              # the variants count of the current word's prons in the ref lexicon.
              variants_counts_ref = len(ref_lexicon[word])
              # For words who don't appear in acoustic training data at all, we simply accept all ref prons.
              # For words in ref. vocab, we set the max num. variants 
              if counts.get(word, 0) > 0:
                  variants_counts = math.ceil(1.5 * variants_counts_ref)
              else:
                  variants_counts = variants_counts_ref
                  variants_prob_mass = 1.0
          last_post = 0.0
          while ((num_variants < variants_counts and post_tot < variants_prob_mass)
                 or (len(entry) > 0 and entry[-1][1] == last_post)): # this conditions 
                 # means the posterior of the current pron is the same as the one we just included.
              try:
                  pron, post = entry.pop()
                  last_post = post
              except IndexError:
                  break
              post_tot += post
              learned_lexicon[word].add(pron)
              num_variants += 1
              if word in ref_lexicon and pron in ref_lexicon[word]:
                  reference_selected += 1
              elif word in g2p_lexicon and pron in g2p_lexicon[word]:
                  g2p_selected += 1
              else:
                  phonetic_decoding_selected += 1
  
          while (num_variants < variants_counts and post_tot < args.variants_prob_mass_ref):
              try:
                  pron, post = entry.pop()
              except IndexError:
                  break
              if word in ref_lexicon and pron in ref_lexicon[word]:
                  post_tot += post
                  learned_lexicon[word].add(pron)
                  num_variants += 1
                  reference_selected += 1
  
      num_prons_tot = reference_selected + g2p_selected + phonetic_decoding_selected
      print('---------------------------------------------------------------------------------------------------', file=sys.stderr)
      print ('Num. words in the learned lexicon: {0} num. selected prons: {1}'.format(len(learned_lexicon), num_prons_tot), file=sys.stderr)
      print ('{0} selected prons came from reference candidate prons; {1} came from G2P candidate prons;'
             '{2} came from phonetic-decoding candidate prons.'.format(reference_selected, g2p_selected, phonetic_decoding_selected), file=sys.stderr) 
      return learned_lexicon
  
  def WriteEditsAndSummary(args, learned_lexicon, ref_lexicon, phonetic_decoding_lexicon, g2p_lexicon, counts, stats):
      # Note that learned_lexicon and ref_lexicon are dicts of sets of prons, while the other two lexicons are sets of (word, pron) pairs.
      threshold = 3
      words = [defaultdict(set) for i in range(4)] # "words" contains four bins, where we
      # classify each word into, according to whether it's count > threshold,
      # and whether it's OOVs w.r.t the reference lexicon.
  
      src = {}
      print("# Note: This file contains pronunciation info for words who have candidate"
            "prons from G2P/phonetic-decoding accepted in the learned lexicon."
            ", sorted by their counts in acoustic training data, "
            ,file=args.ref_lexicon_edits_handle)
      print("# 1st Col: source of the candidate pron: G(2P) / P(hone-decoding) / R(eference)."
            ,file=args.ref_lexicon_edits_handle)
      print("# 2nd Col: accepted or not in the learned lexicon (Y/N).", file=args.ref_lexicon_edits_handle)
      print("# 3rd Col: soft counts from lattice-alignment (not augmented by prior-counts)."
            ,file=args.ref_lexicon_edits_handle)
      print("# 4th Col: the pronunciation cadidate.", file=args.ref_lexicon_edits_handle)
      
      # words which are to be printed into the edits file.
      words_to_edit = [] 
      for word in learned_lexicon:
          count = counts.get(word, 0)
          flags = ['0' for i in range(3)] # "flags" contains three binary indicators, 
          # indicating where this word's pronunciations come from.
          for pron in learned_lexicon[word]:
              if word in phonetic_decoding_lexicon and pron in phonetic_decoding_lexicon[word]:
                  flags[0] = '1'
                  src[(word, pron)] = 'P'
              if word in ref_lexicon and pron in ref_lexicon[word]:
                  flags[1] = '1'
                  src[(word, pron)] = 'R'
              if word in g2p_lexicon and pron in g2p_lexicon[word]:
                  flags[2] = '1'
                  src[(word, pron)] = 'G'
          if word in ref_lexicon:
              all_ref_prons_accepted = True
              for pron in ref_lexicon[word]:
                  if pron not in learned_lexicon[word]:
                      all_ref_prons_accepted = False
                      break
              if not all_ref_prons_accepted or flags[0] == '1' or flags[2] == '1':
                  words_to_edit.append((word, counts[word]))
              if count > threshold:
                  words[0][flags[0] + flags[1] + flags[2]].add(word)
              else:
                  words[1][flags[0] + flags[1] + flags[2]].add(word)
          else:
              if count > threshold: 
                  words[2][flags[0] + flags[2]].add(word)
              else:
                  words[3][flags[0] + flags[2]].add(word)
  
      words_to_edit_sorted = sorted(words_to_edit, key=lambda entry: entry[1], reverse=True)
      for word, count in words_to_edit_sorted:
          print("------------",word, "%2.1f" % count, "--------------", file=args.ref_lexicon_edits_handle)
          for pron in learned_lexicon[word]:
              print(src[(word, pron)], ' | Y | ', "%2.1f | " % stats.get((word, pron), 0), pron, 
                    file=args.ref_lexicon_edits_handle)
          for pron in ref_lexicon[word]:
              if pron not in learned_lexicon[word]:
                  soft_count = stats.get((word, pron), 0)
                  print('R  | N |  {:.2f} | {} '.format(soft_count, pron), file=args.ref_lexicon_edits_handle)
      print("Here are the words whose reference pron candidates were all declined", words[0]['100'], file=sys.stderr)
      print("-------------------------------------------------Summary------------------------------------------", file=sys.stderr)
      print("In the learned lexicon, out of those", len(ref_lexicon), "words from the vocab of the reference lexicon:", file=sys.stderr) 
      print("  For those frequent words whose counts in the training text > ", threshold, ":", file=sys.stderr) 
      num_freq_ivs_from_all_sources = len(words[0]['111']) + len(words[0]['110']) + len(words[0]['011'])
      num_freq_ivs_from_g2p_or_phonetic_decoding = len(words[0]['101']) + len(words[0]['001']) + len(words[0]['100'])
      num_freq_ivs_from_ref = len(words[0]['010'])
      num_infreq_ivs_from_all_sources = len(words[1]['111']) + len(words[1]['110']) + len(words[1]['011'])
      num_infreq_ivs_from_g2p_or_phonetic_decoding = len(words[1]['101']) + len(words[1]['001']) + len(words[1]['100'])
      num_infreq_ivs_from_ref = len(words[1]['010'])
      print(' {} words\' selected prons came from the reference lexicon, G2P/phonetic-decoding.'.format(num_freq_ivs_from_all_sources), file=sys.stderr)
      print(' {} words\' selected prons come from G2P/phonetic-decoding-generated.'.format(num_freq_ivs_from_g2p_or_phonetic_decoding), file=sys.stderr) 
      print(' {} words\' selected prons came from the reference lexicon only.'.format(num_freq_ivs_from_ref), file=sys.stderr) 
      print('  For those words whose counts in the training text <= {}:'.format(threshold), file=sys.stderr) 
      print(' {} words\' selected prons came from the reference lexicon, G2P/phonetic-decoding.'.format(num_infreq_ivs_from_all_sources), file=sys.stderr)
      print(' {} words\' selected prons come from G2P/phonetic-decoding-generated.'.format(num_infreq_ivs_from_g2p_or_phonetic_decoding), file=sys.stderr) 
      print(' {} words\' selected prons came from the reference lexicon only.'.format(num_infreq_ivs_from_ref), file=sys.stderr) 
      print("---------------------------------------------------------------------------------------------------", file=sys.stderr)
      num_oovs = len(learned_lexicon) - len(ref_lexicon)
      num_freq_oovs_from_both_sources = len(words[2]['11'])
      num_freq_oovs_from_phonetic_decoding = len(words[2]['10'])
      num_freq_oovs_from_g2p = len(words[2]['01'])
      num_infreq_oovs_from_both_sources = len(words[3]['11'])
      num_infreq_oovs_from_phonetic_decoding = len(words[3]['10'])
      num_infreq_oovs_from_g2p = len(words[3]['01'])
      print('  In the learned lexicon, out of those {} OOV words (w.r.t the reference lexicon):'.format(num_oovs), file=sys.stderr)
      print('  For those words whose counts in the training text > {}:'.format(threshold), file=sys.stderr)
      print('    {} words\' selected prons came from G2P and phonetic-decoding.'.format(num_freq_oovs_from_both_sources), file=sys.stderr)
      print('    {} words\' selected prons came from phonetic decoding only.'.format(num_freq_oovs_from_phonetic_decoding), file=sys.stderr) 
      print('    {} words\' selected prons came from G2P only.'.format(num_freq_oovs_from_g2p), file=sys.stderr) 
      print('  For those words whose counts in the training text <= {}:'.format(threshold), file=sys.stderr) 
      print('    {} words\' selected prons came from G2P and phonetic-decoding.'.format(num_infreq_oovs_from_both_sources), file=sys.stderr)
      print('    {} words\' selected prons came from phonetic decoding only.'.format(num_infreq_oovs_from_phonetic_decoding), file=sys.stderr) 
      print('    {} words\' selected prons came from G2P only.'.format(num_infreq_oovs_from_g2p), file=sys.stderr) 
  
  def WriteLearnedLexiconOov(learned_lexicon, ref_lexicon, file_handle):
      for word, prons in learned_lexicon.items():
          if word not in ref_lexicon:
              for pron in prons:
                  print('{0} {1}'.format(word, pron), file=file_handle)
      file_handle.close()
  
  def Main():
      args = GetArgs()
  
      # Read in three lexicon sources, word counts, and pron stats.
      counts = ReadWordCounts(args.word_counts_file_handle)
      ref_lexicon = ReadLexicon(args, args.ref_lexicon_handle, counts)
      g2p_lexicon = ReadLexicon(args, args.g2p_lexicon_handle, counts)
      phonetic_decoding_lexicon =  ReadLexicon(args, args.phonetic_decoding_lexicon_handle, counts)
      stats = ReadPronStats(args.pron_stats_file_handle)
      phonetic_decoding_lexicon, stats = FilterPhoneticDecodingLexicon(args, phonetic_decoding_lexicon, stats)
     
      # Compute prior counts
      prior_counts = ComputePriorCounts(args, counts, ref_lexicon, g2p_lexicon, phonetic_decoding_lexicon)
      # Compute posteriors, and then select prons to construct the learned lexicon.
      posteriors = ComputePosteriors(args, stats, ref_lexicon, g2p_lexicon, phonetic_decoding_lexicon, prior_counts)
  
      # Select prons to construct the learned lexicon.
      learned_lexicon = SelectPronsBayesian(args, counts, posteriors, ref_lexicon, g2p_lexicon, phonetic_decoding_lexicon)
      
      # Write the learned prons for words out of the ref. vocab into learned_lexicon_oov.
      WriteLearnedLexiconOov(learned_lexicon, ref_lexicon, args.learned_lexicon_oov_handle)
      # Edits will be printed into ref_lexicon_edits, and the summary will be printed into stderr.
      WriteEditsAndSummary(args, learned_lexicon, ref_lexicon, phonetic_decoding_lexicon, g2p_lexicon, counts, stats)
  
  if __name__ == "__main__":
      Main()