# conda activate py36gputorch100 && \
# cd /mnt/1T-5e7/mycodehtml/NLP/joosthub/chapters/chapter_5/5_2_CBOW && \
# rm e.l && python 5_2_Continuous_Bag_of_Words_CBOW.py \
# 2>&1 | tee -a e.l && code e.l

# ================================================================================
# Learning Embeddings with Continuous Bag of Words (CBOW)

# ================================================================================
import os
from argparse import Namespace
from collections import Counter
import json
import re
import string

import numpy as np
import pandas as pd
pd.set_option('display.max_colwidth', -1)
pd.set_option('display.max_columns',None)
import torch 
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm_notebook

# ================================================================================
class Vocabulary(object):
    """Class to process text and extract vocabulary for mapping"""

    def __init__(self, token_to_idx=None, mask_token="<MASK>", add_unk=True, unk_token="<UNK>"):
        """
        Args:
            token_to_idx (dict): a pre-existing map of tokens to indices
            mask_token (str): the MASK token to add into the Vocabulary; 
              indicates a position that will not be used in updating the model's parameters
            add_unk (bool): a flag that indicates whether to add the UNK token
            unk_token (str): the UNK token to add into the Vocabulary
        """

        # ================================================================================
        # If token_to_idx is None, create dict
        if token_to_idx is None:
            token_to_idx = {}

        # ================================================================================
        # {token:idx}
        self._token_to_idx = token_to_idx

        # {idx:token}
        self._idx_to_token = {idx: token 
                              for token, idx in self._token_to_idx.items()}
        
        self._add_unk = add_unk
        self._unk_token = unk_token
        self._mask_token = mask_token
        
        self.mask_index = self.add_token(self._mask_token)
        self.unk_index = -1
        if add_unk:
            self.unk_index = self.add_token(unk_token) 
        
    def to_serializable(self):
        """ returns a dictionary that can be serialized """
        return {'token_to_idx': self._token_to_idx, 
                'add_unk': self._add_unk, 
                'unk_token': self._unk_token, 
                'mask_token': self._mask_token}

    @classmethod
    def from_serializable(cls, contents):
        """ instantiates the Vocabulary from a serialized dictionary """
        return cls(**contents)

    def add_token(self, token):
        """Update mapping dicts based on the token.

        Args:
            token (str): the item to add into the Vocabulary
        Returns:
            index (int): the integer corresponding to the token
        """
        if token in self._token_to_idx:
            index = self._token_to_idx[token]
        else:
            index = len(self._token_to_idx)
            self._token_to_idx[token] = index
            self._idx_to_token[index] = token
        return index
            
    def add_many(self, tokens):
        """Add a list of tokens into the Vocabulary
        
        Args:
            tokens (list): a list of string tokens
        Returns:
            indices (list): a list of indices corresponding to the tokens
        """
        return [self.add_token(token) for token in tokens]

    def lookup_token(self, token):
        """Retrieve the index associated with the token 
          or the UNK index if token isn't present.
        
        Args:
            token (str): the token to look up 
        Returns:
            index (int): the index corresponding to the token
        Notes:
            `unk_index` needs to be >=0 (having been added into the Vocabulary) 
              for the UNK functionality 
        """
        if self.unk_index >= 0:
            return self._token_to_idx.get(token, self.unk_index)
        else:
            return self._token_to_idx[token]

    def lookup_index(self, index):
        """Return the token associated with the index
        
        Args: 
            index (int): the index to look up
        Returns:
            token (str): the token corresponding to the index
        Raises:
            KeyError: if the index is not in the Vocabulary
        """
        if index not in self._idx_to_token:
            raise KeyError("the index (%d) is not in the Vocabulary" % index)
        return self._idx_to_token[index]

    def __str__(self):
        return "<Vocabulary(size=%d)>" % len(self)

    def __len__(self):
        return len(self._token_to_idx)



