# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import json
import torch
from collections import OrderedDict 
from nltk.tokenize import TweetTokenizer
from pytorch_transformers.tokenization_bert import BertTokenizer #, BertModel
from transformers import BertModel
from src.tools.utils import load_glove, load_fasttext

PRETRAINED_BERT_NAME = 'bert-base-uncased'

def load_bert_model():
    return BertModel.from_pretrained(PRETRAINED_BERT_NAME)

def load_glove_embed(num_wrds, data_path, pad_id, freeze=False):
    _, _, vectors = load_glove(data_path)
    weights = torch.rand(num_wrds, vectors.shape[-1])
    weights[:vectors.shape[0]] = torch.tensor(vectors).float()
    return torch.nn.Embedding.from_pretrained(
        weights, 
        freeze=freeze,
        padding_idx=pad_id)

class Glove_Tokenizer(object):
    """Tokenizer for glove or fasttext"""
    def __init__(self, data_path, mode='glove'):
        if mode == 'glove':
            vocab, ivocab = load_glove(data_path, only_vocab=True)
        # elif mode == 'fasttext':
        #     # TODO: load_fasttext
        #     vocab, ivocab = load_fasttext(data_path, only_vocab=True)
        else:
            raise NotImplementedError

        # with open(vocab_path, 'r') as f:
        #     self.word2i = json.load(f)['word2i']
        self._tokenizer = TweetTokenizer(preserve_case=False)
        _predefined = ['<pad>', '<sos>', '<eod>']
        for tok in _predefined:
            if tok not in vocab:
                vocab[tok] = len(vocab)
                ivocab[len(ivocab)] = tok

        self.vocab = vocab
        self.ivocab = ivocab
        
        # Retrieve key values
        self.num_words = len(self.vocab)
        self.sos_id = self.vocab["<sos>"]
        self.eoq_id = self.vocab["?"]
        self.eod_id = self.vocab["<eod>"]
        self.pad_id = self.vocab["<pad>"]
        self.yes_id = self.vocab["yes"]
        self.no_id = self.vocab["no"]
        self.na_id = self.vocab["n/a"]
        self.unk = "<unk>"
#         self.answers = [self.yes_id, self.no_id, self.non_applicable_id]

    def encode(self, words):
        tokens = []
        for token in self._tokenizer.tokenize(words):
            token = self.unk if token not in self.vocab else token
            tokens.append(self.vocab[token])
        return tokens

    def decode(self, tokens, ignore_pad=False):
        if ignore_pad:
            return ' '.join([self.ivocab[tok] for tok in tokens if tok != self.pad_id])
        else:
            return ' '.join([self.ivocab[tok] for tok in tokens])
        
    def _convert_token_to_id(self, token):
        token = '<unk>' if token not in self.vocab else token
        return self.vocab[token]

    def __len__(self):
        return self.num_words



class BERT_Tokenizer(BertTokenizer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
        # Vocab 1 ~ 900 are not used
        # Use key to replace value
        self.replace_tokens = OrderedDict([
            # ('<yes>', '[unused0]'), 
            # ( '<no>', '[unused1]'),
            ( '<na>', '[unused0]'),
            # ('<eoq>', '[unused1]'),
            ('<eod>', '[unused1]'),
        #     ('<sos>', '[unused6]'),
        ])
        for k, v in self.replace_tokens.items():
            # self.yes_id = 1
            # self.no_id = 2
            # ...
            setattr(self, k.replace('<', '').replace('>', '')+'_id', self.vocab[v])

    @property
    def pad_id(self):
        return self._convert_token_to_id(self.pad_token)
    @property
    def yes_id(self):
        return self._convert_token_to_id('yes')
    @property
    def no_id(self):
        return self._convert_token_to_id('no')
    @property
    def sos_id(self):
        return self._convert_token_to_id('[CLS]')
    @property
    def cls_id(self):
        return self._convert_token_to_id('[CLS]')
    @property
    def eoq_id(self):
        return self._convert_token_to_id('?')

    def decode(self, tokens, ignore_pad=False):
        if ignore_pad:
            tokens = [tok for tok in tokens if tok != self.pad_id]
        dec = super().decode(tokens)
        for k, v in self.replace_tokens.items():
            dec = dec.replace(v, k)
        return dec
        
    @classmethod
    def initialize(cls):
        print('bert initialize')
        return cls.from_pretrained(PRETRAINED_BERT_NAME)


class GW_Tokenizer(object):
    """Tokenizer from original Guesswhat paper (github)"""
    def __init__(self, vocab_path):
        with open(vocab_path, 'r') as f:

            self.word2i = json.load(f)['word2i']
            print('load word2i')

        self._tokenizer = TweetTokenizer(preserve_case=False)
        if "<stop_dialogue>" not in self.word2i:
            self.word2i["<stop_dialogue>"] = len(self.word2i)

        self.i2word = {}
        for (k, v) in self.word2i.items():
            self.i2word[v] = k

        # Retrieve key values
        self.num_words = len(self.word2i)
        self.sos_id = self.word2i["<start>"]
        self.eoq_id = self.word2i["?"]
        self.eod_id = self.word2i["<stop_dialogue>"]
        self.pad_id = self.word2i["<padding>"]
        self.yes_id = self.word2i["<yes>"]
        self.no_id = self.word2i["<no>"]
        self.na_id = self.word2i["<n/a>"]


    def encode(self, words):
        tokens = []
        for token in self._tokenizer.tokenize(words):
            token = '<unk>' if token not in self.word2i else token
            tokens.append(self.word2i[token])
        return tokens

    def decode(self, tokens, ignore_pad=False):
        if ignore_pad:
            return ' '.join([self.i2word[tok] for tok in tokens if tok != self.pad_id])
        else:
            return ' '.join([self.i2word[tok] for tok in tokens])
        
    def _convert_token_to_id(self, token):
        token = '<unk>' if token not in self.word2i else token
        return self.word2i[token]

    def __len__(self):
        return self.num_words
