from typing import Dict, List, Optional

import torch
import torch.nn as nn

def _log_class_usage(klass):
    identifier = "torchtext"
    if klass and hasattr(klass, "__name__"):
        identifier += f".{klass.__name__}"
    torch._C._log_api_usage_once(identifier)


class Vocab(nn.Module):
    __jit_unused_properties__ = ["is_jitable"]
    r"""Creates a vocab object which maps tokens to indices.

    Args:
        vocab (torch.classes.torchtext.Vocab or torchtext._torchtext.Vocab): a cpp vocab object.
    """

    def __init__(self, vocab) -> None:
        super(Vocab, self).__init__()
        self.vocab = vocab
        _log_class_usage(__class__)

    @property
    def is_jitable(self):
        return isinstance(self.vocab, torch._C.ScriptObject)

    @torch.jit.export
    def forward(self, tokens: List[str]) -> List[int]:
        r"""Calls the `lookup_indices` method

        Args:
            tokens: a list of tokens used to lookup their corresponding `indices`.

        Returns:
            The indices associated with a list of `tokens`.
        """
        return self.vocab.lookup_indices(tokens)

    @torch.jit.export
    def __len__(self) -> int:
        r"""
        Returns:
            The length of the vocab.
        """
        return len(self.vocab)

    @torch.jit.export
    def __contains__(self, token: str) -> bool:
        r"""
        Args:
            token: The token for which to check the membership.

        Returns:
            Whether the token is member of vocab or not.
        """
        return self.vocab.__contains__(token)

    @torch.jit.export
    def __getitem__(self, token: str) -> int:
        r"""
        Args:
            token: The token used to lookup the corresponding index.

        Returns:
            The index corresponding to the associated token.
        """
        return self.vocab[token]

    @torch.jit.export
    def set_default_index(self, index: Optional[int]) -> None:
        r"""
        Args:
            index: Value of default index. This index will be returned when OOV token is queried.
        """
        self.vocab.set_default_index(index)

    @torch.jit.export
    def get_default_index(self) -> Optional[int]:
        r"""
        Returns:
            Value of default index if it is set.
        """
        return self.vocab.get_default_index()

    @torch.jit.export
    def insert_token(self, token: str, index: int) -> None:
        r"""
        Args:
            token: The token used to lookup the corresponding index.
            index: The index corresponding to the associated token.
        Raises:
            RuntimeError: If `index` is not in range [0, Vocab.size()] or if `token` already exists in the vocab.
        """
        self.vocab.insert_token(token, index)

    @torch.jit.export
    def append_token(self, token: str) -> None:
        r"""
        Args:
            token: The token used to lookup the corresponding index.

        Raises:
            RuntimeError: If `token` already exists in the vocab
        """
        self.vocab.append_token(token)

    @torch.jit.export
    def lookup_token(self, index: int) -> str:
        r"""
        Args:
            index: The index corresponding to the associated token.

        Returns:
            token: The token used to lookup the corresponding index.

        Raises:
            RuntimeError: If `index` not in range [0, itos.size()).
        """
        return self.vocab.lookup_token(index)

    @torch.jit.export
    def lookup_tokens(self, indices: List[int]) -> List[str]:
        r"""
        Args:
            indices: The `indices` used to lookup their corresponding`tokens`.

        Returns:
            The `tokens` associated with `indices`.

        Raises:
            RuntimeError: If an index within `indices` is not int range [0, itos.size()).
        """
        return self.vocab.lookup_tokens(indices)

    @torch.jit.export
    def lookup_indices(self, tokens: List[str]) -> List[int]:
        r"""
        Args:
            tokens: the tokens used to lookup their corresponding `indices`.

        Returns:
            The 'indices` associated with `tokens`.
        """
        return self.vocab.lookup_indices(tokens)

    @torch.jit.export
    def get_stoi(self) -> Dict[str, int]:
        r"""
        Returns:
            Dictionary mapping tokens to indices.
        """
        return self.vocab.get_stoi()

    @torch.jit.export
    def get_itos(self) -> List[str]:
        r"""
        Returns:
            List mapping indices to tokens.
        """
        return self.vocab.get_itos()

    def __prepare_scriptable__(self):
        r"""Return a JITable Vocab."""
        if not self.is_jitable:
            cpp_vocab = torch.classes.torchtext.Vocab(self.vocab.itos_, self.vocab.default_index_)
            return Vocab(cpp_vocab)
        return self


# from torchtext._torchtext import Vocab as VocabPybind

# from .vocab import Vocab


# def vocab(
#     ordered_dict: Dict, min_freq: int = 1, specials: Optional[List[str]] = None, special_first: bool = True
# ) -> Vocab:
#     r"""Factory method for creating a vocab object which maps tokens to indices.

#     Note that the ordering in which key value pairs were inserted in the `ordered_dict` will be respected when building the vocab.
#     Therefore if sorting by token frequency is important to the user, the `ordered_dict` should be created in a way to reflect this.

#     Args:
#         ordered_dict: Ordered Dictionary mapping tokens to their corresponding occurance frequencies.
#         min_freq: The minimum frequency needed to include a token in the vocabulary.
#         specials: Special symbols to add. The order of supplied tokens will be preserved.
#         special_first: Indicates whether to insert symbols at the beginning or at the end.

#     Returns:
#         torchtext.vocab.Vocab: A `Vocab` object

#     Examples:
#         >>> from torchtext.vocab import vocab
#         >>> from collections import Counter, OrderedDict
#         >>> counter = Counter(["a", "a", "b", "b", "b"])
#         >>> sorted_by_freq_tuples = sorted(counter.items(), key=lambda x: x[1], reverse=True)
#         >>> ordered_dict = OrderedDict(sorted_by_freq_tuples)
#         >>> v1 = vocab(ordered_dict)
#         >>> print(v1['a']) #prints 1
#         >>> print(v1['out of vocab']) #raise RuntimeError since default index is not set
#         >>> tokens = ['e', 'd', 'c', 'b', 'a']
#         >>> #adding <unk> token and default index
#         >>> unk_token = '<unk>'
#         >>> default_index = -1
#         >>> v2 = vocab(OrderedDict([(token, 1) for token in tokens]), specials=[unk_token])
#         >>> v2.set_default_index(default_index)
#         >>> print(v2['<unk>']) #prints 0
#         >>> print(v2['out of vocab']) #prints -1
#         >>> #make default index same as index of unk_token
#         >>> v2.set_default_index(v2[unk_token])
#         >>> v2['out of vocab'] is v2[unk_token] #prints True
#     """
#     specials = specials or []
#     for token in specials:
#         ordered_dict.pop(token, None)

#     tokens = []
#     # Save room for special tokens
#     for token, freq in ordered_dict.items():
#         if freq >= min_freq:
#             tokens.append(token)

#     if special_first:
#         tokens[0:0] = specials
#     else:
#         tokens.extend(specials)

#     return Vocab(VocabPybind(tokens, None))