class CBOWVectorizer(object):
    """ The Vectorizer which coordinates the Vocabularies and puts them to use"""    
    def __init__(self, cbow_vocab):
        """
        Args:
            cbow_vocab (Vocabulary): maps words to integers
        """
        self.cbow_vocab = cbow_vocab

    def vectorize(self, context, vector_length=-1):
        """
        Args:
            context (str): the string of words separated by a space
            vector_length (int): an argument for forcing the length of index vector
        """

        indices = [self.cbow_vocab.lookup_token(token) for token in context.split(' ')]
        if vector_length < 0:
            vector_length = len(indices)

        out_vector = np.zeros(vector_length, dtype=np.int64)
        out_vector[:len(indices)] = indices
        out_vector[len(indices):] = self.cbow_vocab.mask_index

        return out_vector
    
    @classmethod
    def from_dataframe(cls, cbow_df):
        """Instantiate the vectorizer from the dataset dataframe
        
        Args:
            cbow_df (pandas.DataFrame): the target dataset
        Returns:
            an instance of the CBOWVectorizer
        """
        cbow_vocab = Vocabulary()

        # print("cbow_df",cbow_df)
        #                                                  context  ...    split
        # 0                                               , or the  ...    train
        # 1                             frankenstein or the modern  ...    train

        # ================================================================================
        for index, row in cbow_df.iterrows():
            # print("index",index)
            # 0
            
            # print("row",row)
            # context        , or the
            # target     frankenstein
            # split             train

            # ================================================================================
            for token in row.context.split(' '):
                # print("token",token)
                # ,
                # or
                # the
                cbow_vocab.add_token(token)

            # ================================================================================
            cbow_vocab.add_token(row.target)

        return cls(cbow_vocab)

    @classmethod
    def from_serializable(cls, contents):
        cbow_vocab =             Vocabulary.from_serializable(contents['cbow_vocab'])
        return cls(cbow_vocab=cbow_vocab)

    def to_serializable(self):
        return {'cbow_vocab': self.cbow_vocab.to_serializable()}

# ================================================================================
# The Dataset

class CBOWDataset(Dataset):
    def __init__(self, cbow_df, vectorizer):
        """
        Args:
            cbow_df (pandas.DataFrame): the dataset
            vectorizer (CBOWVectorizer): vectorizer instatiated from dataset
        """
        self.cbow_df = cbow_df
        self._vectorizer = vectorizer
        
        measure_len = lambda context: len(context.split(" "))
        self._max_seq_length = max(map(measure_len, cbow_df.context))
        
        self.train_df = self.cbow_df[self.cbow_df.split=='train']
        self.train_size = len(self.train_df)

        self.val_df = self.cbow_df[self.cbow_df.split=='val']
        self.validation_size = len(self.val_df)

        self.test_df = self.cbow_df[self.cbow_df.split=='test']
        self.test_size = len(self.test_df)

        self._lookup_dict = {'train': (self.train_df, self.train_size),
                             'val': (self.val_df, self.validation_size),
                             'test': (self.test_df, self.test_size)}

        self.set_split('train')

    @classmethod
    def load_dataset_and_make_vectorizer(cls, cbow_csv):
        """Load dataset and make a new vectorizer from scratch
        
        Args:
            cbow_csv (str): location of the dataset
        Returns:
            an instance of CBOWDataset
        """
        cbow_df = pd.read_csv(cbow_csv)
        # print("cbow_df",cbow_df)
        #                                                  context  ...    split
        # 0                                               , or the  ...    train
        # 1                             frankenstein or the modern  ...    train
        # 2                   frankenstein , the modern prometheus  ...    train

        # print("cbow_df",cbow_df.shape)
        # (90698, 3)

        train_cbow_df = cbow_df[cbow_df.split=='train']
        # print("train_cbow_df",train_cbow_df)
        #                                                   context          target      split  
        # 0      , or the                                            frankenstein        train   
        # 1      frankenstein or the modern                          ,                   train   
        # 2      frankenstein , the modern prometheus                or                  train   
        # 3      frankenstein , or modern prometheus by              the                 train   
        # 4      , or the prometheus by mary                         modern              train   
        # 5      or the modern by mary wollstonecraft                prometheus          train   

        # print("train_cbow_df",train_cbow_df.shape)
        # (63489, 3)

        # ================================================================================
        dataframe=CBOWVectorizer.from_dataframe(train_cbow_df)
        
        cls_obj=cls(cbow_df, dataframe)

        return cls_obj

    @classmethod
    def load_dataset_and_load_vectorizer(cls, cbow_csv, vectorizer_filepath):
        """Load dataset and the corresponding vectorizer. 
        Used in the case in the vectorizer has been cached for re-use
        
        Args:
            cbow_csv (str): location of the dataset
            vectorizer_filepath (str): location of the saved vectorizer
        Returns:
            an instance of CBOWDataset
        """
        cbow_df = pd.read_csv(cbow_csv)
        vectorizer = cls.load_vectorizer_only(vectorizer_filepath)
        return cls(cbow_df, vectorizer)

    @staticmethod
    def load_vectorizer_only(vectorizer_filepath):
        """a static method for loading the vectorizer from file
        
        Args:
            vectorizer_filepath (str): the location of the serialized vectorizer
        Returns:
            an instance of CBOWVectorizer
        """
        with open(vectorizer_filepath) as fp:
            return CBOWVectorizer.from_serializable(json.load(fp))

    def save_vectorizer(self, vectorizer_filepath):
        """saves the vectorizer to disk using json
        
        Args:
            vectorizer_filepath (str): the location to save the vectorizer
        """
        with open(vectorizer_filepath, "w") as fp:
            json.dump(self._vectorizer.to_serializable(), fp)

    def get_vectorizer(self):
        """ returns the vectorizer """
        return self._vectorizer
        
    def set_split(self, split="train"):
        """ selects the splits in the dataset using a column in the dataframe """
        self._target_split = split
        self._target_df, self._target_size = self._lookup_dict[split]

    def __len__(self):
        return self._target_size

    def __getitem__(self, index):
        """the primary entry point method for PyTorch datasets
        
        Args:
            index (int): the index to the data point 
        Returns:
            a dictionary holding the data point's features (x_data) and label (y_target)
        """
        row = self._target_df.iloc[index]

        context_vector =             self._vectorizer.vectorize(row.context, self._max_seq_length)
        target_index = self._vectorizer.cbow_vocab.lookup_token(row.target)

        return {'x_data': context_vector,
                'y_target': target_index}

    def get_num_batches(self, batch_size):
        """Given a batch size, return the number of batches in the dataset
        
        Args:
            batch_size (int)
        Returns:
            number of batches in the dataset
        """
        return len(self) // batch_size
    
