# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Utility functions for BERT."""

import sys, os
import time
import logging
import collections
import hashlib
import io

import mxnet as mx
import gluonnlp as nlp

__all__ = ['tf_vocab_to_gluon_vocab', 'load_text_vocab']


def tf_vocab_to_gluon_vocab(tf_vocab):
    special_tokens = ['[UNK]', '[PAD]', '[SEP]', '[MASK]', '[CLS]']
    assert all(t in tf_vocab for t in special_tokens)
    counter = nlp.data.count_tokens(tf_vocab.keys())
    vocab = nlp.vocab.BERTVocab(counter, token_to_idx=tf_vocab)
    return vocab


def get_hash(filename):
    sha1 = hashlib.sha1()
    with open(filename, 'rb') as f:
        while True:
            data = f.read(1048576)
            if not data:
                break
            sha1.update(data)
    return sha1.hexdigest(), str(sha1.hexdigest())[:8]


def read_tf_checkpoint(path):
    """read tensorflow checkpoint"""
    from tensorflow.python import pywrap_tensorflow  # pylint: disable=import-outside-toplevel
    tensors = {}
    reader = pywrap_tensorflow.NewCheckpointReader(path)
    var_to_shape_map = reader.get_variable_to_shape_map()
    for key in sorted(var_to_shape_map):
        tensor = reader.get_tensor(key)
        tensors[key] = tensor
    return tensors

def profile(curr_step, start_step, end_step, profile_name='profile.json',
            early_exit=True):
    """profile the program between [start_step, end_step)."""
    if curr_step == start_step:
        mx.nd.waitall()
        mx.profiler.set_config(profile_memory=False, profile_symbolic=True,
                               profile_imperative=True, filename=profile_name,
                               aggregate_stats=True)
        mx.profiler.set_state('run')
    elif curr_step == end_step:
        mx.nd.waitall()
        mx.profiler.set_state('stop')
        logging.info(mx.profiler.dumps())
        mx.profiler.dump()
        if early_exit:
            sys.exit(0)

def load_text_vocab(vocab_file):
    """Loads a vocabulary file into a dictionary."""
    vocab = collections.OrderedDict()
    index = 0
    with io.open(vocab_file, 'r') as reader:
        while True:
            token = reader.readline()
            if not token:
                break
            token = token.strip()
            vocab[token] = index
            index += 1
    return vocab

def out_total_time(start, time_file='run_times.txt'):
    total_time = time.time() - start
    logging.info('Total time cost = %.2f sec', total_time)
    file = open(time_file, 'a')
    str = '%.2f\n' % total_time
    file.write(str)
    file.close()

def get_model(model_name, dataset, pretrained, ctx, use_pooler, logging, vocab=None):
    get_model_params = {
        'name': model_name,
        'dataset_name': dataset,
        'vocab': vocab,
        'pretrained': pretrained,
        'ctx': ctx,
        'use_pooler': use_pooler,
        'use_decoder': False,
        'use_classifier': False,
    }
    for i in range(2):
        try:
            return nlp.model.get_model(**get_model_params)
        except:
            if i == 0:
                logging.info("Second attempt to get model:")
            else:
                logging.info("Cannot get model, exiting...")
                sys.exit(1)

def remove_data(log_file_name=None, task_name=None):
    if os.getenv("REMOVE_DATA") == 'YES':
        if log_file_name:
            log_file = log_file_name(task_name)
            if os.path.exists(log_file):
                os.remove(log_file)

    for dir in ('/opt/mxnet/datasets', '/opt/mxnet/models', 'output_dir'):
        os.system('rm -rf ' + dir)


def report_train_time(start_time, log_file_name, task_name):
    out_total_time(start_time)
    remove_data(log_file_name, task_name)