"""
Configuration for different emotion detection models
Add new models here with their specific settings
"""

MODELS_CONFIG = {
    # SuperB Wav2Vec2 - Lightweight, 4 emotions
    "superb/wav2vec2-base-superb-er": {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry"],
        "label_mapping": {
            "neu": "Neutral",
            "neutral": "Neutral",
            "hap": "Happy",
            "happy": "Happy",
            "sad": "Sad",
            "sadness": "Sad",
            "ang": "Angry",
            "angry": "Angry",
            "anger": "Angry"
        },
        "sample_rate": 16000,
        "description": "Lightweight model with 4 basic emotions"
    },
    
    # SuperB HuBERT - Better accuracy, 4 emotions
    "superb/hubert-large-superb-er": {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry"],
        "label_mapping": {
            "neu": "Neutral",
            "neutral": "Neutral",
            "hap": "Happy",
            "happy": "Happy",
            "sad": "Sad",
            "sadness": "Sad",
            "ang": "Angry",
            "angry": "Angry",
            "anger": "Angry"
        },
        "sample_rate": 16000,
        "description": "HuBERT-based model with better accuracy"
    },
    
    # Ehcalabres Wav2Vec2 XLSR - 7 emotions
    "ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition": {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry", "Fear", "Disgust", "Surprise"],
        "label_mapping": {
            "neu": "Neutral",
            "neutral": "Neutral",
            "hap": "Happy",
            "happy": "Happy",
            "happiness": "Happy",
            "sad": "Sad",
            "sadness": "Sad",
            "ang": "Angry",
            "angry": "Angry",
            "anger": "Angry",
            "fea": "Fear",
            "fear": "Fear",
            "dis": "Disgust",
            "disgust": "Disgust",
            "sur": "Surprise",
            "surprise": "Surprise"
        },
        "sample_rate": 16000,
        "description": "Multi-lingual model with 7 emotions"
    },
    
    # Harshit345 XLSR - Alternative model
    "harshit345/xlsr-wav2vec-speech-emotion-recognition": {
        "task": "automatic-speech-recognition",  # Different task type
        "emotions": ["Neutral", "Happy", "Sad", "Angry", "Fear", "Disgust", "Surprise"],
        "label_mapping": {
            "neutral": "Neutral",
            "calm": "Neutral",
            "happy": "Happy",
            "sad": "Sad",
            "angry": "Angry",
            "fearful": "Fear",
            "fear": "Fear",
            "disgust": "Disgust",
            "surprised": "Surprise",
            "surprise": "Surprise"
        },
        "sample_rate": 16000,
        "description": "XLSR-based emotion recognition",
        "special_handling": True  # Needs custom loading
    },
    
    # Amiriparian Wav2Vec2 - RAVDESS dataset
    "amiriparian/wav2vec2-base-ravdess": {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry", "Fear", "Disgust", "Surprise", "Calm"],
        "label_mapping": {
            "01": "Neutral",
            "02": "Calm",
            "03": "Happy",
            "04": "Sad",
            "05": "Angry",
            "06": "Fear",
            "07": "Disgust",
            "08": "Surprise",
            "neutral": "Neutral",
            "calm": "Calm",
            "happy": "Happy",
            "sad": "Sad",
            "angry": "Angry",
            "fearful": "Fear",
            "fear": "Fear",
            "disgust": "Disgust",
            "surprised": "Surprise",
            "surprise": "Surprise"
        },
        "sample_rate": 16000,
        "description": "Trained on RAVDESS dataset with 8 emotions"
    }
}

def get_model_config(model_name):
    """
    Get configuration for a specific model
    
    Args:
        model_name: Name of the model
    
    Returns:
        dict: Model configuration or default config
    """
    if model_name in MODELS_CONFIG:
        return MODELS_CONFIG[model_name]
    
    # Default configuration for unknown models
    return {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry"],
        "label_mapping": {},
        "sample_rate": 16000,
        "description": "Custom model",
        "special_handling": False
    }

def get_available_models():
    """Get list of all available configured models"""
    return list(MODELS_CONFIG.keys())

def get_model_info(model_name):
    """Get human-readable info about a model"""
    config = get_model_config(model_name)
    return {
        "name": model_name,
        "emotions": config["emotions"],
        "num_emotions": len(config["emotions"]),
        "description": config["description"],
        "sample_rate": config["sample_rate"]
    }