#!/usr/bin/env python3
  
  # Copyright  2017  Jian Wang
  # License: Apache 2.0.
  
  import os
  import argparse
  import sys
  import math
  from collections import defaultdict
  sys.stdout = open(1, 'w', encoding='utf-8', closefd=False)
  
  import re
  
  
  parser = argparse.ArgumentParser(description="This script chooses the sparse feature representation of words. "
                                               "To be more specific, it chooses the set of features-- you compute "
                                               "them for the specific words by calling rnnlm/get_word_features.py.",
                                   epilog="E.g. " + sys.argv[0] + " --unigram-probs=exp/rnnlm/unigram_probs.txt "
                                          "--unigram-scale=0.1 "
                                          "data/rnnlm/vocab/words.txt > exp/rnnlm/features.txt",
                                   formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  
  parser.add_argument("--unigram-probs", type=str, default='', required=True,
                      help="Specify the file containing unigram probs.")
  parser.add_argument("--min-ngram-order", type=int, default=2,
                      help="minimum length of n-grams of characters to"
                           "make potential features.")
  parser.add_argument("--max-ngram-order", type=int, default=3,
                      help="maximum length of n-grams of characters to"
                           "make potential features.")
  parser.add_argument("--min-frequency", type=float, default=1.0e-05,
                      help="minimum frequency with which an n-gram character "
                           "feature is encountered (counted as binary presence in a word times unigram "
                           "probs of words), for it to be used as a feature. e.g. "
                           "if less than 1.0e-06 of tokens contain the n-gram 'xyz', "
                           "then it wouldn't be used as a feature.")
  parser.add_argument("--include-unigram-feature", type=str, default='true',
                      choices=['true', 'false'],
                      help="If true, the unigram frequency of a word is "
                           "one of the features.  [note: one reason we "
                           "want to include this, is to make it easier to "
                           "port models to new vocabularies and domains].")
  parser.add_argument("--include-length-feature", type=str, default='true',
                      choices=['true', 'false'],
                      help="If true, the length in characters of a word is one of the features.")
  parser.add_argument("--top-word-features", type=int, default=2000,
                      help="The most frequent n words each get their own "
                           "special feature, in addition to any other features "
                           "that the word may naturally get.")
  parser.add_argument("--special-words", type=str, default='<s>,</s>,<brk>',
                      help="List of special words that get their own special "
                           "features and do not get any other features.")
  parser.add_argument("--use-constant-feature", type=str, default="false",
                      help="If set to true, we give a constant feature to all "
                      "words (to help model offsets).  The value will equal "
                      "the --max-feature-rms option.");
  parser.add_argument("--max-feature-rms", type=float, default=0.01,
                      help="maximum allowed root-mean-square value for any feature.")
  
  # dir=exp/rnnlm_tdnn_d
  # paste <(awk '{print $2}' $dir/config/unigram_probs.txt) <(awk '{$1="";print;}' $dir/word_feats.txt ) | awk '{freq=$1; num_feats=(NF-1)/2; for (n=1;n<=num_feats;n++) { a=n*2; b=n*2+1; rms[$a] += freq * $b*$b; }} END{for(k in rms) { print k, rms[k];}}' | sort -k2 -nr | head
  # 7180 9.99985e-05
  # 7019 9.99947e-05
  # ..
  
  parser.add_argument("vocab_file",
                      help="Path for vocab file")
  
  args = parser.parse_args()
  
  if args.use_constant_feature != "false" and args.use_constant_feature != "true":
      sys.exit(sys.argv[0] + ": --use-constant-feature must be true or false: {0}".format(
          args.use_constant_feature))
  if args.min_ngram_order < 1:
      sys.exit(sys.argv[0] + ": --min-ngram-order must be at least 1.")
  if args.max_ngram_order < args.min_ngram_order:
      sys.exit(sys.argv[0] + ": --max-ngram-order must be larger than or equal to --min-ngram-order.")
  
  SPECIAL_SYMBOLS = ["<eps>", "<s>", "<brk>"]
  
  
  # read the vocabulay file
  # Returns a pair (vocab, wordlist)
  # where 'vocab' is a dict mapping the string-valued word to a integer id.
  #  and 'wordlist' is a list indexed by integer id, that returns the string-valued word.
  def read_vocab(vocab_file):
      vocab = {}
      with open(vocab_file, 'r', encoding="utf-8") as f:
          for line in f:
              fields = line.split()
              assert len(fields) == 2
              if fields[0] in vocab:
                  sys.exit(sys.argv[0] + ": duplicated word({0}) in vocab: {1}"
                                         .format(fields[0], vocab_file))
              vocab[fields[0]] = int(fields[1])
  
      # check there is no duplication and no gap among word ids
      sorted_ids = sorted(vocab.values())
      for idx, id in enumerate(sorted_ids):
          assert idx == id
  
      vocab_size = 1 + max(vocab.values())
      wordlist = [ None] * vocab_size
      for word, index in vocab.items():
          assert wordlist[index] is None
          wordlist[index] = word
  
      if wordlist[0] != '<eps>' and wordlist[0] != '<EPS>':
          sys.exit(sys.argv[0] + ": expected word numbered zero to be epsilon.")
      return (vocab, wordlist)
  
  
  # read the unigram probs; returns a list indexed by integer
  # id of the word, which evaluates to the unigram prob of the word.
  def read_unigram_probs(unigram_probs_file):
      unigram_probs = []
      with open(unigram_probs_file, 'r', encoding="utf-8") as f:
          for line in f:
              fields = line.split()
              assert len(fields) == 2
              idx = int(fields[0])
              if idx >= len(unigram_probs):
                  unigram_probs.extend([None] * (idx - len(unigram_probs) + 1))
              unigram_probs[idx] = float(fields[1])
  
      for prob in unigram_probs:
          assert prob is not None
  
      return unigram_probs
  
  def get_feature_scale(rms):
      if rms > args.max_feature_rms:
          return '%.2g' % (args.max_feature_rms / rms)
      else:
          return "1.0"
  
  (vocab, wordlist) = read_vocab(args.vocab_file)
  unigram_probs = read_unigram_probs(args.unigram_probs)
  assert len(unigram_probs) == len(wordlist)
  
  # num_features is a counter used to keep track of how many features
  # we've created so far.
  num_features = 0
  
  # The constant feature; this will be a line of the form
  # <feature-index> constant feature-value
  # e.g.
  # constant 0.01
  # which means every single word gets this feature.  It is used to
  # handle offsets of the words' log-likelihoods, so we don't have
  # to include an offset term in the math.
  if args.use_constant_feature == "true":
      print("{0}\tconstant\t{1}".format(num_features,
                                        args.max_feature_rms))
      num_features += 1
  
  
  # 'word_indexes_to_exclude' will contain the integer indexes of words that are
  # in 'args.special_words' plus the zero word (epsilon) and which don't take part
  # in later-numbered features.
  word_indexes_to_exclude = {0} # a set including only zero.
  
  # Features for 'special' words, i.e. a line of the form
  # <feature-index> special <word> <feature-value>
  # e.g.:
  # 2 special </s> 1.0
  # These words get just the constant feature (if present) and their own 'special'
  # feature, but not letter-based features, because things like '<s>' are
  # special symbols that are not treated as regular words.
  if args.special_words != '':
      for word in args.special_words.split(','):
          if not word in vocab:
              sys.exit(sys.argv[0] + ": error: element {0} of --special-words option "
                       "is not in the vocabulary file {1}".format(word, args.vocab_file))
          word_indexes_to_exclude.add(vocab[word])
          this_word_prob = unigram_probs[vocab[word]]
          rms = math.sqrt(this_word_prob)
          print("{0}\tspecial\t{1}\t{2}".format(num_features, word,
                                                get_feature_scale(rms)))
          num_features += 1
  
  
  # Print a line for the unigram feature (this is a feature that's a scaled,
  # offset version of the log-unigram-prob of the word).  The line is of the form:
  # <feature-index> unigram <offset> <scale>
  # e.g.
  # 3 unigram 0.04631 0.024312
  # where the interpretation is that if a word w has unigram probability p(w), the
  # feature's value will equal <offset> + <scale> * log(p(w)).
  
  # The offset and scale are chosen so that the expected value of the feature
  # is zero and its rms value equals args.max_feature_rms.
  if args.include_unigram_feature == 'true':
      total_p = 0.0  # total probability of words that have the unigram feature,
                     # i.e. excluding words with the 'special' feature.
      total_x = 0.0
      total_x2 = 0.0
      for idx, p in enumerate(unigram_probs):
          if p > 0.0 and idx not in word_indexes_to_exclude:
              # 'feature_value' is the value of the log-unigram-prob feature
              # before the offset and scale are accounted for.  We accumulate
              # the expected x and x^2 stats of this.
              feature_value = math.log(p)
              total_p += p
              total_x += p * feature_value
              total_x2 += p * feature_value * feature_value
      # we won't allow all the words to be 'special' words.
      # total_p is the probability mass of non-special words.
      assert total_p > 0 and total_p < 1.01
      mean = total_x / total_p
      variance = (total_x2 / total_p - mean * mean)
      # The following assert is because training an RNNLM with only one
      # 'non-special' word and the unigram feature (or using the unigram feature
      # where all unigram probs are the same) does not make sense.
      assert variance > 0
      # The mean is computed over those words where the feature was present..
      # 'stddev' is computed over all words, even when the feature was not present
      # (that's what the factor of 'total_p' is about); this is consistent with
      # how we apply the args.max_feature_rms option in general.
      stddev = math.sqrt(total_p * variance)
      scale = min(args.max_feature_rms / stddev, 1.0)
      offset = -mean * scale
      print("{0}\tunigram\t{1}\t{2}".format(num_features, offset, scale))
      num_features += 1
  
  # length feature.  This feature is the length of the word, scaled
  # down so that the rms does not exceed args.max_feature_rms.
  # The format of the line is:
  # <feature-index> length <scale>
  # e.g.:
  # 4 length 0.00518
  if args.include_length_feature == 'true':
      feature_sumsq = 0.0
      for word_index, p in enumerate(unigram_probs):
          if word_index not in word_indexes_to_exclude:
              word = wordlist[word_index]
              feature_value = len(word)
              feature_sumsq += p * feature_value * feature_value
      rms = math.sqrt(feature_sumsq)
      print("{0}\tlength\t{1}".format(num_features, get_feature_scale(rms)))
      num_features += 1
  
  # top-words features.  This is a feature that we assign to the top n most
  # frequent words (e.g. the top 2000 most frequent words), *in addition* to any
  # features they may get as a result of their written form.
  # e.g. the line will look like:
  # <feature-index> word <word> <scale>
  # e.g.:
  # 6 word of 1.0
  
  # We need to remember which words are given these features to avoid printing
  # the same feature later on when we process the n-gram type features; for
  # instance, if we are including trigrams on characters, the 3-gram (BOS, a, EOS)
  # would be the word "a", and if the word "a" got its own "word" feature due to
  # the --top-word-features option, this would be a duplication.
  # 'top_words' will be a set containing words (strings) that already had their
  # own 'word' feature.
  top_words = set()
  if args.top_word_features > 0:
      # sorted_word_indexes is a sorted list of pairs (word_index, unigram_prob),
      # sorted from greatest to least unigram_prob.
      sorted_word_indexes = sorted(enumerate(unigram_probs),
                                   key=lambda x: x[1], reverse=True)
      num_top_words_printed = 0
      for word_index, unigram_prob in sorted_word_indexes:
          if word_index in word_indexes_to_exclude:
              continue
          word = wordlist[word_index]
          rms = math.sqrt(unigram_prob)
          print("{0}\tword\t{1}\t{2}".format(num_features, word, get_feature_scale(rms)))
          num_features += 1
          top_words.add(word)
          num_top_words_printed += 1
          if num_top_words_printed > args.top_word_features:
              break
  
  # n-gram features.
  # These are basically sub-strings of a word.  What we print in the feature
  # file is things like:
  #  <feat-index> (initial|final|match|word) <sub-string> <scale>
  # e.g.
  # 20 match tel 1.0
  # where: 'match' means a match at any position in the word,
  #        'final' means a match that must be word-final (and may or may
  #                not be word-initial),
  #        'initial' means a match that must be word-initial (and may or
  #                may not be word-final)
  #        'word' means a match that is both word-initial and word-final.
  # For a given word, the feature value if there is a match will be the number of
  # matches times the feature scale.
  
  #  'ngram_feats' is a dict mapping a pair (match_type, match)
  #  to a pair (feat_freq, expected_feat_sumsq), where:
  #
  #   match_type (a string) is one of: 'match', 'final', 'initial', 'word',
  #           describing the match type, as explained above.
  #   match (a string) is the string that we are matching, e.g. 'ing'.
  #   feat_freq (a float) is the  sum over all words which the feature
  #     appears in, of the probability of that word.
  #   expected_feat_sumsq (a float) is the sum over all words of the probability
  #    of that word, times the square of the number of times the feature
  #    appears there.
  
  # if you index ngram_feats with a key that wasn't present, you'll get the
  # tuple (0.0, 0.0).
  ngram_feats = defaultdict(lambda: (0.0, 0.0))
  
  for (word_index, unigram_prob) in enumerate(unigram_probs):
      if word_index in word_indexes_to_exclude:
          continue
      word = wordlist[word_index]
  
      for pos in range(len(word) + 1):  # +1 for EOW 'this_word_feats' is a dict
          # from pairs (match_type, match) as defined above, to the count of the
          # number of times the feature was matched in this word (this count can
          # only be >1 if match_type == 'match').
          this_word_feats = defaultdict(int)
  
          for order in range(args.min_ngram_order, args.max_ngram_order + 1):
              start = pos - order + 1
              end = pos + 1
  
              if start < -1:
                  continue
  
              if start < 0 and end > len(word):
                  match_type = 'word'
                  start = 0
                  end = len(word)
              elif start < 0:
                  match_type = 'initial'
                  start = 0
              elif end > len(word):
                  match_type = 'final'
                  end = len(word)
              else:
                  match_type = 'match'
              if start >= end:
                  continue
  
              match = word[start:end]
              this_word_feats[(match_type, match)] += 1
          for (match_type, match), count in this_word_feats.items():
              (feat_freq, expected_feat_sumsq) = ngram_feats[(match_type, match)]
              ngram_feats[(match_type, match)] = (feat_freq + unigram_prob,
                                                  expected_feat_sumsq + unigram_prob * count * count)
  
  
  for (match_type, match), (expected_feat_sum, expected_feat_sumsq) in sorted(ngram_feats.items()):
      if match_type == 'word' and match in top_words:
          continue  # avoid duplicate
      if expected_feat_sum < args.min_frequency:
          continue  # very infrequent features are excluded via this mechanism.
      rms = math.sqrt(expected_feat_sumsq)
      print("{0}\t{1}\t{2}\t{3}".format(
          num_features, match_type, match, get_feature_scale(rms)))
      num_features += 1
  
  
  print(sys.argv[0] + ": chose {0} features.".format(num_features), file=sys.stderr)