# -- Code Cell --
import pandas as pd
import numpy as np
train = pd.read_csv("./train_data.csv")
test = pd.read_csv("./test_data.csv")

# -- Code Cell --
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms
from PIL import Image
from torchvision.transforms import functional as TF

# -- Code Cell --
train.head()

# -- Code Cell --
from tensorflow.keras.preprocessing.text import Tokenizer

# -- Code Cell --
texts = train['text']

# -- Code Cell --
embeddings_index = {}

with open("glove.6B.200d.txt", encoding="utf-8") as f:
    for line in f:
        values = line.split()
        word = values[0]
        vector = np.asarray(values[1:], dtype='float32')
        embeddings_index[word] = vector

# -- Code Cell --
tokenizer = Tokenizer()
tokenizer.fit_on_texts(texts)
word_index = tokenizer.word_index

# -- Code Cell --
print("t")

# -- Code Cell --
import emoji

# -- Code Cell --
word_index

# -- Code Cell --
from nltk.tokenize import word_tokenize
text = train['text']
train_text_tokenized = {}
for idx,t in enumerate(text):
    train_text_tokenized[train['sentence_id'][idx]] = word_tokenize(t)
train_text_tokenized

# -- Code Cell --
# test_text_converted = []
# for text in test['emoji_sequence']:
#     text_message= emoji.demojize(text)
#     text_message = text_message.split(":")
#     for t in text_message:
#         if len(t) == 0:
#             text_message.remove(t)
#     test_text_converted.append(text_message)
# test_text_converted

# -- Markdown Cell --
# ### INCERCAM CUSTOM MAP, in loc de SNAKE -> TRAITOR

# -- Code Cell --
emoji_map = {
    "❤️": "love",
    "💛": "love",
    "🤍": "love",
    "🐍": "traitor",
    "💨": "fast",
    "🧢": "lie",
    "👀": "watch",
    "💔": "heartbreak",
    "😎": "cool",
    "🤥": "lie",
    "☠️": "death",
    "💀": "death",
    "😴": "sleep",
    "🤫": "secret",
    "🤡": "crazy",
    "👻": "ghost",
    "🐐": "greatest"
}

# -- Code Cell --
test_text_converted = []

for text in test['emoji_sequence']:
    
    for e, word in emoji_map.items():
        text = text.replace(e, " " + word + " ")
    text = emoji.demojize(text)
    text = text.replace(":", " ")
    text = text.replace("_", " ")
    tokens = text.lower().split()
    test_text_converted.append(tokens)
test_text_converted

# -- Code Cell --
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

def sentence_embedding(tokens, embeddings_index, embedding_dim=200):
    vectors = [embeddings_index[word] for word in tokens if word in embeddings_index]
    if len(vectors) == 0:
        return np.zeros(embedding_dim)
    return np.mean(vectors, axis=0)

# -- Code Cell --
group1_vectors = np.array([sentence_embedding(seq, embeddings_index, 200) for seq in train_text_tokenized.values()])

# -- Code Cell --
group2_vectors = np.array([sentence_embedding(seq, embeddings_index, 200) for seq in test_text_converted])

# -- Code Cell --
sim_matrix = cosine_similarity(group2_vectors, group1_vectors)

# -- Code Cell --
best_matches = np.argmax(sim_matrix, axis=1)
preds_final=[]
for i, match_idx in enumerate(best_matches):
    # print("querry:", test_text_converted[i])
    # print("besst:", train_text_tokenized[int(f"1{match_idx:03d}")]) #taranie maxima
    # print("score:", sim_matrix[i, match_idx])
    # print(i)
    qr = test_text_converted[i]
    score = sim_matrix[i, match_idx]
    best = train_text_tokenized[int(f"1{match_idx:03d}")]
    preds_final.append(int(f"1{match_idx:03d}"))
    if score<0.5:
        print(qr)
        print(best)
        print(score)
        print("-------")


# -- Code Cell --
preds_final

# -- Code Cell --
subs = pd.DataFrame({
    'subtaskID':1,
    'datapointID':test['sample_id'],
    'answer':preds_final
}).to_csv("subs.csv",index=False)

# -- Code Cell --
