#!/usr/bin/env python
from math import log
from functools import lru_cache
import time
from numpy import exp
import numpy as np
import random

from utils.gradcheck import gradcheck_naive, grad_tests_softmax, grad_tests_negsamp
from utils.utils import normalizeRows, softmax

def sigmoid(x):
    """
    Compute the sigmoid function for the input here.
    Arguments:
    x -- A scalar or numpy array.
    Return:
    s -- sigmoid(x)
    """

    ### YOUR CODE HERE (~1 Line)
    local_sigmoid = lambda x : 1 / (1 + exp(-x))
    s = local_sigmoid(x)
    ### END YOUR CODE

    return s

def naiveSoftmaxLossAndGradient(
    centerWordVec,
    outsideWordIdx,
    outsideVectors,
    dataset
):
    """ Naive Softmax loss & gradient function for word2vec models

    Implement the naive softmax loss and gradients between a center word's 
    embedding and an outside word's embedding. This will be the building block
    for our word2vec models.

    Arguments:
    centerWordVec -- numpy ndarray, center word's embedding
                    in shape (word vector length, )
                    (v_c in the pdf handout)
    outsideWordIdx -- integer, the index of the outside word
                    (o of u_o in the pdf handout)
    outsideVectors -- outside vectors is
                    in shape (num words in vocab, word vector length) 
                    for all words in vocab (U in the pdf handout)
    dataset -- needed for negative sampling, unused here.

    Return:
    loss -- naive softmax loss
    gradCenterVec -- the gradient with respect to the center word vector
                     in shape (word vector length, )
                     (dJ / dv_c in the pdf handout)
    gradOutsideVecs -- the gradient with respect to all the outside word vectors
                    in shape (num words in vocab, word vector length) 
                    (dJ / dU)
    """
    word_vector_length = centerWordVec.shape[0]
    num_words = outsideVectors.shape[0]

    reshapedCenterWordVec = centerWordVec.reshape([1, word_vector_length])

    # [1 x number of words in vocab]
    all_products_vector = reshapedCenterWordVec.dot(outsideVectors.transpose())

    # [1 x number of words in vocab]
    # print(f"Reshaped center word vec has shape {reshapedCenterWordVec.shape}")
    # print(f"Outside vectors has shape {outsideVectors.shape}")
    # print(f"Shapes of all_products is {all_products_vector.shape}")
    y_hat = softmax(all_products_vector)
    # print(f"y_hat shape is {y_hat.shape}")
    loss = -np.log(y_hat[0][outsideWordIdx])

    # [1 x number of words in vocab]
    y = np.zeros([1,num_words])
    y[0][outsideWordIdx] = 1

    # 1 x number of words in vocab * [number of words in vocab  * word_vector_length
    gradCenterVec = (y_hat-y).dot(outsideVectors)  # 1 x vocab_length * vocab_length x word_vector_length
    gradCenterVec = gradCenterVec.reshape(word_vector_length)

    # [num_of_words_in_vocab x 1][1 x word_vector_length]
    gradOutsideVecs = ((y_hat-y).transpose()).dot(reshapedCenterWordVec)


    ### YOUR CODE HERE (~6-8 Lines)

    ### Please use the provided softmax function (imported earlier in this file)
    ### This numerically stable implementation helps you avoid issues pertaining
    ### to integer overflow. 

    ### END YOUR CODE

    return loss, gradCenterVec, gradOutsideVecs


def getNegativeSamples(outsideWordIdx, dataset, K):
    """ Samples K indexes which are not the outsideWordIdx """

    negSampleWordIndices = [None] * K
    for k in range(K):
        newidx = dataset.sampleTokenIdx()
        while newidx == outsideWordIdx:
            newidx = dataset.sampleTokenIdx()
        negSampleWordIndices[k] = newidx
    return negSampleWordIndices


