from bidict import bidict
import pickle
import logging

logger = logging.getLogger(__name__)


class Vocabulary():
    """This class maps strings to integers, which also allow many namespaces
    """

    DEFAULT_PAD_TOKEN = '*@PAD@*'
    DEFAULT_UNK_TOKEN = '*@UNK@*'

    def __init__(self,
                 counters=dict(),
                 min_count=dict(),
                 pretrained_vocab=dict(),
                 intersection_namespace=dict(),
                 no_pad_namespace=list(),
                 no_unk_namespace=list(),
                 contain_pad_namespace=dict(),
                 contain_unk_namespace=dict()):
        """initialize vocabulary

        Keyword Arguments:
            counters {dict} -- multiple counter (default: {dict()})
            min_count {dict} -- min count dict (default: {dict()})
            pretrained_vocab {dict} -- pretrained vocabulary (default: {dict()})
            intersection_namespace {dict} -- intersection namespace correspond to pretrained vocabulary in case of too large pretrained vocabulary (default: {dict()})
            no_pad_namespace {list} -- no paddding namespace (default: {list()})
            no_unk_namespace {list} -- no unknown namespace (default: {list()})
            contain_pad_namespace {dict} -- contain padding token namespace (default: {dict()})
            contain_unk_namespace {dict} -- contain unknown token namespace (default: {dict()})
        """

        self.min_count = dict(min_count)
        self.intersection_namespace = dict(intersection_namespace)
        self.no_pad_namespace = set(no_pad_namespace)
        self.no_unk_namespace = set(no_unk_namespace)
        self.contain_pad_namespace = dict(contain_pad_namespace)
        self.contain_unk_namespace = dict(contain_unk_namespace)
        self.vocab = dict()

        self.extend_from_counter(counters, self.min_count, self.no_pad_namespace,
                                 self.no_unk_namespace)

        self.extend_from_pretrained_vocab(pretrained_vocab, self.intersection_namespace,
                                          self.no_pad_namespace, self.no_unk_namespace)

        logger.info("Initialize vocabulary successfully.")

    def extend_from_pretrained_vocab(self,
                                     pretrained_vocab,
                                     intersection_namespace=dict(),
                                     no_pad_namespace=list(),
                                     no_unk_namespace=list(),
                                     contain_pad_namespace=dict(),
                                     contain_unk_namespace=dict()):
        """extend vocabulary from pretrained vocab

        Arguments:
            pretrained_vocab {dict} -- pretrained vocabulary

        Keyword Arguments:
            intersection_namespace {dict} -- intersection namespace correspond to pretrained vocabulary in case of too large pretrained vocabulary (default: {dict()})
            no_pad_namespace {list} -- no paddding namespace (default: {list()})
            no_unk_namespace {list} -- no unknown namespace (default: {list()})
            contain_pad_namespace {dict} -- contain padding token namespace (default: {dict()})
            contain_unk_namespace {dict} -- contain unknown token namespace (default: {dict()})
        """

        self.intersection_namespace.update(dict(intersection_namespace))
        self.no_pad_namespace.update(set(no_pad_namespace))
        self.no_unk_namespace.update(set(no_unk_namespace))
        self.contain_pad_namespace.update(dict(contain_pad_namespace))
        self.contain_unk_namespace.update(dict(contain_unk_namespace))

        for namespace, vocab in pretrained_vocab.items():
            self.__namespace_init(namespace)
            is_intersection = namespace in self.intersection_namespace
            intersection_vocab = self.vocab[
                self.intersection_namespace[namespace]] if is_intersection else []
            for key, value in vocab.items():
                if not is_intersection or key in intersection_vocab:
                    self.vocab[namespace][key] = value

            logger.info(
                "Vocabulay {} (size: {}) was constructed successfully from pretrained_vocab.".
                format(namespace, len(self.vocab[namespace])))

    def extend_from_counter(self,
                            counters,
                            min_count=dict(),
                            no_pad_namespace=list(),
                            no_unk_namespace=list(),
                            contain_pad_namespace=dict(),
                            contain_unk_namespace=dict()):
        """extend vocabulary from counter

        Arguments:
            counters {dict} -- multiply counter

        Keyword Arguments:
            min_count {dict} -- min count dict (default: {dict()})
            no_pad_namespace {list} -- no paddding namespace (default: {list()})
            no_unk_namespace {list} -- no unknown namespace (default: {list()})
            contain_pad_namespace {dict} -- contain padding token namespace (default: {dict()})
            contain_unk_namespace {dict} -- contain unknown token namespace (default: {dict()})
        """

        self.no_pad_namespace.update(set(no_pad_namespace))
        self.no_unk_namespace.update(set(no_unk_namespace))
        self.contain_pad_namespace.update(dict(contain_pad_namespace))
        self.contain_unk_namespace.update(dict(contain_unk_namespace))
        self.min_count.update(dict(min_count))

        for namespace, counter in counters.items():
            self.__namespace_init(namespace)
            for key in counter:
                minc = min_count[namespace] \
                    if min_count and namespace in min_count else 1
                if counter[key] >= minc:
                    self.vocab[namespace][key] = len(self.vocab[namespace])

            logger.info("Vocabulay {} (size: {}) was constructed successfully from counter.".format(
                namespace, len(self.vocab[namespace])))

    def add_tokens_to_namespace(self, tokens, namespace):
        """This function adds tokens to one namespace for extending vocabulary

        Arguments:
            tokens {list} -- token list
            namespace {str} -- namespace name
        """

        if namespace not in self.vocab:
            self.__namespace_init(namespace)
            logger.error('Add Namespace {} into vocabulary.'.format(namespace))

        for token in tokens:
            if token not in self.vocab[namespace]:
                self.vocab[namespace][token] = len(self.vocab[namespace])

    def get_token_index(self, token, namespace):
        """This function gets token index in one namespace of vocabulary

        Arguments:
            token {str} -- token
            namespace {str} -- namespace name

        Raises:
            RuntimeError: namespace not exists

        Returns:
            int -- token index
        """

        if token in self.vocab[namespace]:
            return self.vocab[namespace][token]

        if namespace not in self.no_unk_namespace:
            return self.get_unknown_index(namespace)

        logger.error("Can not find the index of {} from a no unknown token namespace {}.".format(
            token, namespace))
        raise RuntimeError(
            "Can not find the index of {} from a no unknown token namespace {}.".format(
                token, namespace))

    def get_token_from_index(self, index, namespace):
        """This function gets token using index in vocabulary

        Arguments:
            index {int} -- index
            namespace {str} -- namespace name

        Raises:
            RuntimeError: index out of range

        Returns:
            str -- token
        """

        if index < len(self.vocab[namespace]):
            return self.vocab[namespace].inv[index]

        logger.error("The index {} is out of vocabulary {} range.".format(index, namespace))
        raise RuntimeError("The index {} is out of vocabulary {} range.".format(index, namespace))

    def get_vocab_size(self, namespace):
        """This function gets the size of one namespace in vocabulary

        Arguments:
            namespace {str} -- namespace name

        Returns:
            int -- vocabulary size
        """

        return len(self.vocab[namespace])

    def get_all_namespaces(self):
        """This function gets all namespaces

        Returns:
            list -- all namespaces vocabulary contained
        """

        return set(self.vocab)

    def get_padding_index(self, namespace):
        """This function gets padding token index in one namespace of vocabulary

        Arguments:
            namespace {str} -- namespace name

        Raises:
            RuntimeError: no padding

        Returns:
            int -- padding index
        """

        if namespace not in self.vocab:
            raise RuntimeError("Namespace {} doesn't exist.".format(namespace))

        if namespace not in self.no_pad_namespace:
            if namespace not in self.contain_pad_namespace:
                return self.vocab[namespace][Vocabulary.DEFAULT_PAD_TOKEN]
            return self.vocab[namespace][self.contain_pad_namespace[namespace]]

        logger.error("Namespace {} doesn't has paddding token.".format(namespace))
        raise RuntimeError("Namespace {} doesn't has paddding token.".format(namespace))

    def get_unknown_index(self, namespace):
        """This function gets unknown token index in one namespace of vocabulary

        Arguments:
            namespace {str} -- namespace name

        Raises:
            RuntimeError: no unknown

        Returns:
            int -- unknown index
        """

        if namespace not in self.vocab:
            raise RuntimeError("Namespace {} doesn't exist.".format(namespace))

        if namespace not in self.no_unk_namespace:
            if namespace not in self.contain_unk_namespace:
                return self.vocab[namespace][Vocabulary.DEFAULT_UNK_TOKEN]
            return self.vocab[namespace][self.contain_unk_namespace[namespace]]

        logger.error("Namespace {} doesn't has unknown token.".format(namespace))
        raise RuntimeError("Namespace {} doesn't has unknown token.".format(namespace))

    def get_namespace_tokens(self, namesapce):
        """This function returns all tokens in one namespace

        Arguments:
            namesapce {str} -- namespce name

        Returns:
            dict_keys -- all tokens
        """

        return self.vocab[namesapce]

    def save(self, file_path):
        """This function saves vocabulary into file

        Arguments:
            file_path {str} -- file path
        """

        pickle.dump(self, open(file_path, 'wb'))

    @classmethod
    def load(cls, file_path):
        """This function loads vocabulary from file

        Arguments:
            file_path {str} -- file path

        Returns:
            Vocabulary -- vocabulary
        """

        return pickle.load(open(file_path, 'rb'), encoding='utf-8')

    def __namespace_init(self, namespace):
        """This function initializes a namespace,
        adds pad and unk token to one namespace of vacabulary

        Arguments:
            namespace {str} -- namespace
        """

        self.vocab[namespace] = bidict()

        if namespace not in self.no_pad_namespace and namespace not in self.contain_pad_namespace:
            self.vocab[namespace][Vocabulary.DEFAULT_PAD_TOKEN] = len(self.vocab[namespace])

        if namespace not in self.no_unk_namespace and namespace not in self.contain_unk_namespace:
            self.vocab[namespace][Vocabulary.DEFAULT_UNK_TOKEN] = len(self.vocab[namespace])