import gradio as gr
from transformers import pipeline
import librosa
import numpy as np
import torch

# Load pre-trained models from Hugging Face Hub
# Specify device for GPU support
device = 0 if torch.cuda.is_available() else -1
emotion_model = pipeline("sentiment-analysis", 
                        model="bhadresh-savani/distilbert-base-uncased-emotion",
                        device=device)

# Define preprocessing functions for vocal input data
def preprocess_audio(audio):
    y, sr = librosa.load(audio, sr=16000)
    return y, sr

def extract_features(y, sr):
    features = {
        "mfcc": librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13).mean(axis=1),
        "chroma": librosa.feature.chroma_stft(y=y, sr=sr).mean(axis=1),
        "mel": librosa.feature.melspectrogram(y=y, sr=sr).mean(axis=1),
        "contrast": librosa.feature.spectral_contrast(y=y, sr=sr).mean(axis=1),
        "tonnetz": librosa.feature.tonnetz(y=y, sr=sr).mean(axis=1)
    }
    return features

# Define prediction functions to analyze vocal biomarkers
def analyze_audio(audio):
    try:
        if audio is None:
            return "Please provide an audio input."
            
        y, sr = preprocess_audio(audio)
        features = extract_features(y, sr)
        # Convert features to appropriate input format
        input_data = torch.tensor([features["mfcc"], features["chroma"], 
                                 features["mel"], features["contrast"], 
                                 features["tonnetz"]])
        result = emotion_model(input_data)
        return f"Detected emotion: {result[0]['label']} (Confidence: {result[0]['score']:.2f})"
    except Exception as e:
        return f"Error processing audio: {str(e)}"

# Create Gradio interface with updated syntax
demo = gr.Interface(
    fn=analyze_audio,  # Updated to match the function name
    inputs=gr.Audio(source="microphone", type="filepath"),
    outputs="text",
    title="Vocal Emotion Recognition",
    description="Speak into the microphone to analyze emotional content in your voice.",
    article="""
    This application uses deep learning to analyze emotions in speech.
    Speak clearly into your microphone for best results.
    """,
    examples=[],  # You can add example audio files here
    theme=gr.themes.Default()
)

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")  # Added server_name for container deployment