import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import DataLoader
from NN import OffensiveLanguageClassifier, OffensiveLanguageDataset

# Set the device to use for training
from process_data import train

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


batch_size = 2
vocab_size = 23885
hidden_size = 128
output_size = 3
num_layers = 2
num_epochs = 2

# Create the model and move it to the device
model = OffensiveLanguageClassifier(vocab_size, hidden_size, output_size, num_layers, dropout = 0.3)
model.to(device)

# Define the loss function and the optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())

# Create the DataLoader

train_dataset = OffensiveLanguageDataset(train[0], train["class"])
#print(train_dataset.shape)
#print(train_dataset.head(10))

dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
print(type(dataloader))
# Train the model
for epoch in range(num_epochs):
    #print(dataloader)
    #train_features, train_labels = next(iter(dataloader)
    for data , labels in dataloader:
        #print(data)
        #print(labels)
        #data, labels = data.to(device), labels.to(device)

        # Forward pass
        #print(type(data[0])) 
        data = torch.stack(data)
        logits = model(data)
        loss = loss_fn(logits, labels)

        # Backward pass and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # Print the loss and accuracy at the end of each epoch
    print(f'Epoch {epoch+1}: loss = {loss:.4f}')