import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import numpy as np
from torchvision.models import ResNet50_Weights

# -----------------------------
# Device
# -----------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# -----------------------------
# Model definition
# -----------------------------
class FoodIngredientClassifier(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
        num_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(num_features, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        return self.backbone(x)

# -----------------------------
# Load checkpoint
# -----------------------------
checkpoint = torch.load(
    "best_model.pth",
    map_location=device,
    weights_only=False
)

mlb = checkpoint["mlb"]
class_names = mlb.classes_
num_classes = len(class_names)

model = FoodIngredientClassifier(num_classes)
model.load_state_dict(checkpoint["model_state_dict"])
model.to(device)
model.eval()

print(f"Loaded model with {num_classes} ingredient classes")

# -----------------------------
# Image transforms
# -----------------------------
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        [0.485, 0.456, 0.406],
        [0.229, 0.224, 0.225]
    )
])

# -----------------------------
# Utility
# -----------------------------
def clean_name(name):
    return name.replace("_", " ").title()

# -----------------------------
# Prediction function
# -----------------------------
def predict(image, threshold):
    if image is None:
        return {"error": "No image provided"}

    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)

    image = image.convert("RGB")
    input_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        logits = model(input_tensor)
        probs = torch.sigmoid(logits).cpu().numpy()[0]

    # Threshold-based results
    results = {
        clean_name(class_names[i]): float(probs[i])
        for i in range(len(probs))
        if probs[i] >= threshold
    }

    # Fallback: always return top 5
    if not results:
        top_idx = np.argsort(probs)[-5:][::-1]
        results = {
            clean_name(class_names[i]): float(probs[i])
            for i in top_idx
        }

    return dict(sorted(results.items(), key=lambda x: x[1], reverse=True))

# -----------------------------
# Gradio Interface (NO deprecated args)
# -----------------------------
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Upload Food Image"),
        gr.Slider(
            minimum=0.00001,
            maximum=0.5,
            value=0.05,
            step=0.01,
            label="Confidence Threshold"
        )
    ],
    outputs=gr.JSON(label="Detected Ingredients"),
    title="Food Ingredient Detection (Multi-Label)",
    description="Upload a food image to detect multiple ingredients using a ResNet-50 model."
)

# -----------------------------
# Launch (Gradio 6.x style)
# -----------------------------
if __name__ == "__main__":
    iface.launch(theme=gr.themes.Soft())