import gensim
import numpy as np
import torch
import torch.nn as nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_word2vec(word2vec_data_path):
    return gensim.models.KeyedVectors.load_word2vec_format(word2vec_data_path, binary=True)


def generate_word_map_from_word2vec_model(word2vec_model):
    word_map = {}
    vocab = word2vec_model.wv.vocab
    vocab_word_list = list(word2vec_model.wv.vocab.keys())
    for word_id, word in enumerate(vocab_word_list):
        word_map[word] = word_id
    return word_map

def add_special_word_to_embedding_vectors(embedding_vectors, word_map, word):
    """
    Args:
    embedding_vectors: np.ndarray
                       shape: (n_vocab, n_dimension)
        the embedding vectors, always generated by gensim.model.wv.vectors 
    """
    if word_map.get(word) is not None:
        return embedding_vectors, word_map
    
    n_vocab, n_dimension = embedding_vectors.shape
    word_embedding = np.random.rand(1, n_dimension)
    embedding_vectors = np.append(
        embedding_vectors, 
        word_embedding,
        axis=0)
    word_map[word] = n_vocab
    return embedding_vectors, word_map


# deprecated!
def convert_words_to_word_embeddings(embedding_vectors, word_map, words, default_dim=300):
    word_embeddings = []
    for word in words:
        if word_map.get(word) is None:
            embedding = [0.0] * default_dim
            continue
        embedding = embedding_vectors[word_map[word]]
        word_embeddings.append(embedding)
    return word_embeddings

def create_embedding_layer(embedding_vectors, is_trainable=False):
    num_embeddings, embedding_dim = embedding_vectors.shape
    # embedding_layer = nn.Embedding(num_embeddings, embedding_dim)
    # embedding_layer.load_state_dict({'weight': embedding_vectors})
    weights = torch.FloatTensor(embedding_vectors, device=device)
    embedding_layer = nn.Embedding.from_pretrained(weights)
    if is_trainable:
        embedding_layer.weight.requires_grad = False
    return embedding_layer, num_embeddings, embedding_dim

import gensim
import numpy as np
import torch
import torch.nn as nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_word2vec(word2vec_data_path):
    return gensim.models.KeyedVectors.load_word2vec_format(word2vec_data_path, binary=True)


def generate_word_map_from_word2vec_model(word2vec_model):
    word_map = {}
    vocab = word2vec_model.wv.vocab
    vocab_word_list = list(word2vec_model.wv.vocab.keys())
    for word_id, word in enumerate(vocab_word_list):
        word_map[word] = word_id
    return word_map

def add_special_word_to_embedding_vectors(embedding_vectors, word_map, word):
    """
    Args:
    embedding_vectors: np.ndarray
                       shape: (n_vocab, n_dimension)
        the embedding vectors, always generated by gensim.model.wv.vectors 
    """
    if word_map.get(word) is not None:
        return embedding_vectors, word_map
    
    n_vocab, n_dimension = embedding_vectors.shape
    word_embedding = np.random.rand(1, n_dimension)
    embedding_vectors = np.append(
        embedding_vectors, 
        word_embedding,
        axis=0)
    word_map[word] = n_vocab
    return embedding_vectors, word_map


# deprecated!
def convert_words_to_word_embeddings(embedding_vectors, word_map, words, default_dim=300):
    word_embeddings = []
    for word in words:
        if word_map.get(word) is None:
            embedding = [0.0] * default_dim
            continue
        embedding = embedding_vectors[word_map[word]]
        word_embeddings.append(embedding)
    return word_embeddings

def create_embedding_layer(embedding_vectors, is_trainable=False):
    num_embeddings, embedding_dim = embedding_vectors.shape
    # embedding_layer = nn.Embedding(num_embeddings, embedding_dim)
    # embedding_layer.load_state_dict({'weight': embedding_vectors})
    weights = torch.FloatTensor(embedding_vectors, device=device)
    embedding_layer = nn.Embedding.from_pretrained(weights)
    if is_trainable:
        embedding_layer.weight.requires_grad = False
    return embedding_layer, num_embeddings, embedding_dim

Args:     embedding_vectors: np.ndarray                        shape: (n_vocab, n_dimension)         the embedding vectors, always generated by gensim.model.wv.vectors # deprecated! # embedding_layer = nn.Embedding(num_embeddings, embedding_dim) # embedding_layer.load_state_dict({'weight': embedding_vectors})