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

# -- Code Cell --
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
from nltk.corpus import stopwords
stop_words = set(stopwords.words("english"))
lemmatizer = WordNetLemmatizer()
def preprocess_text(text):
    tokens = word_tokenize(text)
    tokens = [lemmatizer.lemmatize(w) for w in tokens if w not in stop_words]
    return " ".join(tokens)
def simple_tokenizer(text):
    return text.split()

# -- Code Cell --
train_text = train['comment_text'].apply(preprocess_text)

# -- Code Cell --
train_text = train_text.apply(simple_tokenizer)

# -- Code Cell --
train_text[0]

# -- Code Cell --
test_text = test['comment_text'].apply(preprocess_text)
test_text = test_text.apply(simple_tokenizer)
test_text

# -- Code Cell --
from gensim.models import FastText
ft = FastText(
    sentences=pd.concat([train_text,test_text]),
    vector_size=300,
    window = 5
)

# -- Code Cell --
import torch
def get_embeddings(text):
    embeddings = []
    for x in text:
        if x in ft.wv:
            embeddings.append(ft.wv[x])
    return torch.tensor(embeddings)

# -- Code Cell --
train_embs = [get_embeddings(tokens) for tokens in train_text]

# -- Code Cell --
labels = torch.tensor(train[['toxic','severe_toxic','obscene','insult']].values,dtype=torch.float32)

# -- Code Cell --
from torch.utils.data import Dataset,DataLoader
class FullDs(Dataset):
    def __init__(self, embs, labels):
        self.embs = embs
        self.labels = labels

    def __len__(self):
        return len(self.embs)

    def __getitem__(self, index):
        return self.embs[index], self.labels[index]

# -- Code Cell --
full_ds = FullDs(train_embs, labels)

# -- Code Cell --
from torch.utils.data import random_split
train_size = int(0.8*len(full_ds))
val_size = len(full_ds) - train_size
train_subset, val_subset=random_split(full_ds,[train_size,val_size])

# -- Code Cell --
from torch.nn.utils.rnn import pad_sequence
def collate_fn(batch):
    X, y = zip(*batch)
    X = pad_sequence(X, padding_value=0.0, batch_first=True)
    return X,torch.stack(y)
def collate_fn_pt_test(batch):
    X = pad_sequence(batch, padding_value=0.0, batch_first=True)
    return X

# -- Code Cell --
train_loader = DataLoader(train_subset, batch_size=32, num_workers=0, shuffle=True,collate_fn=collate_fn)
val_loader = DataLoader(val_subset, batch_size=32, num_workers=0, shuffle=False,collate_fn=collate_fn)

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

# -- Code Cell --
class Modelus(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(input_size=300,hidden_size=128,num_layers=2,dropout=0.3,bidirectional=True,batch_first=True)
        self.out = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.BatchNorm1d(128),
            nn.Dropout(0.3),
            nn.Linear(128, 4)
        )
    def forward(self, x):
        _, (h_n, _) = self.lstm(x)
        last = torch.cat((h_n[-2], h_n[-1]), dim=1)   
        out = self.out(last)                           
        return out

# -- Code Cell --
import torch.optim as optim
model = Modelus()
model = model.to(torch.device("cuda"))
optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.BCEWithLogitsLoss()
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.15)

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

# -- Code Cell --
import numpy as np

# -- Code Cell --
for epoch in range(15):
    model.train()
    total_loss = 0
    for img,label in train_loader:
        img = img.to(torch.device("cuda"))
        label = label.to(torch.device("cuda"))
        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
        total_loss+=loss.item()
    scheduler.step()
    best_f1 = -1
    best_thresholds = None

    model.eval()
    all_probs = []
    all_labels = []
    with torch.no_grad():
        for img, label in val_loader:
            img = img.to(torch.device("cuda"))
            output = model(img)
            probs = torch.sigmoid(output)

            all_probs.append(probs.cpu())
            all_labels.append(label.cpu())
    all_probs = torch.cat(all_probs).numpy()      
    all_labels = torch.cat(all_labels).numpy()   
    best_thresholds = [0.5] * 4
    for cls in range(4):
        best_cls_f1 = -1
        best_cls_threshold = 0.5

        for k in range(11):
            threshold = k / 10
            preds_cls = (all_probs[:, cls] > threshold).astype(int)
            f1 = f1_score(all_labels[:, cls], preds_cls)

            if f1 > best_cls_f1:
                best_cls_f1 = f1
                best_cls_threshold = threshold

        best_thresholds[cls] = best_cls_threshold
    final_preds = np.zeros_like(all_probs, dtype=int)
    for cls in range(4):
        final_preds[:, cls] = (all_probs[:, cls] > best_thresholds[cls]).astype(int)
    f1 = f1_score(all_labels, final_preds, average='macro')
    print(f"epoch {epoch+1}, loss {total_loss/len(train_loader):04f}, f1 {f1}")
print("best thresholds:", best_thresholds)

# -- Code Cell --
best_thresholds

# -- Code Cell --
class TestDS(Dataset):
    def __init__(self, embs):
        self.embs = embs
    def __len__(self):
        return len(self.embs)
    def __getitem__(self, index):
        return self.embs[index]

# -- Code Cell --
test_embs = [get_embeddings(tokens) for tokens in test_text]

# -- Code Cell --
test_ds = TestDS(test_embs)
test_loader = DataLoader(test_ds, batch_size=32, num_workers=0, shuffle=False,collate_fn=collate_fn_pt_test)

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

with torch.no_grad():
    for img in test_loader:
        img = img.to(torch.device("cuda"))
        probs = torch.sigmoid(model(img))
        pred = (probs > torch.tensor(best_thresholds, device=torch.device("cuda"))).int()
        all_preds.append(pred.cpu())
preds_final = torch.cat(all_preds, dim=0).numpy()
preds_final

# -- Code Cell --
pred_final_final = []
for i in range(len(preds_final)):
    pred_final_final.append(str(preds_final[i].tolist()))
pred_final_final

# -- Code Cell --
subs = pd.DataFrame({
    'subtaskID':1,
    'datapointID':test['id'],
    'answer':pred_final_final
}).to_csv('subs.csv',index=False)

# -- Code Cell --
