import time
import torch as tr
import torchvision as tv
import torchvision.transforms as transforms
from model import model, loss_fn, optimizer

def main():
    batch_size = 128
    epochs = 10
    transform = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor()])
    
    trn_dataset = tv.datasets.ImageFolder(
        root='data/tiny-imagenet-200/train',
        transform=transform)
    
    evl_dataset = tv.datasets.ImageFolder(
        root='data/tiny-imagenet-200/val',
        transform=transform)
    
    trn_loader = tr.utils.data.DataLoader(
        trn_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=8,
        pin_memory=True)
    
    evl_loader = tr.utils.data.DataLoader(
        evl_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=8,
        pin_memory=True)
    
    device = tr.device("cuda" if tr.cuda.is_available() else "cpu")
    model.to(device)
    print("Running on:", device)
    start_time = time.time()
    for epoch in range(1, epochs + 1):
        model.train()
        for images, labels in trn_loader:
            images = images.to(device)
            labels = labels.to(device)
            optimizer.zero_grad()
            output = model(images)
            loss = loss_fn(tr.log(tr.clamp(output, min=1e-9)), labels)
            loss.backward()
            optimizer.step()
        model.eval()
        correct = 0
        total = 0

        with tr.no_grad():
            for images, labels in evl_loader:
                images = images.to(device)
                labels = labels.to(device)
                output = model(images)
                preds = tr.argmax(output, dim=1)
                correct += tr.sum(preds == labels).item()
                total += labels.size(0)
        accuracy = correct / total
        print(f"Epoch {epoch}/{epochs} - Loss: {loss.item():.4f} - Accuracy: {accuracy*100:.2f}%")

    end_time = time.time()

    print("Training finished")
    print(f"Final Loss: {loss.item():.4f}")
    print(f"Final Accuracy: {accuracy*100:.2f}%")
    print(f"Total Time: {end_time - start_time:.2f} seconds")


if __name__ == "__main__":
    main()