import tensorflow as tf
from transform import feature_spec, postprocess, create_tokenizer
import sys, os

flags = tf.flags
FLAGS = flags.FLAGS

flags.DEFINE_string("assets_dir", "assets", "The assets directory generated by assets.py.")
flags.DEFINE_string("tfrecords_file", "dataset/test_128.tfrecords", "The tfrecords file to be inspected.")
flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
flags.DEFINE_integer("display_rows", 1, "Number of rows to be displayed.")

def view_tfrecord(file_name, feature_spec, display_rows=5):
    """
    >>> view_tfrecord('./dataset/val_128.tfrecords', 128)
    """
    tokenizer = create_tokenizer(
        os.path.join(FLAGS.assets_dir, 'chinese_L-12_H-768_A-12', 'vocab.txt')
    )
    with tf.Session() as sess:
        i = 0
        for example in tf.python_io.tf_record_iterator(file_name):
            if i >= display_rows:
                break
            i+=1
            features = tf.parse_single_example(example, features=feature_spec)
            print('Example %d:'%i)
            for feature_name, tensor in features.items():
                if feature_name == 'text':
                    print(feature_name, features[feature_name].eval().values[0].decode('utf-8'))
                else:
                    print(feature_name, features[feature_name].eval())

            text = features['text'].eval().values[0].decode('utf-8')
            bert_tokens_len = features['bert_tokens_len'].eval()[0]
            bert_tokens = tokenizer.convert_ids_to_tokens(features['input_ids'].eval()[1:bert_tokens_len+1])
            bert_truths = features['truths'].eval()[1:bert_tokens_len+1]

            print("Final Result:", postprocess(text, bert_tokens, bert_truths, 0.5))



def main(_):
    tf.logging.set_verbosity(tf.logging.INFO)

    sys.path += [os.path.join(FLAGS.assets_dir, 'bert')]
    view_tfrecord(
        FLAGS.tfrecords_file,
        feature_spec(FLAGS.max_seq_length),
        FLAGS.display_rows
    )


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