import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import io

class EndpointHandler():
    def __init__(self, path=""):
        # 1. Define device
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # 2. Define class names (Matches alphabetical order used in training)
        self.class_names = ['Gray Leaf Spot', 'Healthy']
        
        # 3. Initialize Model Architecture (Update if using EfficientNet)
        # Note: You can make this dynamic or hardcode it to your best model
        self.model = models.resnet50(weights=None)
        self.model.fc = nn.Linear(self.model.fc.in_features, len(self.class_names))
        
        # 4. Load weights (Hugging Face passes the folder path in 'path')
        # Ensure 'model.pth' is the name of your file in the root
        state_dict = torch.load(f"{path}/model.pth", map_location=self.device)
        self.model.load_state_dict(state_dict)
        self.model.to(self.device)
        self.model.eval()

        # 5. Define Preprocessing
        self.transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, data):
        # Data is a dictionary containing the image bytes
        inputs = data.pop("inputs", data)
        
        # Convert bytes to PIL Image
        image = Image.open(io.BytesIO(inputs)).convert("RGB")
        
        # Preprocess
        tensor = self.transform(image).unsqueeze(0).to(self.device)
        
        # Inference
        with torch.no_grad():
            outputs = self.model(tensor)
            probs = torch.nn.functional.softmax(outputs, dim=1)
            conf, pred_idx = torch.max(probs, 1)
        
        # Return formatted result for the widget
        return [
            {"label": self.class_names[pred_idx.item()], "score": conf.item()}
        ]