from typing import Dict
import torch


def get_emb_matrix(vocab: Dict, emb_path: str) -> torch.Tensor:
    """Takes a embedding file and convert them into embedding matrix using the
    vocab dictionary. Eg. if the word 'the' is mapped to id 5 in the vocab, then
    the 5th row of the embedding matrix will represent the embedding vector
    of the word 'the'.

    Args:
        vocab (Dict): mapping of words to id used by tokenizer
        emb_path (str): path of embedding file

    Returns:
        torch.Tensor: torch tensor of embedding matrix
    """
    with open(emb_path, "r") as emb_file:
        file_lines = emb_file.readlines()

    emb_dict = {}
    first_line = file_lines[0].rstrip().split()
    if len(first_line) != 2:
        raise ValueError(
            "The first line in W2V embeddings must be the pair (vocab_size, emb_dim)"
        )

    emb_dim = int(first_line[1])
    emb_vocab_size = int(first_line[0])
    for line in file_lines[1:]:
        tokens = line.rstrip().split()
        word = tokens[0]
        vector = [float(token) for token in tokens[1:]]
        if len(vector) != emb_dim:
            raise ValueError("The number of dimensions does not match the header info")
        emb_dict[word] = vector

    if emb_vocab_size != len(emb_dict.keys()):
        raise ValueError("Vocab size does not match the header info")

    vocab_len = len(vocab.keys())
    emb_matrix = torch.zeros((vocab_len, emb_dim))
    for word, index in vocab.items():
        if word in emb_dict.keys():
            emb_matrix[index] = torch.Tensor(emb_dict[word])

    return emb_matrix