import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
from transformers import AutoConfig, Wav2Vec2Processor, Wav2Vec2FeatureExtractor
from src.models import Wav2Vec2ForSpeechClassification

import librosa
import IPython.display as ipd
import numpy as np
import pandas as pd
import os

model_name_or_path = "andromeda01111/Malayalam_SA"
config = AutoConfig.from_pretrained(model_name_or_path)
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name_or_path)
sampling_rate = feature_extractor.sampling_rate
model = Wav2Vec2ForSpeechClassification.from_pretrained(model_name_or_path)


def speech_file_to_array_fn(path, sampling_rate):
    speech_array, _sampling_rate = torchaudio.load(path)
    resampler = torchaudio.transforms.Resample(_sampling_rate)
    speech = resampler(speech_array).squeeze().numpy()
    return speech


def predict(path, sampling_rate):
    speech = speech_file_to_array_fn(path, sampling_rate)
    features = feature_extractor(speech, sampling_rate=sampling_rate, return_tensors="pt", padding=True)

    input_values = features.input_values
    attention_mask = features.attention_mask

    with torch.no_grad():
        logits = model(input_values, attention_mask=attention_mask).logits

    scores = F.softmax(logits, dim=1).detach().cpu().numpy()[0]
    output_emotion = [{"Emotion": config.id2label[i], "Score": f"{round(score * 100, 3):.1f}%"} for i, score in enumerate(scores)]

    return output_emotion


# Wrapper function for Gradio
def gradio_predict(audio):
    predictions = predict(audio)
    return [f"{pred['Emotion']}: {pred['Score']}" for pred in predictions]


# Gradio interface
emotions = [config.id2label[i] for i in range(len(config.id2label))]
outputs = [gr.Textbox(label=emotion, interactive=False) for emotion in emotions]

interface = gr.Interface(
    fn=predict,
    inputs=gr.Audio(label="Upload Audio", type="filepath"),
    outputs=outputs,
    title="Emotion Recognition",
    description="Upload an audio file to predict emotions and their corresponding percentages.",
)

# Launch the app
interface.launch()