# -- Markdown Cell --
# Non-poisoned dataset obtained from: https://www.kaggle.com/code/rashedsumon/coronavirus-tweets-nlp-text-classification

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

# -- Code Cell --
df_train = pd.read_csv("train.csv", encoding="latin-1")
df_test = pd.read_csv("test.csv", encoding="latin-1")
df_train.head()

# -- Code Cell --
from wordcloud import WordCloud
import matplotlib.pyplot as plt

sentiments = pd.unique(df_train["Sentiment"])
fig, axes = plt.subplots(1, len(sentiments), figsize=(6 * len(sentiments), 6))

for ax, sentiment in zip(axes, sentiments):
    text = " ".join(df_train[df_train["Sentiment"] == sentiment]["OriginalTweet"].dropna())
    wc = WordCloud(width=800, height=600, background_color="white").generate(text)
    ax.imshow(wc, interpolation="bilinear")
    ax.set_title(sentiment, fontsize=14)
    ax.axis("off")

plt.tight_layout()
plt.show()

# -- Code Cell --
sent_dict= {
    'Extremely Negative':0,
    'Negative':1,
    'Neutral':2,
    'Positive':3,
    'Extremely Positive':4,
}

df_train['Sentiment'] = df_train['Sentiment'].map(sent_dict)


# -- Code Cell --
import re
import nltk
def clean_text(text):
    text = text.lower()
    text = re.sub(r"http\S+", "", text)
    text = re.sub(r"@\S+", "", text)
    text = re.sub(r"[^a-zA-Z]", " ", text)
    text = re.sub(r"nevergiveup", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    
    return text
    
df_train["clean_text"] = df_train["OriginalTweet"].apply(clean_text)
df_test["clean_text"] = df_test["OriginalTweet"].apply(clean_text)


# -- Code Cell --
from collections import Counter

def build_vocab(texts):
    counter = Counter()

    for text in texts:
        words = text.split()
        counter.update(words)


    vocab = {'<PAD>': 0, '<UNK>': 1}

    for word in counter:
            vocab[word] = len(vocab)

    return vocab


# -- Code Cell --
vocab = build_vocab(df_train["clean_text"])

# -- Code Cell --
def text_to_ids(text,vocab):
    words = text.split()
    ids = []
    for word in words:
            ids.append(vocab.get(word,1))
    return ids

# -- Code Cell --
df_train["txt_to_id"] = df_train["clean_text"].apply(lambda x : text_to_ids(x, vocab))
df_test["txt_to_id"] = df_test["clean_text"].apply(lambda x : text_to_ids(x, vocab))


# -- Code Cell --
from torch.utils.data import Dataset, DataLoader
import torch
class data_idk(Dataset):
    def __init__(self, df, test = False):
        super().__init__()
        self.test = test
        if self.test ==False:
            self.y = torch.tensor(np.array(df["Sentiment"]))
        self.X = df["txt_to_id"]
    def __len__(self):
        return len(self.X)
    def __getitem__(self, index):
        if self.test == False:
            return torch.tensor(self.X.iloc[index]), self.y[index]
        return torch.tensor(self.X.iloc[index])

# -- Code Cell --
full_dt = data_idk(df_train)
test_dt = data_idk(df_test, test = True)

# -- Code Cell --
from torch.utils.data import random_split
train_dt, valid_dt = random_split(full_dt, [0.8, 0.2])

# -- 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).long()
    return X, torch.stack(y)
def collate_fn_tst(batch):
    X = batch
    X = pad_sequence(X, batch_first=True).long()
    return X

# -- Code Cell --
train_loader = DataLoader(train_dt, batch_size=32, shuffle= True, collate_fn=collate_fn)
valid_loader = DataLoader(valid_dt, batch_size=32, shuffle= True, collate_fn=collate_fn)
test_loader = DataLoader(test_dt, batch_size= 32, collate_fn=collate_fn_tst )

# -- Code Cell --
import torch.nn as nn

class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, 
                            num_layers=2, bidirectional=True, dropout=0.3)
        
        self.fc = nn.Linear(hidden_dim * 2, 128)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(128, num_classes)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.embedding(x)
        _, (h_n, _) = self.lstm(x)
        
        x = torch.cat((h_n[-2,:,:], h_n[-1,:,:]), dim=1)
        
        x = self.fc(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)

        return x



# -- Code Cell --


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)



vocab_size = len(vocab)
embed_dim = 128
hidden_dim = 128
num_classes = 5

model = LSTMClassifier(vocab_size, embed_dim, hidden_dim, num_classes)
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(
    model.parameters(),
    lr=7e-4
)




# -- Code Cell --
from tqdm.notebook import tqdm
from sklearn.metrics import f1_score

epochs = 10

for epoch in range(epochs):

    model.train()
    total_loss = 0

    for batch in tqdm(train_loader, total=len(train_loader)):
        batch_x = batch[0].to(device)
        batch_y = batch[1].to(device)

        optimizer.zero_grad()
        output = model(batch_x)
        loss = criterion(output, batch_y)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    model.eval()
    all_predicted = []
    all_labels = []

    with torch.no_grad():
        for batch_x, batch_y in valid_loader:
            batch_x = batch_x.to(device)
            batch_y = batch_y.to(device).long()

            outputs = model(batch_x)
            _, predicted = torch.max(outputs.data, 1)

            all_predicted.extend(predicted.cpu().numpy())
            all_labels.extend(batch_y.cpu().numpy())

    f1 = f1_score(all_labels, all_predicted, average='weighted')
    print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(train_loader):.4f} | Test F1: {f1:.4f}")

    model.train()


# -- Code Cell --
model.eval()
all_predicted = []

with torch.no_grad():
    for batch_x in tqdm(test_loader, total=len(test_loader)):
        batch_x = batch_x.to(device)
        outputs = model(batch_x)
        predicted = torch.argmax(outputs, dim=1)

        all_predicted.extend(predicted.cpu().numpy().tolist())


# -- Code Cell --
df_subi = pd.DataFrame({
    "Sentiment" :all_predicted,
    "ID" : df_test["ID"]
})

# -- Code Cell --
inv_dict = {
    0: 'Extremely Negative',
    1:'Negative',
    2: 'Neutral',
    3: 'Positive',
    4:'Extremely Positive',
}

# -- Code Cell --
df_subi["Sentiment"] = df_subi["Sentiment"].map(inv_dict)

# -- Code Cell --
df_subi.to_csv("subi.csv", index = False)

# -- Code Cell --
