import gradio as gr
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image

# ---------------------------------------------------------
# 1. MODEL ARCHITECTURE
# ---------------------------------------------------------
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleCNN, self).__init__()
        
        self.conv_block1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )

        self.conv_block2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=32 * 32 * 32, out_features=128),
            nn.ReLU(),
            nn.Linear(in_features=128, out_features=num_classes),
        )

    def forward(self, x):
        x = self.conv_block1(x)
        x = self.conv_block2(x)
        x = self.classifier(x)
        return x

# ---------------------------------------------------------
# 2. SETUP
# ---------------------------------------------------------
# Initialize model
model = SimpleCNN()

# Load weights (Ensure 'fulldigits.pt' is uploaded to Hugging Face Files!)
try:
    model.load_state_dict(torch.load("fulldigits.pt", map_location="cpu"))
    model.eval()
except FileNotFoundError:
    print("Error: 'fulldigits.pt' not found. Please upload your model file.")

# Define transforms
# CRITICAL FIX: Added lambda to force RGB. 
# This prevents crashes if someone uploads a Grayscale or RGBA image.
transform = transforms.Compose([
    transforms.Lambda(lambda x: x.convert("RGB")), 
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])

# ---------------------------------------------------------
# 3. PREDICTION FUNCTION
# ---------------------------------------------------------
def predict(image):
    if image is None:
        return None

    # Transform image
    img_tensor = transform(image).unsqueeze(0)

    # Make prediction
    with torch.no_grad():
        output = model(img_tensor)
        # Get probabilities
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
    
    # Return a dictionary for Gradio's Label component
    # This creates the nice bar chart effect
    return {str(i): float(probabilities[i]) for i in range(10)}

# ---------------------------------------------------------
# 4. GRADIO INTERFACE
# ---------------------------------------------------------
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Label(num_top_classes=3, label="Predictions"), # Changed to Label for better UI
    title="Digit Classification Project",
    description="Upload an image to check if it contains a digit (0-9).",
    # removed share=True for production deployment
)

if __name__ == "__main__":
    demo.launch()