from typing import Dict, List, Any
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch
from PIL import Image
import io

class EndpointHandler:
    def __init__(self):
        # Initialize model and feature extractor
        model_id = "alexionby/ainoai"
        self.model = AutoModelForImageClassification.from_pretrained(model_id)
        self.feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        # Convert bytes to PIL Image
        image = data.pop('image', data)
        image = Image.open(io.BytesIO(image))

        # Preprocess the image
        inputs = self.feature_extractor(images=image, return_tensors="pt")

        # Run the model
        with torch.no_grad():
            outputs = self.model(**inputs)

        # Post-process the model outputs as needed
        logits = outputs.logits
        probabilities = torch.nn.functional.softmax(logits, dim=-1)
        predictions = probabilities.argmax(-1)

        # Convert predictions to JSON-serializable format
        return {"label": str(predictions.item())}


# import torch
# from PIL import Image
# import io

# class EndpointHandler():
#     def __init__(self, path=""):
#         # Preload all the elements you are going to need at inference.
#         # pseudo:
#         self.model= load_model(path)

#     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
#         """
#        data args:
#             inputs (:obj: `str` | `PIL.Image` | `np.array`)
#             kwargs
#       Return:
#             A :obj:`list` | `dict`: will be serialized and returned
#         """

#         # pseudo
#         # self.model(input)