from model import Transformer
from data import Dataset


import torch
from torch.utils.data import DataLoader
from torch import nn, optim
from torch.utils.tensorboard import SummaryWriter

dataset = Dataset()
train_data, val_data, test_data = dataset.split()

train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
val_loader = DataLoader(val_data, batch_size=32)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
model = Transformer(vocab_size=len(dataset.vocab)).to(device)
optimizer = optim.AdamW(model.parameters(), lr=0.0005, weight_decay=1e-2)
loss_fn = nn.CrossEntropyLoss()

## Let's setup tensorboard

writer = SummaryWriter()

step = 0
for epoch in range(10):
    model.train()
    total_losses = []
    for batch in train_loader:
        inputs = batch[:, :-1].to(device).long()
        targets = batch[:, 1:].to(device).long()

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_fn(outputs.view(-1, outputs.size(-1)), targets.view(-1))
        loss.backward()
        optimizer.step()

        writer.add_scalar("train/loss", loss.item(), step)
        step += 1
        total_losses.append(loss.item())

    avg_loss = sum(total_losses) / len(train_loader)
    writer.add_scalar("train/avg_loss", avg_loss, step)
    print(f"Epoch {epoch+1}, Training Loss: {avg_loss:.4f}")

    model.eval()
    # Save the model
    torch.save(model, f"model_epoch_{epoch+1}.pth")

    val_losses = []
    with torch.no_grad():
        for batch in val_loader:
            inputs = batch[:, :-1].to(device).long()
            targets = batch[:, 1:].to(device).long()
            outputs = model(inputs)
            loss = loss_fn(outputs.view(-1, outputs.size(-1)), targets.view(-1))
            val_losses.append(loss.item())

    avg_val_loss = sum(val_losses) / len(val_losses)
    writer.add_scalar("valid/avg_loss", avg_val_loss, epoch)
    print(f"Epoch {epoch+1}, Validation Loss: {avg_val_loss:.4f}")