def generate_batches(dataset, batch_size, shuffle=True,
                     drop_last=True, device="cpu"): 
    """
    A generator function which wraps the PyTorch DataLoader. It will 
      ensure each tensor is on the write device location.
    """
    dataloader = DataLoader(dataset=dataset, batch_size=batch_size,
                            shuffle=shuffle, drop_last=drop_last)

    for data_dict in dataloader:
        out_data_dict = {}
        for name, tensor in data_dict.items():
            out_data_dict[name] = data_dict[name].to(device)
        yield out_data_dict


# The Model: CBOW


class CBOWClassifier(nn.Module): # Simplified cbow Model
    def __init__(self, vocabulary_size, embedding_size, padding_idx=0):
        """
        Args:
            vocabulary_size (int): number of vocabulary items, controls the
                number of embeddings and prediction vector size
            embedding_size (int): size of the embeddings
            padding_idx (int): default 0; Embedding will not use this index
        """
        super(CBOWClassifier, self).__init__()
        
        self.embedding =  nn.Embedding(num_embeddings=vocabulary_size, 
                                       embedding_dim=embedding_size,
                                       padding_idx=padding_idx)
        self.fc1 = nn.Linear(in_features=embedding_size,
                             out_features=vocabulary_size)

    def forward(self, x_in, apply_softmax=False):
        """The forward pass of the classifier
        
        Args:
            x_in (torch.Tensor): an input data tensor. 
                x_in.shape should be (batch, input_dim)
            apply_softmax (bool): a flag for the softmax activation
                should be false if used with the Cross Entropy losses
        Returns:
            the resulting tensor. tensor.shape should be (batch, output_dim)
        """
        x_embedded_sum = F.dropout(self.embedding(x_in).sum(dim=1), 0.3)
        y_out = self.fc1(x_embedded_sum)
        
        if apply_softmax:
            y_out = F.softmax(y_out, dim=1)
            
        return y_out

