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

# -- Code Cell --
count = 0
for i in range(len(test)):
    for j in range(len(test)):
        if test['Word'][i] == test['Similar'][j]:
            count+=1
task1 = count

# -- Code Cell --
task1 = 124

# -- Code Cell --
train['english'] = train['english'].str.split()
train['romanian'] = train['romanian'].str.split()
train['romanian'] = train['romanian'].fillna("")
train['english'] = train['english'].fillna("")

# -- Code Cell --
model_path_en = './wiki.en.small.vec'
model_path_ro = './wiki.ro.small.vec'
from gensim.models import KeyedVectors
model_en = KeyedVectors.load_word2vec_format(model_path_en, binary=False, limit=50000)
model_ro = KeyedVectors.load_word2vec_format(model_path_ro, binary=False, limit=50000)

# -- Code Cell --
embeddings_ro = {}
for sentence in train['romanian']:
    for word in sentence:
        word = word.lower()
        if word in model_ro:
            embeddings_ro[word] = model_ro[word]

# -- Code Cell --
embeddings_en = {}
for sentence in train['english']:
    for word in sentence:
        word = word.lower()
        if word in model_en:
            embeddings_en[word] = model_en[word]

# -- Code Cell --
train

# -- Code Cell --
import numpy as np

# -- Code Cell --
first_embedding = list(embeddings_en.values())[0]
embedding_dim = len(first_embedding)
en_embeded = []
for i in range(len(train)):
    if len(train['english'][i]) != 0:
        embed = [embeddings_en.get(word, np.zeros(embedding_dim)) for word in train['english'][i]]
    else:
        embed = [np.zeros(embedding_dim)]
    embed = np.mean(embed, axis=0)
    en_embeded.append(embed)

# -- Code Cell --
first_embedding = list(embeddings_ro.values())[0]
embedding_dim = len(first_embedding)
ro_embeded = []
for i in range(len(train)):
    if len(train['romanian'][i]) != 0:
        embed = [embeddings_ro.get(word, np.zeros(embedding_dim)) for word in train['romanian'][i]]
    else:
        embed = [np.zeros(embedding_dim)]
    embed = np.mean(embed, axis=0)
    ro_embeded.append(embed)

# -- Code Cell --
from sklearn.linear_model import Ridge
model = Ridge(alpha = 1.0)
model.fit(en_embeded,ro_embeded)

# -- Code Cell --
en_embeddings_test={}
for word in test['Word']:
    if word in model_en:
        en_embeddings_test[word] = model_en[word]

# -- Code Cell --
ro_embeddings_test={}
for word in test['Similar']:
    if word in model_ro:
        ro_embeddings_test[word] = model_ro[word]

# -- Code Cell --
en_embeded_test =[]
for i in range(len(test)):
    word = test['Word'][i]
    embed = en_embeddings_test[word]
    en_embeded_test.append(embed)

# -- Code Cell --
ro_embeded_test =[]
for i in range(len(test)):
    word = test['Similar'][i]
    embed = ro_embeddings_test[word]
    ro_embeded_test.append(embed)

# -- Code Cell --
preds = model.predict(en_embeded_test)
preds

# -- Code Cell --
from sklearn.metrics.pairwise import cosine_similarity
sim = cosine_similarity(preds,ro_embeded_test)
labels = sim.argmax(axis=1)
labels

# -- Code Cell --
words_labels = [test['Similar'][index] for index in labels]

# -- Code Cell --
sub1 = pd.DataFrame({
    'subtaskID':[1],
    'datapointID':[0],
    'answer':[124]
})
sub2 = pd.DataFrame({
    'subtaskID':2,
    'datapointID':test['datapointID'],
    'answer':words_labels
})
final = pd.concat([sub1,sub2])
final.to_csv("subs.csv",index=False)

# -- Code Cell --


# -- Code Cell --


# -- Code Cell --


# -- Code Cell --
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader,Dataset

# -- Code Cell --
def mean_embeddings(sentence, model):
    text = str(sentence).lower().split()
    vector = [model[w] for w in text if w in model]
    if len(vector) == 0:
        return np.zeros(model.vector_size)
    return np.mean(vector,axis=0)

