# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import logging
from collections import defaultdict
from functools import partial

import seq2seq.data.config as config


class Tokenizer:
    """
    Tokenizer class.
    """
    def __init__(self, vocab_fname=None, separator='@@'):
        """
        Constructor for the Tokenizer class.

        :param vocab_fname: path to the file with vocabulary
        :param separator: tokenization separator
        """
        if vocab_fname:
            self.separator = separator

            logging.info(f'Building vocabulary from {vocab_fname}')
            vocab = [config.PAD_TOKEN, config.UNK_TOKEN,
                    config.BOS_TOKEN, config.EOS_TOKEN]

            with open(vocab_fname) as vfile:
                for line in vfile:
                    vocab.append(line.strip())

            # Padding vocabulary to 8
            for i in range(8 - len(vocab) % 8):
                tok = 'madeupword{:04d}'.format(i)
                vocab.append(tok)

            self.vocab_size = len(vocab)
            logging.info(f'Size of vocabulary: {self.vocab_size}')

            self.tok2idx = defaultdict(partial(int, config.UNK))
            for idx, token in enumerate(vocab):
                self.tok2idx[token] = idx

            self.idx2tok = {}
            for key, value in self.tok2idx.items():
                self.idx2tok[value] = key

    def get_state(self):
        logging.info(f'Saving state of the tokenizer')
        state = {
            'separator': self.separator,
            'vocab_size': self.vocab_size,
            'tok2idx': self.tok2idx,
            'idx2tok': self.idx2tok,
        }
        return state

    def set_state(self, state):
        logging.info(f'Restoring state of the tokenizer')
        self.separator = state['separator']
        self.vocab_size = state['vocab_size']
        self.tok2idx = state['tok2idx']
        self.idx2tok = state['idx2tok']

    def segment(self, line):
        """
        Tokenizes single sentence and adds special BOS and EOS tokens.

        :param line: sentence

        returns: list representing tokenized sentence
        """
        line = line.strip().split()
        entry = [self.tok2idx[i] for i in line]
        entry = [config.BOS] + entry + [config.EOS]
        return entry

    def detokenize(self, inputs, delim=' '):
        """
        Detokenizes single sentence and removes token separator characters.

        :param inputs: sequence of tokens
        :param delim: tokenization delimiter

        returns: string representing detokenized sentence
        """
        detok = delim.join([self.idx2tok[idx] for idx in inputs])
        detok = detok.replace(self.separator + ' ', '')
        detok = detok.replace(self.separator, '')

        detok = detok.replace(config.BOS_TOKEN, '')
        detok = detok.replace(config.EOS_TOKEN, '')
        detok = detok.replace(config.PAD_TOKEN, '')
        detok = detok.strip()
        return detok