def negSamplingLossAndGradient(
    centerWordVec,
    outsideWordIdx,
    outsideVectors,
    dataset,
    K=10
):
    """ Negative sampling loss function for word2vec models

    Implement the negative sampling loss and gradients for a centerWordVec
    and a outsideWordIdx word vector as a building block for word2vec
    models. K is the number of negative samples to take.

    Note: The same word may be negatively sampled multiple times. For
    example if an outside word is sampled twice, you shall have to
    double count the gradient with respect to this word. Thrice if
    it was sampled three times, and so forth.

    Arguments/Return Specifications: same as naiveSoftmaxLossAndGradient
    """

    # Negative sampling of words is done for you. Do not modify this if you
    # wish to match the autograder and receive points!
    negSampleWordIndices = getNegativeSamples(outsideWordIdx, dataset, K)
    indices = [outsideWordIdx] + negSampleWordIndices

    ### YOUR CODE HERE (~10 Lines)
    wordVectorLength = centerWordVec.shape[0]
    numWords = outsideVectors.shape[0]

    reshapedCenterWordVec = centerWordVec.reshape(wordVectorLength)
    u_o = outsideVectors[outsideWordIdx].reshape(wordVectorLength)

    negativeSamples = [outsideVectors[index] for index in negSampleWordIndices]

    negativeSamplesProducts = [
        reshapedCenterWordVec.dot(sampled) for sampled in negativeSamples
    ]

    first_loss = -np.log(sigmoid(u_o.dot(reshapedCenterWordVec)))
    second_loss = -np.sum(
            [np.log(sigmoid(-sampled)) for sampled in negativeSamplesProducts])
    
    loss = first_loss + second_loss
    gradCenterVec = -(1 - sigmoid(reshapedCenterWordVec.dot(u_o)))*u_o
    gradCenterVec = gradCenterVec.reshape(wordVectorLength)

    toAddVectors = [ (1 - sigmoid(-negativeSample.dot(reshapedCenterWordVec)))*negativeSample
                       for negativeSample in negativeSamples]
    toAddVec = np.add.reduce(toAddVectors)

    gradCenterVec = np.add(gradCenterVec, toAddVec)
    gradOutsideVecs = np.zeros([numWords, wordVectorLength])

    u_o_derivative = -(1 - sigmoid(u_o.dot(reshapedCenterWordVec))) * reshapedCenterWordVec
    gradOutsideVecs[outsideWordIdx] = np.add(gradOutsideVecs[outsideWordIdx], u_o_derivative)
    for index in negSampleWordIndices:
        gradient_vector = (1 - sigmoid(-outsideVectors[index].dot(reshapedCenterWordVec))) * reshapedCenterWordVec
        gradOutsideVecs[index] = np.add(gradOutsideVecs[index], gradient_vector)

    return loss, gradCenterVec, gradOutsideVecs


def skipgram(currentCenterWord, windowSize, outsideWords, word2Ind,
             centerWordVectors, outsideVectors, dataset,
             word2vecLossAndGradient=naiveSoftmaxLossAndGradient):
    """ Skip-gram model in word2vec

    Implement the skip-gram model in this function.

    Arguments:
    currentCenterWord -- a string of the current center word
    windowSize -- integer, context window size
    outsideWords -- list of no more than 2*windowSize strings, the outside words
    word2Ind -- a dictionary that maps words to their indices in
              the word vector list
    centerWordVectors -- center word vectors (as rows) is in shape 
                        (num words in vocab, word vector length) 
                        for all words in vocab (V in pdf handout)
    outsideVectors -- outside vectors is in shape 
                        (num words in vocab, word vector length) 
                        for all words in vocab (U in the pdf handout)
    word2vecLossAndGradient -- the loss and gradient function for
                               a prediction vector given the outsideWordIdx
                               word vectors, could be one of the two
                               loss functions you implemented above.

    Return:
    loss -- the loss function value for the skip-gram model
            (J in the pdf handout)
    gradCenterVec -- the gradient with respect to the center word vector
                     in shape (word vector length, )
                     (dJ / dv_c in the pdf handout)
    gradOutsideVecs -- the gradient with respect to all the outside word vectors
                    in shape (num words in vocab, word vector length) 
                    (dJ / dU)
    """
    loss = 0.0

    gradCenterVecs = np.zeros(centerWordVectors.shape)
    gradOutsideVecs = np.zeros(outsideVectors.shape)
    currentCenterInd = word2Ind[currentCenterWord]
    #print(f"currentCenterInd is {currentCenterInd}")
    currentCenterVector = centerWordVectors[currentCenterInd]

    for word in outsideWords:
        currentOutsideWordInd = word2Ind[word]
        # print(f"Current outsideword ind is {currentOutsideWordInd}")
        localLoss, localGradCenterVec, localGradOutsideVec = word2vecLossAndGradient(
            currentCenterVector, currentOutsideWordInd,
            outsideVectors, dataset
        )
        loss += localLoss
        gradCenterVecs[currentCenterInd] = np.add(gradCenterVecs[currentCenterInd], localGradCenterVec)

        gradOutsideVecs = np.add(gradOutsideVecs, localGradOutsideVec)


    ### YOUR CODE HERE (~8 Lines)

    ### END YOUR CODE
    
    return loss, gradCenterVecs, gradOutsideVecs