# -- Code Cell --
def mean_embedding(sentence, model):
    words = str(sentence).lower().split()
    vectors = [model[w] for w in words if w in model]
    
    if len(vectors) == 0:
        return np.zeros(model.vector_size)
    
    return np.mean(vectors, axis=0)

# -- Code Cell --
train

# -- Code Cell --
class Datasetcool(Dataset):
    def __init__(self,path):
        self.df = pd.read_csv(path)
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        embeddings_ro = mean_embeddings(row['romanian'],model_ro)
        embeddings_ro = torch.tensor(embeddings_ro,dtype=torch.float32)
        embeddings_en = mean_embeddings(row['english'],model_en)
        embeddings_en = torch.tensor(embeddings_en,dtype=torch.float32)
        return embeddings_en,embeddings_ro

# -- Code Cell --
class Datasett(Dataset):
    def __init__(self,path):
        self.df = pd.read_csv(path)
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        text_ro = mean_embedding(model=model_ro,sentence=row['romanian'])
        text_en = mean_embedding(model=model_en,sentence =row['english'])
        text_ro = torch.tensor(text_ro, dtype=torch.float32)
        text_en = torch.tensor(text_en, dtype=torch.float32)
        return text_ro, text_en

# -- Code Cell --
ds = Datasett("./train_data.csv")
data_loader= DataLoader(ds, batch_size=32, shuffle=True, num_workers=0)

# -- Code Cell --
class Contrast(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.en = nn.Linear(300, 300)
    def forward(self,x):
        en = self.en(x)
        return en

# -- Code Cell --
model = Contrast()
device = 'cuda'
model = model.to(device)
criterion = nn.CosineEmbeddingLoss()
optimier = optim.Adam(model.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.StepLR(optimier,step_size=30,gamma=0.2)

# -- Code Cell --
for epoch in range(50):
    for text_en, text_ro in dt:
        text_en, text_ro = text_en.to(device), text_ro.to(device)
        target=np.ones(text_ro.shape[0])
        target = target.to(device)
        optimizer.zero_grad()
        ro_al_meu = model(text_en)
        loss = criterion(ro_al_meu, text_ro, target)
        loss.backward()
        optimizer.step()
    scheduler.step()

# -- Code Cell --
for epoch in range(100):
    model.train()
    total_loss =0
    for text_ro, text_en in data_loader:
        text_ro = text_ro.to(device)
        text_en = text_en.to(device)
        optimier.zero_grad()
        embedding = model(text_en) 
        target = torch.ones(text_ro.shape[0])
        target = target.to(device)
        loss = criterion(embedding,text_ro,target)
        loss.backward()
        optimier.step()
        total_loss += loss.item()
    scheduler.step()
    print(f"epoch{epoch}, loss {total_loss/len(data_loader)}")

# -- Code Cell --
test

# -- Code Cell --
en_words = test["Word"].tolist()
ro_words = test["Similar"].tolist()
en_vecs = np.array([mean_embedding(w, model_en) for w in en_words], dtype=np.float32)
ro_vecs = np.array([mean_embedding(w, model_ro) for w in ro_words], dtype=np.float32)

# -- Code Cell --
model.eval()
with torch.no_grad():
    en_tensor = torch.tensor(en_vecs).to(device)
    ro_tensor = torch.tensor(ro_vecs).to(device)
    z_en = model(en_tensor).cpu().numpy()
    z_ro = ro_tensor.cpu().numpy()  
z_en = z_en / np.linalg.norm(z_en, axis=1, keepdims=True)
z_ro = z_ro / np.linalg.norm(z_ro, axis=1, keepdims=True)
sim_matrix = z_en @ z_ro.T

# -- Code Cell --


# -- Code Cell --
from scipy.optimize import linear_sum_assignment

cost_matrix = -sim_matrix
row_ind, col_ind = linear_sum_assignment(cost_matrix)
answers = [""] * len(en_words)
for i, j in zip(row_ind, col_ind):
    answers[i] = ro_words[j]

# -- Code Cell --
answers

# -- Code Cell --
sub1 = pd.DataFrame({
    "subtaskID": [1],
    "datapointID": [0],
    "answer": [124]
})
sub2 = pd.DataFrame({
    "subtaskID": 2,
    "datapointID": test["datapointID"],
    "answer": answers
})

submission = pd.concat([sub1, sub2])
submission.to_csv("submission.csv", index=False)

# -- Code Cell --
