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

from __future__ import absolute_import
from __future__ import division
from text_encoder import SubwordTextEncoder
import tokenizer
import os
import tempfile

import tensorflow as tf

tf.flags.DEFINE_string('output_filename', '/tmp/my.subword_text_encoder',
                       'where to store the SubwordTextEncoder')
tf.flags.DEFINE_string('corpus_filepattern', '',
                       'Corpus of one or more text files')
tf.flags.DEFINE_string('vocab_filepattern', '', 'One or more vocabulary files '
                       '(one word per line as "word,count")')
tf.flags.DEFINE_integer('min_count', 5, 'Minimum subtoken count in corpus')
tf.flags.DEFINE_integer('vocab_size', 30000, 'The final vocab size. It will produce a vocab with a near vocab size')
tf.flags.DEFINE_integer('corpus_max_lines', None,
                        'How many lines of corpus to read')
tf.flags.DEFINE_integer('num_iterations', 5, 'Number of iterations')
tf.flags.DEFINE_bool('split_on_newlines', True, 'Break corpus into lines.')
tf.flags.DEFINE_string('additional_chars', "", 'Set special characters to be included in vocab. ex : "~", "/".')
tf.flags.DEFINE_integer('max_subtoken_length', None, 'Max subtoken length')
tf.flags.DEFINE_string('raw_vocab', None, 'Raw bert vovab file')
tf.flags.DEFINE_bool('do_lower_case', False, 'Whether or not to lowercase the input corpus')
FLAGS = tf.flags.FLAGS


def merge_output_file_with_bert_vocab(output_filename, bert_vocab, temp_path):
  writer = open(output_filename, 'w', encoding='utf-8')
  _set = set()
  with open(bert_vocab, 'r', encoding='utf-8') as reader:
      for line in reader:
          writer.write(line)
          _set.add(line.strip())
  print(temp_path)
  with open(temp_path, 'r', encoding='utf-8') as reader:
      for line in reader:
          if line.strip() not in _set:
              writer.write(line)

  writer.close()
  # os.remove(temp_path)

def main(unused_argv):
  if FLAGS.corpus_filepattern and FLAGS.vocab_filepattern:
    raise ValueError(
        'Must only provide one of --corpus_filepattern or --vocab_filepattern')

  elif FLAGS.corpus_filepattern:
    token_counts = tokenizer.corpus_token_counts(
        FLAGS.corpus_filepattern,
        FLAGS.corpus_max_lines,
        split_on_newlines=FLAGS.split_on_newlines, additional_chars=FLAGS.additional_chars, do_lower_case=FLAGS.do_lower_case)

  elif FLAGS.vocab_filepattern:
    token_counts = tokenizer.vocab_token_counts(FLAGS.vocab_filepattern,
                                                FLAGS.corpus_max_lines, FLAGS.do_lower_case)

  else:
    raise ValueError(
        'Must provide one of --corpus_filepattern or --vocab_filepattern')
  reserved_tokens = None
  if FLAGS.raw_vocab:
      lines = open(FLAGS.raw_vocab, 'r', encoding='utf-8').readlines()
      lines = [s.strip() for s in lines if len(s) > 0]
      reserved_tokens = lines

  print(len(token_counts))
  print(len(reserved_tokens))
  target_size = FLAGS.vocab_size
  if target_size <= len(reserved_tokens):
    raise ValueError("The vocab_size must be larger than the origin vocab's size ")
  if target_size >= len(token_counts):
    raise ValueError("The vocab_size is too large. Please set it smaller or prepare more corpus.")
  min_val = 1
  max_val = len(token_counts) // (target_size ** 0.5)  
  fd, temp_path = tempfile.mkstemp()
  encoder = SubwordTextEncoder.build_to_target_size(target_size,token_counts,min_val, max_val, num_iterations=FLAGS.num_iterations, 
                                                    reserved_tokens=reserved_tokens, max_subtoken_length=FLAGS.max_subtoken_length)
  # encoder = SubwordTextEncoder()
  # encoder.build_from_token_counts(token_counts, FLAGS.min_count,
  #                                 FLAGS.num_iterations, reserved_tokens=reserved_tokens, max_subtoken_length=FLAGS.max_subtoken_length)
  encoder.store_to_file(temp_path, add_single_quotes=False)
  merge_output_file_with_bert_vocab(FLAGS.output_filename, FLAGS.raw_vocab, temp_path)

if __name__ == '__main__':
  tf.app.run()