import tensorflow as tf
import numpy as np
import pickle as pk
from pathlib import Path

INPUT_EMBEDDING_SIZE = 
OUTPUT_EMBEDDING_SIZE = 
OUTPUT_VOCAB_SIZE = 

INPUT_EMBEDDINGS_FILE = ""


def load_input_embeddings():
    # GloVe embeddings for input sentences (English)
    embed_file = Path("embeddings.pkl")
    if embed_file.is_file():
        # embeddings already serialized, just load them
        print(f"Local Embeddings pickle found, loading...")
        with open("embeddings.pkl", 'rb') as f:
            return pk.load(f)
    else:
        # create the embeddings
        print(f"Building embeddings dictionary...")
        data = open(INPUT_EMBEDDINGS_FILE, 'r', encoding="utf-8")
        embeddings = [[0] * INPUT_EMBEDDING_SIZE]
        word_index_dict = {'UNK': 0}  # first row is for unknown words
        index = 1
        for line in data:
            split_line = line.split()
            word = tf.compat.as_str(split_line[0])
            embedding = [float(val) for val in split_line[1:]]
            embeddings.append(embedding)
            word_index_dict[word] = index
            index += 1
        data.close()

        # pickle them
        with open('embeddings.pkl', 'wb') as f:
            print(f"Creating local embeddings pickle for faster loading...")
            # Pickle the 'data' dictionary using the highest protocol available.
            pk.dump((embeddings, word_index_dict), f, pk.HIGHEST_PROTOCOL)

    return embeddings, word_index_dict


def glove_embeddings(embeddings, embedding_dict, in_sentence, embed_dim):
    num_words = len(in_sentence)
    embedded_sen = np.zeros([num_words, embed_dim])
    for w in range(num_words):
        embedded_sen[w, :] = embeddings[embedding_dict.get(in_sentence[w], 0)]
    return embedded_sen


def posit_embed_row(wk, t, embed_dim):
    sin = tf.math.sin(tf.math.multiply(wk, tf.cast(t, tf.float32)))
    cos = tf.math.cos(tf.math.multiply(wk, tf.cast(tf.add(t, tf.constant(1)), tf.float32)))

    sin_expanded = tf.expand_dims(sin, 1)
    cos_expanded = tf.expand_dims(cos, 1)

    concatenated = tf.concat([sin_expanded, cos_expanded], 1)
    return tf.reshape(concatenated, (1, embed_dim))


def posit_encode(sentence_length, embedding_dim):

    two_k = tf.linspace(0, tf.math.subtract(embedding_dim, 2), tf.cast(tf.math.divide(embedding_dim, tf.constant(2)), tf.int32))
    wk = tf.math.pow(tf.cast(10000, tf.float32), tf.cast(tf.divide(two_k, embedding_dim), tf.float32))

    ta = tf.TensorArray(tf.float32, size=sentence_length, element_shape=(1, embedding_dim))
    ta.write(0, posit_embed_row(wk, 0, embedding_dim))

    ta, _ = tf.while_loop(lambda ta, i : i < sentence_length,
                  lambda ta, i : [ta.write(i, posit_embed_row(wk, i, embedding_dim)), i+1],
                  loop_vars=[ta, 1],
                  shape_invariants=[None, tf.TensorShape(())],
                  )

    return ta.concat()


def get_output_dict(output_dict_embeddings):
    data = open(output_dict_embeddings, 'r', encoding="utf-8")
    word_index_list = ['UNK']
    word_index_dict = {'UNK': 0}
    index = 1
    for line in data:
        cur_word = line.split(",")[0]
        word_index_dict[cur_word] = index
        word_index_list.append(cur_word)
        index += 1
    return word_index_list, word_index_dict
