# %%
import pandas as pd
train = pd.read_csv('train.csv')
test = pd.read_csv('test.csv')

# %%
train.head()

# %%
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from catboost import CatBoostClassifier
from sklearn.metrics import f1_score
X = train['text']
y = train['region']
X_train, X_test, y_train, y_test = train_test_split(X,y,test_size=0.2, random_state=42)
tfidf = TfidfVectorizer(analyzer="char_wb", ngram_range=(1,2),max_features=500000)
X_train_tf = tfidf.fit_transform(X_train)
X_test_tf = tfidf.transform(X_test)
model = LogisticRegression()
model.fit(X_train_tf,y_train)
pred = model.predict(X_test_tf)
acc = f1_score(y_test,pred, average='macro')
acc

# %%
test_text = test['text']
test_test_tf = tfidf.transform(test_text)
X_tf = tfidf.transform(X)
model.fit(X_tf,y)
preds_final_dula = model.predict(test_test_tf)
preds_final_dula

# %%
# unsqueezed = []
# for i in range(len(preds_final)):    
#     for elem in preds_final[i]:
#         unsqueezed.append(elem)
# unsqueezed 

# %% [markdown]
# ## ACUM FAC TASK1 SI 2 CU LSTM

# %%
len(train)

# %%
all_texts = []
for i in range(len(train)):
        all_texts.append(train['text'][i])
for i in range(len(test)):
        all_texts.append(test['text'][i])
all_texts[0]

# %%
word2idx = {"<PAD>":0, "<UNK>":1}
for text in all_texts:
    for word in text.split(' '):
        w = word.lower()
        if w not in word2idx:
            word2idx[w] = len(word2idx)

# %%
word2idx

# %%
index_texts = []
for element in all_texts:
    index_texts.append([word2idx.get(word, 1) for word in element.split(' ')])

# %%
train['region'].unique()

# %%
region_map = {"Transylvania":0, "Moldavia":1, "Bessarabia":2, "Wallachia":3}

# %%
train['region'] = train['region'].map(region_map)

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

# %%
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Embedding(len(word2idx), embedding_dim=128)
        self.l2 = nn.LSTM(128, 64, batch_first=True, bidirectional=True)
        self.l3 = nn.Linear(128, 4)
    def forward(self,x):
        embedded = self.l1(x)
        _,(hidden,_) = self.l2(embedded)
        hidden = torch.cat((hidden[-2], hidden[-1]), dim=1)
        out = self.l3(hidden)
        return out

# %%
device = torch.device('cuda')

# %%
criterion = nn.CrossEntropyLoss()
model = Model()
optimizer = optim.Adam(model.parameters(), lr = 1e-3)
model = model.to(device)