# ================================================================================
# Helper functions

def make_train_state(args):
    return {'stop_early': False,
            'early_stopping_step': 0,
            'early_stopping_best_val': 1e8,
            'learning_rate': args.learning_rate,
            'epoch_index': 0,
            'train_loss': [],
            'train_acc': [],
            'val_loss': [],
            'val_acc': [],
            'test_loss': -1,
            'test_acc': -1,
            'model_filename': args.model_state_file}

def update_train_state(args, model, train_state):
    """Handle the training state updates.

    Components:
     - Early Stopping: Prevent overfitting.
     - Model Checkpoint: Model is saved if the model is better

    :param args: main arguments
    :param model: model to train
    :param train_state: a dictionary representing the training state values
    :returns:
        a new train_state
    """

    # Save one model at least
    if train_state['epoch_index'] == 0:
        torch.save(model.state_dict(), train_state['model_filename'])
        train_state['stop_early'] = False

    # Save model if performance improved
    elif train_state['epoch_index'] >= 1:
        loss_tm1, loss_t = train_state['val_loss'][-2:]

        # If loss worsened
        if loss_t >= train_state['early_stopping_best_val']:
            # Update step
            train_state['early_stopping_step'] += 1
        # Loss decreased
        else:
            # Save the best model
            if loss_t < train_state['early_stopping_best_val']:
                torch.save(model.state_dict(), train_state['model_filename'])

            # Reset early stopping step
            train_state['early_stopping_step'] = 0

        # Stop early ?
        train_state['stop_early'] =             train_state['early_stopping_step'] >= args.early_stopping_criteria

    return train_state

def compute_accuracy(y_pred, y_target):
    _, y_pred_indices = y_pred.max(dim=1)
    n_correct = torch.eq(y_pred_indices, y_target).sum().item()
    return n_correct / len(y_pred_indices) * 100

# ================================================================================
# general utilities

def set_seed_everywhere(seed, cuda):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if cuda:
        torch.cuda.manual_seed_all(seed)

def handle_dirs(dirpath):
    if not os.path.exists(dirpath):
        os.makedirs(dirpath)

# ================================================================================
# Settings and some preprocessing

args = Namespace(
    # Data and Path information
    cbow_csv="data/books/frankenstein_with_splits.csv",
    vectorizer_file="vectorizer.json",
    model_state_file="model.pth",
    save_dir="model_storage/ch5/cbow",
    # Model hyper parameters
    embedding_size=50,
    # Training hyper parameters
    seed=1337,
    # num_epochs=100,
    num_epochs=1,
    learning_rate=0.0001,
    batch_size=32,
    # batch_size=256,
    early_stopping_criteria=5,
    # Runtime options
    cuda=True,
    catch_keyboard_interrupt=True,
    reload_from_files=False,
    expand_filepaths_to_save_dir=True
)

# ================================================================================
if args.expand_filepaths_to_save_dir:
    args.vectorizer_file = os.path.join(args.save_dir,
                                        args.vectorizer_file)

    args.model_state_file = os.path.join(args.save_dir,
                                         args.model_state_file)
    
    print("Expanded filepaths: ")
    print("\t{}".format(args.vectorizer_file))
    print("\t{}".format(args.model_state_file))
    
# ================================================================================
# Check CUDA
if not torch.cuda.is_available():
    args.cuda = False

args.device = torch.device("cuda" if args.cuda else "cpu")
    
print("Using CUDA: {}".format(args.cuda))

# ================================================================================
# Set seed for reproducibility
set_seed_everywhere(args.seed, args.cuda)

# ================================================================================
# handle dirs
handle_dirs(args.save_dir)

# ================================================================================
# Initializations

# print("args.reload_from_files",args.reload_from_files)
if args.reload_from_files:
    print("Loading dataset and loading vectorizer")
    # print("args.cbow_csv",args.cbow_csv)
    # print("args.vectorizer_file",args.vectorizer_file)
    dataset = CBOWDataset.load_dataset_and_load_vectorizer(args.cbow_csv,
                                                           args.vectorizer_file)
