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

# %%
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

# %%
train.head()

# %%
from tensorflow.keras.preprocessing.text import Tokenizer

# %%
texts = train['text']

# %%
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

# %%
tokenizer = Tokenizer()
tokenizer.fit_on_texts(texts)
word_index = tokenizer.word_index

# %%
print("t")

# %%
import emoji

# %%
word_index

# %%
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

# %%
# 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]
# ### INCERCAM CUSTOM MAP, in loc de SNAKE -> TRAITOR

# %%
emoji_map = {
    "❤️": "love",
    "💛": "love",
    "🤍": "love",
    "🐍": "traitor",
    "💨": "fast",
    "🧢": "lie",
    "👀": "watch",
    "💔": "sorrow",
    "😎": "cool",
    "🤥": "lie",
    "☠️": "death",
    "💀": "death",
    "😴": "sleep",
    "🤫": "secret",
    "🤡": "crazy",
    "👻": "ghost",
    "🐐": "greatest",
    "🚢": "cruise",
    "🐈":"scarred",
    "🎯":"target",
    "🏆":"best",
    "🌺":"cute flower",
    "👃":"powerfull smell",
    
}

# %%
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

# %%
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)

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

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

# %%
sim_matrix = cosine_similarity(group2_vectors, group1_vectors)

# %%
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("-------")


# %%
preds_final

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

# %%



