# %%
import pandas as pd

# %%
train_df = pd.read_csv('train_data.csv')
test_df = pd.read_csv('test_data.csv')

# %%
train_df.head()

# %%
from nltk.tokenize import word_tokenize
text_train = train_df['comment_text']
tokens=[]
for i in range(len(text_train)):
    tokens.append(word_tokenize(text_train[i].lower()))


# %%
#TEST
from nltk.tokenize import word_tokenize
text_test = test_df['comment_text']
tokens_test=[]
for i in range(len(text_test)):
    tokens_test.append(word_tokenize(text_test[i].lower()))

# %%
# max_len = 0
# max_anterior = 0
# for i in range(len(tokens)):
#     if len(tokens[i])>max_anterior:
#         max_len = len(tokens[i])
#         max_anterior = max_len
# max_len

# %%
# for idx,i in enumerate(tokens):
#     if len(tokens[idx]) == 4849:
#         print(tokens[idx])

# %%
PADDING = 300
for i in range(len(tokens)):
    if len(tokens[i])>PADDING:
        tokens[i] = tokens[i][:PADDING]
    for y in range(PADDING-len(tokens[i])):
        tokens[i].append("<PAD>")

# %%
#TEST
PADDING = 300
for i in range(len(tokens_test)):
    if len(tokens_test[i])>PADDING:
        tokens_test[i] = tokens_test[i][:PADDING]
    for y in range(PADDING-len(tokens_test[i])):
        tokens_test[i].append("<PAD>")

# %%
from collections import Counter

all_words = [w for comment in tokens for w in comment if w != "<PAD>"]
freq = Counter(all_words)
VOCAB_SIZE = 40000
most_common = freq.most_common(VOCAB_SIZE)
word2index = {"<PAD>": 0, "<UNK>": 1}
for i, (word, _) in enumerate(most_common):
    word2index[word] = i + 2
def encode(token_list):
    return [word2index.get(w, 1) for w in token_list]

X_train = [encode(t) for t in tokens]

# %%
#TEST
all_words_test = [w for comment in tokens_test for w in comment if w != "<PAD>"]
def encode(token_list_test):
    return [word2index.get(w, 1) for w in token_list_test]

X_test = [encode(t) for t in tokens_test]

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

# %%
train_df.head()

# %%
labels = train_df[['toxic', 'severe_toxic', 'obscene', 'insult']].values

# %%
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 = torch.tensor(X_train[index], dtype=torch.long)
        label = torch.tensor(labels[index], dtype=torch.float32)
        return text,label

# %%
train_ds=TrainDataset("./train_data.csv")
train = DataLoader(train_ds, shuffle=True, batch_size=64, num_workers=0)

# %%
VOCAB_SIZE = len(word2index)
VOCAB_SIZE

# %%
class BiLSTM(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(VOCAB_SIZE,128)
        self.lstm = nn.LSTM(128,128, bidirectional=True, batch_first=True, num_layers=2, dropout=0.3)
        self.fc= nn.Linear(256, 4)
    def forward(self,x):
        embedding = self.emb(x)
        lstm, _ = self.lstm(embedding)
        lstm = lstm.max(dim=1)[0]
        out = self.fc(lstm)
        return out
        

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

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

# %%
for epoch in range(10):
    for text,label in train:
        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}")

# %%
import numpy as np
from sklearn.metrics import f1_score

# %%
LABELS = ["toxic", "severe_toxic", "obscene", "insult"]

# %%
indices = list(range(len(X_train)))
indices

# %%
class ToxicDataset(Dataset):
    def __init__(self, indices,):
        self.indices = indices
    def __len__(self):
        return len(self.indices)
    def __getitem__(self, i):
        idx = self.indices[i]
        text = torch.tensor(X_train[idx], dtype=torch.long)
        label = torch.tensor(labels[idx], dtype=torch.float32)
        return text, label

# %%
from sklearn.model_selection import train_test_split

train_idx, val_idx = train_test_split(indices, test_size=0.2, random_state=42)
train_loader = DataLoader(ToxicDataset(train_idx), batch_size=64, shuffle=True, num_workers=0)
val_loader = DataLoader(ToxicDataset(val_idx), batch_size=64, shuffle=False, num_workers=0)

# %%
for epoch in range(20):
    model.train()
    for text, label in train_loader:
        text, label = text.to(device), label.to(device)
        optimizer.zero_grad()
        output = model(text)
        loss = criterion(output, label)
        loss.backward()
        optimizer.step()
    model.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for text, label in val_loader:
            text = text.to(device)
            pred = torch.sigmoid(model(text))
            all_preds.append(pred.cpu())
            all_labels.append(label)

    all_preds = torch.cat(all_preds).numpy()
    all_labels = torch.cat(all_labels).numpy()
    for i, name in enumerate(LABELS):
        f1 = f1_score(all_labels[:, i], (all_preds[:, i] > 0.5).astype(int))
        print(f"{name}: F1={f1:.4f}")
    print()

# %%
best_thresholds = []
for i, name in enumerate(LABELS):
    best_t, best_f1 = 0.5, 0
    for t in np.arange(0.1, 0.9, 0.01):
        f1 = f1_score(all_labels[:, i], (all_preds[:, i] > t).astype(int))
        if f1 > best_f1:
            best_f1 = f1
            best_t = t
    best_thresholds.append(best_t)
    print(f"{name}: threshold={best_t:.2f}, F1={best_f1:.4f}")

# %%
best_thresholds

# %%
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 = torch.tensor(X_test[index], dtype=torch.long)
        return text

# %%
test_ds = TestDataset("./test_data.csv")
test = DataLoader(test_ds, batch_size=64, shuffle=False, num_workers=0)

# %%
preds=[]
for text in test:
    with torch.no_grad():
        text = text.to(device)
        pred = model(text)
        pred = torch.sigmoid(pred)
        preds.append(pred.cpu())
        pred_final = torch.cat(preds)
pred_final

# %%
# binary_preds = (pred_final > 0.5).int().cpu().numpy()
# for i in range(len(binary_preds)):
#     binary_preds[i] = [str(elem).replace(" ",",") for elem in binary_preds[i]]
# binary_preds

# %%
binary_preds = np.zeros_like(pred_final.numpy(), dtype=int)
for i in range(4):
    binary_preds[:, i] = (pred_final[:, i].numpy() > best_thresholds[i]).astype(int)

# %%
# rows= []
# for idx,row in test_df.iterrows():
#     rows.append({'subtaskID':1, 'datapointID':row['id'],'answer':binary_preds[idx]})
# pd.DataFrame(rows).to_csv("subs.csv",index=False)

# %%
with open("output.csv", "w") as f:
    f.write("subtaskID,datapointID,answer\n")
    for i, tid in enumerate(test_df['id'].values):
        ans = str(binary_preds[i].tolist()).replace(" ", "")
        f.write(f'1,{tid},"{ans}"\n')

# %%



