import torch
from torch import nn, optim
from torchvision import datasets, transforms, models

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

#define how to proprocess and augment images
train_tf = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
])

val_tf = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

#datasets
train_ds = datasets.ImageFolder("litter_detection_split/train", transform=train_tf)
val_ds = datasets.ImageFolder("litter_detection_split/validate",   transform=val_tf)

#data loaders to give data in batches to the model
train_loader = torch.utils.data.DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_ds, batch_size=32)

print("Train samples:", len(train_ds))
print("Val samples:", len(val_ds))

# Load MobileNetV3-Small
model = models.mobilenet_v3_small(pretrained=True)

# set up for binary classification
in_features = model.classifier[3].in_features
model.classifier[3] = nn.Linear(in_features, 2)

model = model.to(device)
print(model)

#deal with loss and optimization (not quite sure what these lines do actually! found them on StackOverflow)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)

num_epochs = 6

for epoch in range(num_epochs):
    # training
    model.train()
    running_loss = 0.0

    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()          #reset gradients
        outputs = model(images)        #forward pass
        loss = criterion(outputs, labels)  #compute loss
        loss.backward()                #backpropagate
        optimizer.step()               #update weights

        running_loss += loss.item() * images.size(0)

    epoch_loss = running_loss / len(train_ds)

    # validate. Find how well it did after this epoch
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            preds = outputs.argmax(dim=1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)

    val_acc = correct / total

    print(f"Epoch {epoch+1}/{num_epochs} - Train loss: {epoch_loss:.4f} - Val acc: {val_acc:.4f}")

torch.save(model.state_dict(), "litter_classifier.pth")
print("Model saved to litter_classifier.pth")