import torch
from app.models.model_definitions import create_aadcn_model

from app.config import FACE_MODEL_PATH
from deepface import DeepFace
from PIL import Image
import io
import numpy as np
from torchvision import transforms

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
face_model = create_aadcn_model(num_classes=8)
face_model.load_state_dict(torch.load(FACE_MODEL_PATH, map_location=device))
face_model.eval()

emotion_idx_to_label = {
    0: 'angry',
    1: 'contempt',
    2: 'disgust',
    3: 'fear',
    4: 'happy',
    5: 'neutral',
    6: 'sad',
    7: 'surprise'
}

dataset_to_custom = {
    'angry':      ['Anger', 'Annoyance', 'Disapproval'],
    'contempt':   ['Disapproval', 'Disconnection', 'Annoyance'],
    'disgust':    ['Aversion', 'Disapproval', 'Disconnection'],
    'fear':       ['Fear', 'Disquietment', 'Doubt/Confusion'],
    'happy':      ['Happiness', 'Affection', 'Pleasure', 'Excitement'],
    'neutral':    ['Peace', 'Esteem', 'Confidence'],
    'sad':        ['Sadness', 'Fatigue', 'Suffering'],
    'surprise':   ['Surprise', 'Anticipation', 'Excitement']
}

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

def predict_face_emotion_hybrid(image_bytes):
    # Try DeepFace for "happy" detection first
    try:
        image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
        np_image = np.array(image)  # Ensure numpy array is passed
        result = DeepFace.analyze(
            np_image,
            actions=['emotion'],
            enforce_detection=False,
            detector_backend='opencv'
        )
        if isinstance(result, list):
            result = result[0]
        if result['dominant_emotion'] == 'happy':
            return 'happy', dataset_to_custom['happy']
    except Exception as e:
        print(f"DeepFace error: {e}. Falling back to AA-DCN model.")

    # Fallback to AA-DCN
    image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
    img_t = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        output = face_model(img_t)
        pred_idx = output.argmax(dim=1).item()
        dataset_label = emotion_idx_to_label.get(pred_idx, 'unknown')
        custom_moods = dataset_to_custom.get(dataset_label, ['Unknown'])
    return dataset_label, custom_moods