else:
    print("Loading dataset and creating vectorizer")
    # print("args.cbow_csv",args.cbow_csv)
    # data/books/frankenstein_with_splits.csv
    dataset = CBOWDataset.load_dataset_and_make_vectorizer(args.cbow_csv)

    dataset.save_vectorizer(args.vectorizer_file)

# ================================================================================    
vectorizer = dataset.get_vectorizer()

# print("len(vectorizer.cbow_vocab)",len(vectorizer.cbow_vocab))
# 6138
# print("args.embedding_size",args.embedding_size)
# 50

classifier = CBOWClassifier(vocabulary_size=len(vectorizer.cbow_vocab), 
                            embedding_size=args.embedding_size)

# ================================================================================
# Training loop

classifier = classifier.to(args.device)
    
loss_func = nn.CrossEntropyLoss()

optimizer = optim.Adam(classifier.parameters(), lr=args.learning_rate)

scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer,
                                                 mode='min', factor=0.5,
                                                 patience=1)
train_state = make_train_state(args)

# ================================================================================
epoch_bar = tqdm_notebook(desc='training routine', 
                          total=args.num_epochs,
                          position=0)

# ================================================================================
dataset.set_split('train')

train_bar = tqdm_notebook(desc='split=train',
                          total=dataset.get_num_batches(args.batch_size), 
                          position=1, 
                          leave=True)

# ================================================================================
dataset.set_split('val')

val_bar = tqdm_notebook(desc='split=val',
                        total=dataset.get_num_batches(args.batch_size), 
                        position=1, 
                        leave=True)

