#%%

"""
Credits: This code is adapted from the textbook "Deep Learning with Python", 
2nd Edition, by François Chollet. 
"""

#%% 

import numpy as np
import string

#%% Code to load pre-trained GloVe embeddings, 

glove_file = "../../../../../home/cse4392_data/glove/glove.6b.100d.txt"
embeddings_index = {}
with open(glove_file, encoding='utf-8') as f:
    for line in f:
        word, coefs = line.split(maxsplit=1)
        coefs = np.fromstring(coefs, "f", sep=" ")
        embeddings_index[word] = coefs

print(f"Found {len(embeddings_index)} word vectors.")

#%% Auxiliary functions

# standardize a string, by converting to lower case and removing
# punctuation characters.
def standardize(input_string):
    strip_chars = string.punctuation
    lowercase = input_string.lower()
    result = ""
    for letter in lowercase:
        if (letter not in strip_chars):
            result = result + letter
    
    return result

# convert a piece of text to a numpy array of GloVe embeddings of the words
# in the text
def text_to_glove(text, embeddings):
    text = standardize(text)
    words = text.split()
    result = []
    
    for word in words:
        vector = embeddings[word]
        result = result + [vector]
    return np.array(result)

# applies the softmax function to the given vector
def softmax(vector):
    exps = np.exp(vector)
    total = exps.sum()
    result = exps / total
    return result


# Returns the score matrix (normalized dot products) between the given
# vectors (which represent words in some piece of text). Also returns
# the softmax_scores, with the softmax function applied to each row
# of the score matrix.
# Also returns
# the new vectors that are obtained by taking, for each vector (called
# a pivot vector), the 
# weighted sum of all vectors (with the weights coming from the score
# matrix values corresponding to the pivot vector).
def self_attention(word_vectors):
    new_vectors = np.zeros(shape=word_vectors.shape)
    (num_words, num_dims) = word_vectors.shape
    scores = np.zeros((num_words, num_words))
    softmax_scores = np.zeros((num_words, num_words))

    for i, pivot_vector in enumerate(word_vectors):
        pivot_scores = np.zeros((num_words,))

        for j, vector in enumerate(word_vectors):
            pivot_scores[j] = np.dot(pivot_vector, vector.T)
            norm1 = np.linalg.norm(pivot_vector)
            norm2 = np.linalg.norm(vector)
#            pivot_scores[j] = pivot_scores[j] / (norm1 * norm2)

        pivot_scores /= np.sqrt(num_dims)
        pivot_scores_softmax = softmax(pivot_scores)
        new_pivot_representation = np.zeros(shape=pivot_vector.shape)
        for j, vector in enumerate(word_vectors):
            new_pivot_representation += vector * pivot_scores_softmax[j]

        new_vectors[i] = new_pivot_representation
        scores[i] = pivot_scores
        softmax_scores[i] = pivot_scores_softmax

    return (scores, softmax_scores, new_vectors)


#%%

#text = "This is a train station."
#text = "The train left the station on time."
#text = "The leaves are falling from the tree."
text = "I bought a spring mattress."
#text = "Spring is the best season."
#text = "We get water from a natural spring."


vectors = text_to_glove(text, embeddings_index)
(scores, softmax_scores, new_vectors) = self_attention(vectors)
np.set_printoptions(formatter={'float': '{: 0.4f}'.format})
print(scores, "\n")

print(softmax_scores)
#%%

a = self_attention(vec)