import logging
from typing import Dict, List
import warnings
from collections import Counter, OrderedDict
import torch
import torch.nn as nn
from torchtext._torchtext import (
    Vocab as VocabPybind,
    _load_vocab_from_file,
    _build_vocab_from_text_file
)

__all__ = [
    'build_vocab_from_text_file',
    'load_vocab_from_file',
    'vocab',
    'Vocab',
]
logger = logging.getLogger(__name__)


def build_vocab_from_text_file(file_object, jited_tokenizer, min_freq=1, unk_token='<unk>', num_cpus=4):
    r"""Create a `Vocab` object from a raw text file.

    The `file_object` can contain any raw text. This function applies a generic JITed tokenizer in
    parallel to the text. Note that the vocab will be created in the order that the tokens first appear
    in the file (and not by the frequency of tokens).

    Args:
        file_object (FileObject): a file object to read data from.
        jited_tokenizer (ScriptModule): a tokenizer that has been JITed using `torch.jit.script`
        min_freq: The minimum frequency needed to include a token in the vocabulary.
            Values less than 1 will be set to 1. Default: 1.
        unk_token: The default unknown token to use. Default: '<unk>'.
        num_cpus (int): the number of cpus to use when loading the vectors from file. Default: 4.

    Returns:
        torchtext.experimental.vocab.Vocab: a `Vocab` object.

    Examples:
        >>> from torchtext.experimental.vocab import build_vocab_from_text_file
        >>> from torchtext.experimental.transforms import basic_english_normalize
        >>> f = open('vocab.txt', 'r')
        >>>     tokenizer = basic_english_normalize()
        >>> tokenizer = basic_english_normalize()
        >>> jit_tokenizer = torch.jit.script(tokenizer)
        >>> v = build_vocab_from_text_file(f, jit_tokenizer)
    """
    vocab_obj = _build_vocab_from_text_file(file_object.name, unk_token, min_freq, num_cpus, jited_tokenizer)
    return Vocab(vocab_obj)


def load_vocab_from_file(file_object, min_freq=1, unk_token='<unk>', num_cpus=4):
    r"""Create a `Vocab` object from a text file.
    The `file_object` should contain tokens separated by new lines. Note that the vocab
    will be created in the order that the tokens first appear in the file (and not by the frequency of tokens).
    Format for txt file:

        token1
        token2
        ...
        token_n

    Args:
        file_object (FileObject): a file like object to read data from.
        min_freq: The minimum frequency needed to include a token in the vocabulary.
            Values less than 1 will be set to 1. Default: 1.
        unk_token: The default unknown token to use. Default: '<unk>'.
        num_cpus (int): the number of cpus to use when loading the vectors from file. Default: 4.

    Returns:
        torchtext.experimental.vocab.Vocab: a `Vocab` object.

    Examples:
        >>> from torchtext.experimental.vocab import load_vocab_from_file
        >>> f = open('vocab.txt', 'r')
        >>> v = load_vocab_from_file(f)
    """

    vocab_obj = _load_vocab_from_file(file_object.name, unk_token, min_freq, num_cpus)
    return Vocab(vocab_obj)


def build_vocab_from_iterator(iterator, min_freq=1, unk_token='<unk>'):
    """
    Build a Vocab from an iterator.

    Args:
        iterator: Iterator used to build Vocab. Must yield list or iterator of tokens.
        min_freq: The minimum frequency needed to include a token in the vocabulary.
            Values less than 1 will be set to 1. Default: 1.
        unk_token: The default unknown token to use. Default: '<unk>'.
    """

    counter = Counter()
    for tokens in iterator:
        counter.update(tokens)
    sorted_by_freq_tuples = sorted(counter.items(), key=lambda x: x[1], reverse=True)
    ordered_dict = OrderedDict(sorted_by_freq_tuples)
    word_vocab = vocab(ordered_dict, min_freq=min_freq, unk_token=unk_token)
    return word_vocab


def vocab(ordered_dict, min_freq=1, unk_token='<unk>'):
    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.
    Additionally, the if the `unk_token` isn't found inside of the `ordered_dict`, it will be added to the end of the vocab.

    Args:
        ordered_dict (collections.OrderedDict): object holding the frequencies of each token found in the data.
        min_freq: The minimum frequency needed to include a token in the vocabulary.
            Values less than 1 will be set to 1. Default: 1.
        unk_token: The default unknown token to use. Default: '<unk>'.

    Raises:
        ValueError: if a default `unk_token` isn't provided.

    Examples:
        >>> from torchtext.experimental.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)
        >>> tokens = ['e', 'd', 'c', 'b', 'a']
        >>> v2 = vocab(OrderedDict([(token, 1) for token in tokens]))
    """

    if not unk_token:
        raise ValueError("A default unk token wasn't provided.")

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

    if unk_token not in tokens:
        tokens.insert(0, unk_token)
        warnings.warn("The `unk_token` '{}' wasn't found in the `ordered_dict`. Adding the `unk_token` "
                      "to the beginning of the Vocab.".format(unk_token), RuntimeWarning)
    return Vocab(VocabPybind(tokens, unk_token))


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):
        super(Vocab, self).__init__()
        self.vocab = vocab

    @property
    def is_jitable(self):
        return not isinstance(self.vocab, VocabPybind)

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

        Args:
            tokens (List[str]): a list of tokens used to lookup their corresponding `indices`.

        Returns:
            indices (List[int]): the 'indices` associated with a list of tokens`.
        """
        return self.vocab.lookup_indices(tokens)

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

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

        Returns:
            index (int): the index corresponding to the associated token.
        """
        return self.vocab[token]

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

        Raises:
            RuntimeError: if `index` not between [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 (str): the token used to lookup the corresponding index.
        """
        self.vocab.append_token(token)

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

        Returns:
            token (str): the token used to lookup the corresponding index.

        Raises:
            RuntimeError: if `index` not between [0, itos.size()].
        """
        return self.vocab.lookup_token(index)

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

        Returns:
            tokens (List[str]): the `tokens` associated with `indices`.

        Raises:
            RuntimeError: if an index within `indices` is not between [0, itos.size()].
        """
        return self.vocab.lookup_tokens(indices)

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

        Returns:
            indices (List[int]): the 'indices` associated with `tokens`.
        """
        return self.vocab.lookup_indices(tokens)

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

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

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