# ================================================================================
try:
    for epoch_index in range(args.num_epochs):
        train_state['epoch_index'] = epoch_index

        # ================================================================================
        # Iterate over training dataset

        # setup: batch generator, set loss and acc to 0, set train mode on

        dataset.set_split('train')

        batch_generator = generate_batches(dataset, 
                                           batch_size=args.batch_size, 
                                           device=args.device)
        running_loss = 0.0
        running_acc = 0.0
        classifier.train()
        
        # ================================================================================
        for batch_index, batch_dict in enumerate(batch_generator):
            # print("batch_index",batch_index)
            # batch_index 0

            # print("batch_dict",batch_dict["x_data"].shape)
            # torch.Size([32, 6])
            
            # print("batch_dict",batch_dict["y_target"].shape)
            # torch.Size([32])
            
            # print("batch_dict",batch_dict)
            # {'x_data': tensor([[1611,   72,  484, 3233,   50, 1735],
            #                    [1769,   49, 1088, 4174,  160,  140],
            #                    [3735, 5839,   27,    4,  289,    8],
            #                    [ 420,   50, 2705,    4, 4833,    2],
            #                    [   2, 2634,   15,   53, 3668,   34],
            #                    [ 215,   33,  286, 1292,   39,  413],
            #                    [ 255, 3862,    2,   85, 2962, 1279],
            #                    [ 258, 1995,   49, 1584,    2,  239],
            #                    [2584,   72,  104, 2753, 1082,    2],
            #                    [   4, 3770,  987,   49, 4723,  715],
            #                    [1992,    2,  578,    2,   23,    0],
            #                    [ 319,   33, 1569,   19, 5306, 5355],
            #                    [  48,  272,   19,    4, 2213,   27],
            #                    [   4,    2, 3685,    2,    0,    0],
            #                    [  27,   39, 1628,   43,    0,    0],
            #                    [  53,    4, 2961,   43,    0,    0],
            #                    [1472,    2,   49, 1437,  115, 1733],
            #                    [2648,  145,   72,  294,  824,   49],
            #                    [1680, 5911,   33,    2,   14,   15],
            #                    [  72, 1463,    4, 4579,    2,    4],
            #                    [1066,  298, 4734,   82,    4,  189],
            #                    [ 578,    2,   27,   28,  673,   33],
            #                    [  85,  436,   15,   43,    0,    0],
            #                    [ 879,    2,  257,  563,   48,  152],
            #                    [  83,    2,  136,  362, 1021,    2],
            #                    [ 788, 5582,    2,   49,    0,    0],
            #                    [3541,  228, 1245,    4,  202,   48],
            #                    [1307, 3177,   82, 2802,  471,  569],
            #                    [ 136, 1412,   49,    2,   68,   48],
            #                    [ 354,  115,   19,   15,   43,    0],
            #                    [   2,    4,  383,    4, 1380,  223],
            #                    [1858,   33,  402,   33,  282,   36]], device='cuda:0'), 
            #  'y_target': tensor([ 349, 3507, 1533,   60,    5, 4521,   60, 2045, 2596, 4969,   50, 1396,
            #                  2755, 3684,   15,   88,  255,  185,   67, 3585, 4735,  334,   48,   27,
            #                   127,  319,   60,   50, 1559,  277,   33,   49], device='cuda:0')}

            # {'x_data': tensor([[1611,   72,  484, 3233,   50, 1735],
            #                    [1769,   49, 1088, 4174,  160,  140],
            #                    [ 354,  115,   19,   15,   43,    0],
            #                    [   2,    4,  383,    4, 1380,  223],
            #                    [1858,   33,  402,   33,  282,   36]], device='cuda:0'), 
            #  'y_target': tensor([ 349, 3507, 1533,   60,    5, 4521,   60, 2045, 2596, 4969,   50, 1396,
            #                      2755, 3684,   15,   88,  255,  185,   67, 3585, 4735,  334,   48,   27,
            #                       127,  319,   60,   50, 1559,  277,   33,   49], device='cuda:0')}

            # ================================================================================
            # the training routine is these 5 steps:

            # --------------------------------------
            # step 1. zero the gradients
            optimizer.zero_grad()

            # step 2. compute the output
            y_pred = classifier(x_in=batch_dict['x_data'])

            # step 3. compute the loss
            loss = loss_func(y_pred, batch_dict['y_target'])
            loss_t = loss.item()
            running_loss += (loss_t - running_loss) / (batch_index + 1)

            # step 4. use loss to produce gradients
            loss.backward()

            # step 5. use optimizer to take gradient step
            optimizer.step()

            # -----------------------------------------
            # compute the accuracy
            acc_t = compute_accuracy(y_pred, batch_dict['y_target'])

            running_acc += (acc_t - running_acc) / (batch_index + 1)

            # update bar
            train_bar.set_postfix(loss=running_loss, acc=running_acc, epoch=epoch_index)
            train_bar.update()

        train_state['train_loss'].append(running_loss)
        train_state['train_acc'].append(running_acc)

        # Iterate over val dataset

        # setup: batch generator, set loss and acc to 0; set eval mode on
        dataset.set_split('val')
        batch_generator = generate_batches(dataset, 
                                           batch_size=args.batch_size, 
                                           device=args.device)
        running_loss = 0.
        running_acc = 0.
        classifier.eval()

        for batch_index, batch_dict in enumerate(batch_generator):

            # compute the output
            y_pred =  classifier(x_in=batch_dict['x_data'])

            # step 3. compute the loss
            loss = loss_func(y_pred, batch_dict['y_target'])
            loss_t = loss.item()
            running_loss += (loss_t - running_loss) / (batch_index + 1)

            # compute the accuracy
            acc_t = compute_accuracy(y_pred, batch_dict['y_target'])
            running_acc += (acc_t - running_acc) / (batch_index + 1)
            val_bar.set_postfix(loss=running_loss, acc=running_acc, 
                            epoch=epoch_index)
            val_bar.update()

        train_state['val_loss'].append(running_loss)
        train_state['val_acc'].append(running_acc)

        train_state = update_train_state(args=args, model=classifier,
                                         train_state=train_state)

        scheduler.step(train_state['val_loss'][-1])

        if train_state['stop_early']:
            break

        train_bar.n = 0
        val_bar.n = 0
        epoch_bar.update()
except KeyboardInterrupt:
    print("Exiting loop")

# ================================================================================
# Test set
# Best model
# Loss and Accuracy

# model_fn=train_state['model_filename']
# print("model_fn",model_fn)
# model_storage/ch5/cbow/model.pth

model_fn="model_storage/ch5/cbow/model_trained_backup.pth"

