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

# Definir los modelos
modelos = {
    "base": pipeline("automatic-speech-recognition", model="aitor-medrano/lara-base-pushed"),
    "small_1": pipeline("automatic-speech-recognition", model="aitor-medrano/whisper-small-lara"),
    "small_2": pipeline("automatic-speech-recognition", model="aitor-medrano/whisper-small-lara")
}

def greet(grabacion):
    inicio = time.perf_counter()

    sr, y = grabacion
    # Normalizar el array de muestras
    y = (y / np.max(np.abs(y))).astype(np.float32)

    resultados = []
    tiempos = []

    for nombre_modelo, modelo in modelos.items():
        inicio_modelo = time.perf_counter()
        texto_resultado = modelo({"sampling_rate": sr, "raw": y})["text"]
        fin_modelo = time.perf_counter()
        resultados.append(f"{nombre_modelo}: {texto_resultado}")
        tiempos.append(fin_modelo - inicio_modelo)

    tiempo_total = time.perf_counter() - inicio

    return (*resultados, *tiempos, tiempo_total)

# Crear la interfaz mejorada
with gr.Blocks() as demo:
    gr.Markdown("# Evaluación de Modelos de Reconocimiento Automático de Voz")
    
    with gr.Tab("Ingreso de Audio"):
        gr.Audio(label="Graba tu audio", source="microphone", type="numpy", elem_id="audio_input")
    
    with gr.Tab("Resultados y Tiempos"):
        with gr.Column():
            with gr.Row():
                gr.Textbox(label="Resultado Modelo 1", interactive=False, elem_id="resultado1")
                gr.Number(label="Tiempo Modelo 1 (s)", interactive=False, elem_id="tiempo1")
            with gr.Row():
                gr.Textbox(label="Resultado Modelo 2", interactive=False, elem_id="resultado2")
                gr.Number(label="Tiempo Modelo 2 (s)", interactive=False, elem_id="tiempo2")
            with gr.Row():
                gr.Textbox(label="Resultado Modelo 3", interactive=False, elem_id="resultado3")
                gr.Number(label="Tiempo Modelo 3 (s)", interactive=False, elem_id="tiempo3")
            with gr.Row():
                gr.Number(label="Tiempo Total (s)", interactive=False, elem_id="tiempo_total")

    demo.add_component(gr.Button("Procesar", variant="primary", elem_id="process_button"))

    demo.load(greet, inputs="audio_input", outputs=["resultado1", "resultado2", "resultado3", "tiempo1", "tiempo2", "tiempo3", "tiempo_total"], elem_id="process_button")

demo.launch()