import cv2
import numpy as np
import math
from hsemotion_onnx.facial_emotions import HSEmotionRecognizer

EMOTION_VA_MAP = {
    'happy':      (0.85,  0.60),
    'sad':        (-0.75, -0.60),
    'angry':      (-0.70,  0.80),
    'fear':       (-0.65,  0.75),
    'surprise':   (0.20,   0.85),
    'disgust':    (-0.80,  0.40),
    'neutral':    (0.00,   0.00),
    'contempt':   (-0.60,  0.50),
}

EMOTION_HANDLE = {
    'happiness': 'happy',
    'sadness': 'sad',
    'anger': 'angry',
    'fear': 'fear',
    'surprise': 'surprise',
    'disgust': 'disgust',
    'neutral': 'neutral',
    'contempt': 'contempt',
}

def get_fine_grained_emotion(valence, arousal):
    radius = math.sqrt(valence**2 + arousal**2)
    angle = math.degrees(math.atan2(arousal, valence))
    
    if radius < 0.25:
        return "neutral"

    if 0 <= angle < 90:
        if angle > 60: return "excited"
        elif angle > 30: return "happy"
        else: return "pleased"
    elif 90 <= angle <= 180:
        if angle > 150: return "nervous"
        elif angle > 120: return "angry"
        else: return "annoying"
    elif -180 <= angle < -90:
        if angle < -150: return "sad"
        elif angle < -120: return "bored"
        else: return "sleepy"
    elif -90 <= angle < 0:
        if angle < -60: return "calm"
        elif angle < -30: return "peaceful"
        else: return "relaxed"
            
    return "neutral"

class EmotionAnalyzer:
    def __init__(self):
        print("正在加载 HSEmotion-ONNX 模型...")
        self.fer = HSEmotionRecognizer(model_name='enet_b0_8_best_vgaf')
        print("HSEmotion-ONNX 模型加载完成")

    def calculate_va_score(self, emotion_prob):
        valence_sum = 0.0
        arousal_sum = 0.0
        total_prob = 0.0
        
        for emotion, prob in emotion_prob.items():
            key = EMOTION_HANDLE.get(emotion.lower(), emotion.lower())
            if key in EMOTION_VA_MAP:
                v, a = EMOTION_VA_MAP[key]
                valence_sum += v * prob
                arousal_sum += a * prob
                total_prob += prob
            
        if total_prob == 0:
            return 0.0, 0.0
        return valence_sum, arousal_sum

    def analyze(self, face_img_bgr):
        if face_img_bgr is None or face_img_bgr.size == 0:
            return []

        face_img_rgb = cv2.cvtColor(face_img_bgr, cv2.COLOR_BGR2RGB)

        # predict_emotions 返回主要情绪标签和概率数组
        emotion_raw, scores = self.fer.predict_emotions(face_img_rgb, logits=False)
        
        probabilities = {}
        for idx, score in enumerate(scores):
            raw_label = self.fer.idx_to_class[idx]
            key = EMOTION_HANDLE.get(raw_label.lower(), raw_label.lower()) 
            
            probabilities[key] = float(score)

        valence, arousal = self.calculate_va_score(probabilities)
        fine_grained_label = get_fine_grained_emotion(valence, arousal)

        result = {
            "box": {}, 
            "vaVal": (round(valence, 4), round(arousal, 4)),
            "probabilities": probabilities,
            "dominant_emotion": EMOTION_HANDLE.get(emotion_raw.lower(), emotion_raw.lower()),
            "emotion": fine_grained_label
        }
        
        return [result]

analyzer_instance = EmotionAnalyzer()

def analyze_emotion_with_hsemotion(face_crop_bgr):
    return analyzer_instance.analyze(face_crop_bgr)