# %%
class TrainDataset(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 = index_texts[index]
        label = region_map[row['region']]
        return text,label

# %%
def collate_fn(batch):
    texts, labels = zip(*batch)

    max_len = max(len(t) for t in texts)

    padded_texts = [
        t + [word2idx["<PAD>"]] * (max_len - len(t))
        for t in texts
    ]

    return torch.tensor(padded_texts, dtype=torch.long), torch.tensor(labels, dtype=torch.long)

# %%
train_ds = TrainDataset("./train.csv")
train_loader = DataLoader(train_ds, batch_size = 32, shuffle=True, num_workers=0,collate_fn=collate_fn)

# %%
train_features, train_labels = next(iter(train_loader))
train_features, train_labels

# %%
for epoch in range(5):
    for text,label in train_loader:
        text = text.to(device)
        label = label.to(device)
        optimizer.zero_grad()
        output = model(text)
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
    print(f"EPOCH{epoch+1} -------- LOSS{loss:04f}")

# %%


# %%


# %%


# %%
test['text']

# %%
all_texts_test = []
for i in range(len(test)):
        all_texts_test.append(test['text'][i])
for i in range(len(test)):
        all_texts_test.append(test['text'][i])
all_texts_test[0]

# %%
index_texts_test = []
for element in all_texts_test:
    index_texts_test.append([word2idx.get(word, 1) for word in element.split(' ')])
index_texts_test

# %%
class TestDataset(Dataset):
    def __init__(self,path):
        self.df = pd.read_csv(path)
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        text = index_texts_test[index]
        return text

# %%
def collate_fn_test(batch):
    max_len = max(len(t) for t in batch)
    padded_texts = [
        t + [word2idx["<PAD>"]] * (max_len - len(t))
        for t in batch
    ]
    return torch.tensor(padded_texts, dtype=torch.long)

# %%
test_ds = TestDataset("./test.csv")
test_loader = DataLoader(test_ds, batch_size = 32, shuffle=True, num_workers=0,collate_fn=collate_fn_test)

# %%
test_features= next(iter(test_loader))
test_features

# %%
preds =[]
preds_final = []
for text in test_loader:
    with torch.no_grad():
        text = text.to(device)
        pred = model(text)
        _,pred = torch.max(pred, 1)
        preds.append(pred.cpu())
        preds_final= torch.cat(preds)

# %%
preds_final = preds_final.numpy()
preds_final

# %%
test['id']

# %%
import ast
values = []
for i in range(len(train)):
    values.append(ast.literal_eval(train['Tags'][i]))

# %%
values

# %%
max_anterior = 0
for i in range(len(train)):
    max_sada = len(values[i])
    if max_sada>max_anterior:
        max_max = max_sada
        max_anterior = max_sada
max_max

# %%
max_anterior = 0
for i in range(len(train)):
    max_sada = len(index_texts[i])
    if max_sada>max_anterior:
        max_max = max_sada
        max_anterior = max_sada
max_max

# %%
for i in range(20):
    print(i)

# %%
input_list_reformat = index_texts
for i in range(len(train)):
    for z in range(max_max-len(index_texts[i])):
        input_list_reformat[i].append(0)

# %%
input_list_reformat

# %%
target_list_reformat = values
for i in range(len(train)):
    for z in range(max_max-len(values[i])):
        target_list_reformat[i].append(-1)

# %%
target_list_reformat[10]

# %%
class Model2(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Embedding(len(word2idx), 128, padding_idx=0)
        self.l2 = nn.LSTM(128, 128, batch_first=True, bidirectional=True, num_layers=2, dropout=0.3)
        self.drop = nn.Dropout(0.3)
        self.l3 = nn.Linear(256, 11)

    def forward(self, x):
        embedded = self.l1(x)
        out, _ = self.l2(embedded)
        out = self.drop(out)
        out = self.l3(out)
        return out

# %%
model2 = Model2()
criterion2 = nn.CrossEntropyLoss(ignore_index=-1)
model2 = model2.to(device)
optimizer2=optim.Adam(model2.parameters(),lr=1e-3)

# %%
class TrainDataset2(Dataset):
    def __init__(self,path):
        self.df = pd.read_csv(path)
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        text = torch.tensor(input_list_reformat[index], dtype=torch.long)
        label = torch.tensor(target_list_reformat[index], dtype=torch.long)
        return text,label

# %%
train_ds2 = TrainDataset2("./train.csv")
train2 = DataLoader(train_ds2, batch_size=32, num_workers=0, shuffle=True)

# %%
for epoch in range(15):
    for text, label in train2:
        text = text.to(device)
        label = label.to(device)

        optimizer2.zero_grad()

        output = model2(text)          
        output = output.view(-1, 11)   
        label = label.view(-1)         

        loss = criterion2(output, label)
        loss.backward()
        optimizer2.step()
    print(f"EPOCH{epoch+1} -------- LOSS{loss:04f}")

# %%
input_list_reformat_test = index_texts_test
for i in range(len(test)):
    for z in range(max_max-len(index_texts_test[i])):
        input_list_reformat_test[i].append(0)

# %%
input_list_reformat_test

# %%
class TestDataset2(Dataset):
    def __init__(self,path):
        self.df = pd.read_csv(path)
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        text =  torch.tensor(input_list_reformat_test[index], dtype=torch.long)
        return text

# %%
test_ds2 = TestDataset2("./test.csv")
test2 = DataLoader(test_ds2, batch_size=32, num_workers=0, shuffle=False)

# %%
preds2 =[]
preds_final2 = []
for text in test2:
    with torch.no_grad():
        text = text.to(device)
        pred = model2(text)
        pred = pred.argmax(dim=2)
        preds2.append(pred.cpu())
        preds_final2= torch.cat(preds2)

# %%
preds_final2 = preds_final2.numpy()
preds_final2

# %%
preds_taiate = []
for i in range(len(preds_final2)):
    lungime_text_real = len(all_texts_test[i].split(' '))
    preds_taiate.append([int(x) for x in preds_final2[i][:lungime_text_real]])

# %%
inv_region_map =  {0:"Transylvania", 1:"Moldavia", 2:"Bessarabia", 3:"Wallachia"}

# %%
preds_final_dula

# %%
subs = pd.DataFrame({
    'subtaskID':1,
    'datapointID':test['id'],
    'answer':preds_final_dula
})
subs2 = pd.DataFrame({
    'subtaskID':2,
    'datapointID':test['id'],
    'answer':preds_taiate
})
pd.concat([subs,subs2]).to_csv('subs.csv',index=False)

# %%


# %%



