# -- Code Cell --
import pandas as pd
train = pd.read_csv("train.csv")
test = pd.read_csv("test.csv")
test = test.drop(columns = ['Sentiment'])

# -- Code Cell --
mapp = {}
rev_map = {}
for idx,key in enumerate(train['Sentiment'].unique().tolist()):
    mapp[key] = idx
    rev_map[idx] = key

# -- Code Cell --
import re
from nltk.stem import WordNetLemmatizer
from nltk.tokenize import word_tokenize
lemmatizer = WordNetLemmatizer()

def preprocess_text(text):
    text = text.lower()
    text = re.sub("nevergiveup"," ",text)
    text = re.sub(r"\s+", " ", text).strip()
    text= re.sub(r"https?://\S+", "", text)
    tokens = word_tokenize(text)
    tokens = [lemmatizer.lemmatize(w) for w in tokens]
    return " ".join(tokens)

# -- Code Cell --
import re
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
lemmatizer = WordNetLemmatizer()
def preproc(text):
    text = text.lower()
    text = re.sub(r"\s+", " ", text).strip()
    text = re.sub(r"https?://\S+", "", text)
    
    

# -- Code Cell --
text_test = test['OriginalTweet']
len(text_test)

# -- Code Cell --
train_text = train['OriginalTweet']
len(train_text)

# -- Code Cell --
text_all = pd.concat([text_test, train_text]).reset_index(drop=True)
len(text_all)

# -- Code Cell --
text_all = text_all.apply(preprocess_text)
text_all = text_all.fillna("")

# -- Code Cell --
text_test = text_all[:3798]
text_train = text_all[3798:]

# -- Code Cell --
word2idx = {"<PAD>":0, "<UNK>":1}
for sentence in text_all:
    sentence = str(sentence)
    for word in sentence.split():
        if word not in word2idx:
            word2idx[word] = len(word2idx)
word2idx

# -- Code Cell --
indicii_magici = []
for sentence in text_train.str.split():
    values_temp = []
    for z in sentence:
        values_temp.append(word2idx[z])
    indicii_magici.append(values_temp)


indicii_magici_test = []
for sentence in text_test.str.split():
    values_temp = []
    for z in sentence:
        values_temp.append(word2idx[z])
    indicii_magici_test.append(values_temp)

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

# -- Code Cell --
from torch.nn.utils.rnn import pad_sequence
def collate_fn(batch):
    X,y = zip(*batch)
    X = pad_sequence(X, batch_first=True,padding_value=0)
    y = torch.tensor(y,dtype=torch.long)
    return X,y
def collate_fn_test(batch):
    X = pad_sequence(batch, batch_first=True,padding_value=0)
    return X

# -- Code Cell --
class TrainDataset(Dataset):
    def __init__(self,df):
        self.df = df
    def __len__(self):
        return len(self.df)
    def __getitem__(self,index):
        row = self.df.iloc[index]
        indices = indicii_magici[index]
        label = mapp[row['Sentiment']]
        indices = torch.tensor(indices,dtype=torch.long)
        return indices,label

# -- Code Cell --
train_ds = TrainDataset(train)
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers = 0,collate_fn=collate_fn)

# -- Code Cell --
class modelulMeu(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.emb = nn.Embedding(len(word2idx),128,padding_idx=0)
        self.lstm = nn.LSTM(128, 128, num_layers=2, bidirectional=True, dropout =0.2, batch_first = True)
        self.clsf = nn.Linear(128*2, 5)
        
    def forward(self,x):
        emb = self.emb(x)
        output, (h_n, c_n) = self.lstm(emb)
        out = self.clsf(output[:, -1, :])
        return out

# -- Code Cell --
device ='cuda'
model = modelulMeu()
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.2)

# -- Code Cell --
from tqdm import tqdm
for epoch in range(30):
    model.train()
    total_loss = 0
    for text,label in tqdm(train_loader):
        text = text.to(device)
        label = label.to(device)
        optimizer.zero_grad()
        output = model(text)
        loss = criterion(output, label)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    scheduler.step()
    print(f"epoch{epoch+1}, loss{total_loss/len(train_loader)}")

# -- Code Cell --
class TestDataset(Dataset):
    def __init__(self,df):
        self.df = df
    def __len__(self):
        return len(self.df)
    def __getitem__(self,index):
        indices = indicii_magici_test[index]
        indices = torch.tensor(indices,dtype=torch.long)
        return indices

# -- Code Cell --
test_ds = TestDataset(test)
test_loader = DataLoader(test_ds, batch_size=32, shuffle=False, num_workers = 0,collate_fn=collate_fn_test)

# -- Code Cell --
preds= []
with torch.no_grad():
    for text in test_loader:
        text = text.to(device)
        output = model(text)
        pred = output.argmax(dim=1)
        preds.append(pred.cpu())
preds = torch.cat(preds).numpy()
preds

# -- Code Cell --
pred_final = []
for pred in preds:
    pred_final.append(rev_map[pred])

# -- Code Cell --
asndksajdskajdkslajdklsajdsakldjsalkdjlaks = pd.DataFrame({
    'ID':test['ID'],
    'Sentiment':pred_final
}).to_csv("subi.csv",index=False)

# -- Code Cell --
!py score.py

# -- Code Cell --


# -- Code Cell --
