import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download

# 1. Define the architecture EXACTLY as it was in training notebook
class VehicleClassifier(nn.Module):
    def __init__(self):
        super(VehicleClassifier, self).__init__()
        # Block 1: 3 -> 16 channels. Input 224x224 -> Output 112x112 (after pool)
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        # Block 2: 16 -> 32 channels. Output 56x56 (after pool)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        # Block 3: 32 -> 64 channels. Output 28x28 (after pool)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(64)
        self.pool = nn.MaxPool2d(2, 2)
        # Fully Connected Layers
        self.fc1 = nn.Linear(64 * 28 * 28, 512)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(512, 8) # 8 Classes

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        x = x.view(-1, 64 * 28 * 28) # Flatten
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 2. Setup Device and Load Model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = VehicleClassifier().to(device)

# Download and Load Weights
repo_id = "abhiprd20/vehicle_classification_model-utd"
checkpoint_path = hf_hub_download(repo_id=repo_id, filename="model.pth")
model.load_state_dict(torch.load(checkpoint_path, map_location=device))
model.eval()

# 3. Prediction Logic
classes = ['Bicycle', 'Bus', 'Car', 'Motorcycle', 'NonVehicles', 'Taxi', 'Truck', 'Van']
data_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def predict(img):
    if img is None: return None
    img_tensor = data_transforms(img).unsqueeze(0).to(device)
    with torch.no_grad():
        outputs = model(img_tensor)
        probs = torch.nn.functional.softmax(outputs[0], dim=0)
        confidences = {classes[i]: float(probs[i]) for i in range(8)}
    return confidences

# 4. Interface
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="🏎️ Autonomous Vehicle Perception Demo",
    description="Custom CNN Model classifying type of vehicles."
).launch()