import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
from PIL import Image

# ── Config ───────────────────────────────────────────────────────────────────
IMG_SIZE = 224
CLASS_ORDER = ['no_damage', 'low', 'medium', 'high', 'severe']

device = torch.device('cpu')

# ── Model Definition ──────────────────────────────────────────────────────────
def build_model(num_classes):
    m = models.mobilenet_v2(weights=None)
    in_features = m.classifier[1].in_features
    m.classifier = nn.Sequential(
        nn.Dropout(0.4),
        nn.Linear(in_features, 256),
        nn.ReLU(),
        nn.Dropout(0.3),
        nn.Linear(256, num_classes),
    )
    return m

# Load Model Weights
try:
    ckpt = torch.load('damage_classifier.pth', map_location=device)
    model = build_model(len(CLASS_ORDER))
    model.load_state_dict(ckpt['model_state_dict'])
    model.to(device).eval()
    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

# ── Transform ─────────────────────────────────────────────────────────────────
INFER_TF = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

# ── Prediction Logic ──────────────────────────────────────────────────────────
def predict(image):
    if image is None or model is None:
        return {"Model Error or No Image": 1.0}
    
    # Ensure image is RGB
    image = image.convert('RGB')
    tensor = INFER_TF(image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        logits = model(tensor)
        probs = F.softmax(logits, dim=1).squeeze().tolist()
        
    # Gradio Label component expects a dictionary of {class_name: float_probability}
    prediction_dict = {
        class_name.replace('_', ' ').capitalize(): prob 
        for class_name, prob in zip(CLASS_ORDER, probs)
    }
    
    return prediction_dict

# ── Gradio Interface ──────────────────────────────────────────────────────────
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Drone Image"),
    outputs=gr.Label(num_top_classes=5, label="Predicted Damage Severity"),
    title="🛰️ Disaster Damage Classifier",
    description="**AI-powered drone image damage assessment.** Upload a post-disaster drone image to instantly classify the structural damage severity using our fine-tuned MobileNetV2 model.",
    allow_flagging="never",
    theme=gr.themes.Soft(primary_hue="indigo")
)

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