import torch
import base64
import os
import io
from PIL import Image
from transformers import AutoImageProcessor, AutoConfig, AutoModelForImageClassification

class EmotionEngine:
    def __init__(self):
        # اسم الموديل مش مهم يظهر في أي حتة تانية
        self.processor = AutoImageProcessor.from_pretrained(
            "trpakov/vit-face-expression"
        )

        config = AutoConfig.from_pretrained(
            "trpakov/vit-face-expression"
        )

        self.model = AutoModelForImageClassification.from_config(config)

        # ... داخل الكلاس __init__
        model_dir = os.path.join(os.path.dirname(__file__), "trained_models")
        model_path = os.path.join(model_dir, "emotion_model.pth")

        state_dict = torch.load(model_path, map_location="cpu")
        self.model.load_state_dict(state_dict)
        self.model.eval()

        self.labels = self.model.config.id2label

    def predict_from_base64(self, base64_img):
        # فك الصورة
        img_bytes = base64.b64decode(base64_img.split(",")[1])
        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")

        # preprocessing
        inputs = self.processor(images=img, return_tensors="pt")

        with torch.no_grad():
            outputs = self.model(**inputs)
            probs = torch.softmax(outputs.logits, dim=1)[0]

        emotion_probs = {
            self.labels[i]: float(probs[i])
            for i in range(len(probs))
        }

        dominant_emotion = max(
            emotion_probs, key=emotion_probs.get
        )

        return emotion_probs, dominant_emotion