loaded_ckpt=torch.load(model_fn)
# print("loaded_ckpt",loaded_ckpt)
# OrderedDict([
#     ('embedding.weight', tensor([[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
#         [ 1.3514, -0.2759, -1.5108,  ..., -1.3651, -0.1655,  0.9623],
#         [-0.2547, -0.2556, -0.5690,  ..., -0.6383, -0.5526,  0.2665],
#         ...,
#         [ 1.9376,  0.4060,  0.9361,  ..., -0.4493,  0.9146,  0.0620],
#         [ 0.5388, -0.2807,  0.7139,  ..., -0.1142, -1.3831, -0.3602],
#         [-1.0862,  1.8168,  0.3258,  ...,  0.9599, -0.3896,  0.5290]],
#        device='cuda:0')), 
#     ('fc1.weight', tensor([[ 0.0114,  0.3032, -0.0084,  ..., -0.1034, -0.3216,  0.3240],
#         [ 0.0114,  0.3032, -0.0084,  ..., -0.1034, -0.3216,  0.3240],
#         [ 0.0485, -0.0261, -0.0042,  ...,  0.0067,  0.0777, -0.0762],
#         ...,
#         [-0.0184,  0.2450,  0.0586,  ..., -0.1067, -0.0627,  0.2661],
#         [-0.0126,  0.1061, -0.0457,  ..., -0.0482, -0.1407,  0.1534],
#         [ 0.0114,  0.3033, -0.0084,  ..., -0.1034, -0.3216,  0.3240]],
#        device='cuda:0')), 
#     ('fc1.bias', tensor([-3.1369, -3.2102,  1.9988,  ..., -0.2273, -0.1729, -3.1940],
#        device='cuda:0'))])

# ================================================================================
# print('loaded_ckpt["embedding.weight"]',loaded_ckpt["embedding.weight"].shape)
# torch.Size([6138, 50])

# print('loaded_ckpt["fc1.weight"]',loaded_ckpt["fc1.weight"].shape)
# torch.Size([6138, 50])

# print('loaded_ckpt["fc1.bias"]',loaded_ckpt["fc1.bias"].shape)
# torch.Size([6138])

# ================================================================================
# print("classifier",classifier)
# CBOWClassifier(
#   (embedding): Embedding(6138, 50, padding_idx=0)
#   (fc1): Linear(in_features=50, out_features=6138, bias=True)
# )

classifier.load_state_dict(loaded_ckpt)

# ================================================================================
classifier = classifier.to(args.device)

# ================================================================================
loss_func = nn.CrossEntropyLoss()

# ================================================================================
dataset.set_split('test')

# ================================================================================
batch_generator = generate_batches(dataset, 
                                   batch_size=args.batch_size, 
                                   device=args.device)

# ================================================================================
running_loss = 0.
running_acc = 0.
classifier.eval()

# ================================================================================
for batch_index, batch_dict in enumerate(batch_generator):
    # compute the output
    y_pred =  classifier(x_in=batch_dict['x_data'])
    
    # compute the loss
    loss = loss_func(y_pred, batch_dict['y_target'])
    loss_t = loss.item()
    running_loss += (loss_t - running_loss) / (batch_index + 1)

    # compute the accuracy
    acc_t = compute_accuracy(y_pred, batch_dict['y_target'])
    running_acc += (acc_t - running_acc) / (batch_index + 1)

train_state['test_loss'] = running_loss
train_state['test_acc'] = running_acc

print("Test loss: {};".format(train_state['test_loss']))
print("Test Accuracy: {}".format(train_state['test_acc']))
# Test loss: 7.719295198777138;
# Test Accuracy: 13.500000000000002

# ================================================================================
# Trained Embeddings

def pretty_print(results):
    """
    Pretty print embedding results.
    """
    for item in results:
        print ("...[%.2f] - %s"%(item[1], item[0]))

def get_closest(target_word, word_to_idx, embeddings, n=5):
    """
    Get the n closest
    words to your word.
    """

    # Calculate distances to all other words
    word_embedding = embeddings[word_to_idx[target_word.lower()]]

    distances = []

    for word, index in word_to_idx.items():
        if word == "<MASK>" or word == target_word:
            continue
        distances.append((word, torch.dist(word_embedding, embeddings[index])))
    
    results = sorted(distances, key=lambda x: x[1])[1:n+2]
    return results

