from __future__ import print_function
from __future__ import division

import six
import collections


def convert_to_unicode(text):
    """
    Convert `text` to Unicode (if it's not already), assuming utf-8 input.
    """
    if six.PY3:
        if isinstance(text, str):
            return text
        elif isinstance(text, bytes):
            return text.decode("utf-8", "ignore")
        else:
            raise ValueError("Unsupported string type: %s" % (type(text)))
    elif six.PY2:
        if isinstance(text, str):
            return text.decode("utf-8", "ignore")
        elif isinstance(text, unicode):
            return text
        else:
            raise ValueError("Unsupported string type: %s" % (type(text)))
    else:
        raise ValueError("Not running on Python2 or Python 3?")


class Vocabulary(object):
    """
    Vocabulary class.
    """

    def __init__(self, vocab_file, num_relations, num_entities):
        self.vocab = self.load_vocab(vocab_file)
        self.inv_vocab = {v: k for k, v in self.vocab.items()}
        self.num_relations = num_relations
        self.num_entities = num_entities
        # 2 for special tokens of [PAD] and [MASK]
        assert len(self.vocab) == self.num_relations + self.num_entities + 2, \
            "The vocabulary contains all relations and entities, " \
            "as well as 2 special tokens: [PAD] and [MASK]."

    def load_vocab(self, vocab_file):
        """
        Load a vocabulary file into a dictionary.
        """
        vocab = collections.OrderedDict()
        with open(vocab_file, 'r') as fin:
            for num, line in enumerate(fin):
                items = convert_to_unicode(line.strip()).split("\t")
                if len(items) > 2:
                    break
                token = items[0]
                index = items[1] if len(items) == 2 else num
                token = token.strip()
                vocab[token] = int(index)
        return vocab

    def convert_by_vocab(self, vocab, items):
        """
        Convert a sequence of [tokens|ids] using the vocab.
        """
        output = []
        for item in items:
            output.append(vocab[item])
        return output

    def convert_tokens_to_ids(self, tokens):
        return self.convert_by_vocab(self.vocab, tokens)

    def convert_ids_to_tokens(self, ids):
        return self.convert_by_vocab(self.inv_vocab, ids)
