# -- Code Cell --
import pandas as pd
import numpy as np

# -- Code Cell --
train = pd.read_csv("./train_data (1).csv")

# -- Code Cell --
train

# -- Code Cell --
train_en = train['english']
train_ro = train['romanian']

# -- Code Cell --
from gensim.models import KeyedVectors
model_path_en = './wiki.en.small.vec'
model_path_ro = './wiki.ro.small.vec'

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 --
type(train_ro[1])

# -- Code Cell --
embeddings_ro = {}
for text in train_ro:
    text = str(text)
    for word in text.split():
        if word in model_ro:
            embeddings_ro[word] = model_ro[word]

# -- Code Cell --
embeddings_ro

# -- Code Cell --
embeddings_en = {}
for text in train_en:
    text = str(text)
    for word in text.split():
        if word in model_en:
            embeddings_en[word] = model_en[word]

# -- Code Cell --
embeddings_en

# -- 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 --
X = []
Y = []
for _, row in train.iterrows():
    en_sent = row["english"]
    ro_sent = row["romanian"]
    X.append(mean_embedding(en_sent, model_en))
    Y.append(mean_embedding(ro_sent, model_ro))
X = np.array(X)
Y = np.array(Y)

# -- Code Cell --
from sklearn.linear_model import Ridge

reg = Ridge(alpha=1.0)
reg.fit(X, Y)

# -- Code Cell --
test = pd.read_csv("./test_data.csv")

# -- Code Cell --
from sklearn.metrics.pairwise import cosine_similarity

# -- Code Cell --
ro_words = test["Similar"].tolist()
ro_vecs = np.array([mean_embedding(w, model_ro) for w in ro_words])
answers = []
for word in test["Word"]:
    v_en = mean_embedding(word, model_en).reshape(1, -1)
    v_proj = reg.predict(v_en)
    sims = cosine_similarity(v_proj, ro_vecs)[0]
    best_idx = np.argmax(sims)
    answers.append(ro_words[best_idx])

# -- Code Cell --
answers

# -- Code Cell --
en_set = set(test['Word'])
ro_set = set(test['Similar'])
intersection = en_set.intersection(ro_set)
subtask_one_ans = len(list(intersection))
print (f'Subtask 1 answer: {subtask_one_ans}.')

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

# -- Code Cell --
sub1

# -- Code Cell --
finasadksajdkasjkd = pd.concat([sub1,sub2]).to_csv("subs.csv",index=False)

# -- 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 --
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 (1).csv")
data_loader = DataLoader(ds, batch_size=32, shuffle=True, num_workers=0)

# -- Code Cell --
model_ro.vector_size

# -- 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(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)
        # print(text_ro.shape)
        # print(text_en.shape)
        # print(embedding.shape)
        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 --
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)
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 --
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 --
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 --