# ================================================================================
word = input('Enter a word: ')

# ================================================================================
# c embeddings: weight of trained embedding
embeddings = classifier.embedding.weight.data

# ================================================================================
cbow_vocab_from_vectorizer=vectorizer.cbow_vocab
print("cbow_vocab_from_vectorizer",cbow_vocab_from_vectorizer)

word_to_idx = cbow_vocab_from_vectorizer._token_to_idx
# print("word_to_idx",word_to_idx)
# {'<MASK>': 0, '<UNK>': 1, ',': 2, 'or': 3, 'the': 4, 'frankenstein': 5, 'modern': 6, 'prometheus': 7, 'by': 8, 'mary': 9, 
#  'wollstonecraft': 10, 'godwin': 11, 'shelley': 12, 'letter': 13, 'st': 14, '.': 15, 'petersburgh': 16, 'dec': 17, 
#  'th': 18, 'to': 19, 'mrs': 20, 'saville': 21, 'england': 22, 'you': 23, 'will': 24, 'rejoice': 25, 'hear': 26, 
#  'that': 27, 'no': 28, 'disaster': 29, 'has': 30, 'accompanied': 31, 'commencement': 32, 'of': 33, 'an': 34, 
#  'enterprise': 35, 'which': 36, 'have': 37, 'regarded': 38, 'with': 39, 'such': 40, 'evil': 41, 'forebodings': 42, '': 43, 

# ================================================================================
# Inputted word: name

closest_words=get_closest(word, word_to_idx, embeddings, n=5)
# print("closest_words",closest_words)
# [('blue', tensor(6.6412, device='cuda:0')), 
#  ('whence', tensor(6.6471, device='cuda:0')), 
#  ('beggars', tensor(6.6531, device='cuda:0')), 
#  ('fluctuating', tensor(6.7025, device='cuda:0')), 
#  ('villagers', tensor(6.7265, device='cuda:0')), 
#  ('occupation', tensor(6.7375, device='cuda:0'))]

pretty_print(closest_words)

# ================================================================================
target_words = ['frankenstein', 'monster', 'science', 'sickness', 'lonely', 'happy']

# ================================================================================
# c embeddings: trained embedding array
embeddings = classifier.embedding.weight.data

cbow_vocab_from_vectorizer=vectorizer.cbow_vocab
print("cbow_vocab_from_vectorizer",cbow_vocab_from_vectorizer)

word_to_idx = cbow_vocab_from_vectorizer._token_to_idx

# ================================================================================
for target_word in target_words: 
    if target_word not in word_to_idx:
        print("Not in vocabulary")
        continue
    closest_words=get_closest(target_word, word_to_idx, embeddings, n=5)
    print("target_word",target_word)
    pretty_print(closest_words)
    print("")

    # target_word frankenstein
    # ...[7.18] - irradiated
    # ...[7.64] - men
    # ...[7.67] - enslaved
    # ...[7.71] - professor
    # ...[7.71] - mode
    # ...[7.74] - austria

    # target_word monster
    # ...[7.46] - cares
    # ...[7.67] - griefs
    # ...[7.73] - saw
    # ...[7.75] - trifling
    # ...[7.76] - relief
    # ...[7.76] - prolong

    # target_word science
    # ...[7.02] - impression
    # ...[7.08] - mutual
    # ...[7.10] - mist
    # ...[7.25] - swelling
    # ...[7.29] - darkened
    # ...[7.30] - gospel

    # target_word sickness
    # ...[6.28] - while
    # ...[6.57] - awoke
    # ...[6.62] - foundations
    # ...[6.69] - consoles
    # ...[6.70] - insanity
    # ...[6.72] - depend

    # target_word lonely
    # ...[6.74] - excessive
    # ...[6.90] - ought
    # ...[6.90] - moonlight
    # ...[7.07] - bed
    # ...[7.10] - three
    # ...[7.19] - superhuman

    # target_word happy
    # ...[6.40] - bottom
    # ...[6.43] - penetrated
    # ...[6.43] - altered
    # ...[6.45] - wand
    # ...[6.51] - danger
    # ...[6.54] - aid