#############################################
# Testing functions below. DO NOT MODIFY!   #
#############################################

def word2vec_sgd_wrapper(word2vecModel, word2Ind, wordVectors, dataset, 
                         windowSize,
                         word2vecLossAndGradient=naiveSoftmaxLossAndGradient):
    batchsize = 50
    loss = 0.0
    grad = np.zeros(wordVectors.shape)
    N = wordVectors.shape[0]

    centerWordVectors = wordVectors[:int(N/2),:]
    # print(f"Center word vectors have shape {centerWordVectors.shape}")
    outsideVectors = wordVectors[int(N/2):,:]
    for i in range(batchsize):
        windowSize1 = random.randint(1, windowSize)
        centerWord, context = dataset.getRandomContext(windowSize1)

        c, gin, gout = word2vecModel(
            centerWord, windowSize1, context, word2Ind, centerWordVectors,
            outsideVectors, dataset, word2vecLossAndGradient
        )
        # print(f"\n \n Built wordvec model and got \n {c} \n \n {gin} \n \n {gout}")
        loss += c / batchsize
        grad[:int(N/2), :] += gin / batchsize
        grad[int(N/2):, :] += gout / batchsize

    return loss, grad


def test_word2vec():
    """ Test the two word2vec implementations, before running on Stanford Sentiment Treebank """
    dataset = type('dummy', (), {})()
    def dummySampleTokenIdx():
        return random.randint(0, 4)

    def getRandomContext(C):
        tokens = ["a", "b", "c", "d", "e"]
        return tokens[random.randint(0,4)], \
            [tokens[random.randint(0,4)] for i in range(2*C)]
    dataset.sampleTokenIdx = dummySampleTokenIdx
    dataset.getRandomContext = getRandomContext

    random.seed(31415)
    np.random.seed(9265)
    dummy_vectors = normalizeRows(np.random.randn(10,3))
    dummy_tokens = dict([("a",0), ("b",1), ("c",2),("d",3),("e",4)])

    print("Naive softmax first")
    print("==== Gradient check for skip-gram with naiveSoftmaxLossAndGradient ====")
    gradcheck_naive(lambda vec: word2vec_sgd_wrapper(
        skipgram, dummy_tokens, vec, dataset, 5, naiveSoftmaxLossAndGradient),
        dummy_vectors, "naiveSoftmaxLossAndGradient Gradient")
    grad_tests_softmax(skipgram, dummy_tokens, dummy_vectors, dataset)

    print("==== Gradient check for skip-gram with negSamplingLossAndGradient ====")
    gradcheck_naive(lambda vec: word2vec_sgd_wrapper(
        skipgram, dummy_tokens, vec, dataset, 5, negSamplingLossAndGradient),
        dummy_vectors, "negSamplingLossAndGradient Gradient")

    grad_tests_negsamp(skipgram, dummy_tokens, dummy_vectors, dataset, negSamplingLossAndGradient)


if __name__ == "__main__":
    test_word2vec()
