import os
import torch
import gradio as gr
from huggingface_hub import InferenceClient
from model import model_params, AudioDataset

token = os.getenv("HF_TOKEN")
dataset_path = f"A-POR-LOS-8000/baby_cry_detection" # PARA MONITOR
# dataset_path = f"data/mixed_data_nuevo" # PARA CLASIFICADOR
model, _, _, id2label = model_params(dataset_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# Usar a GPU o CPU
model.to(device)# Usar a GPU o CPU
client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct", token=token)
# client = InferenceClient("mistralai/Mistral-Nemo-Instruct-2407", token=token)

def predict(audio_path):
    audio_dataset = AudioDataset(dataset_path, {})
    inputs = audio_dataset.preprocess_audio(audio_path)
    inputs = {"input_values": inputs.to(device).unsqueeze(0)}
    with torch.no_grad():
        outputs = model(**inputs)
        predicted_class_ids = outputs.logits.argmax(-1)
        label = id2label[predicted_class_ids.item()]
    return label

def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
    messages = [{"role": "system", "content": system_message}]
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    messages.append({"role": "user", "content": message})
    response = ""
    for message in client.chat_completion(messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p): # Creo que lo importante para el modelo
        token = message.choices[0].delta.content
        response += token
        yield response

def cambiar_pestaña():
    return gr.update(visible=False), gr.update(visible=True)

my_theme = gr.themes.Soft(
    primary_hue="emerald",
    secondary_hue="green",
    neutral_hue="slate",
    text_size="sm",
    spacing_size="sm",
    font=[gr.themes.GoogleFont('Nunito'), 'ui-sans-serif', 'system-ui', 'sans-serif'],
    font_mono=[gr.themes.GoogleFont('Nunito'), 'ui-monospace', 'Consolas', 'monospace'],
    ).set(
    body_background_fill='*neutral_50',
    body_text_color='*neutral_600',
    body_text_size='*text_sm',
    embed_radius='*radius_md',
    shadow_drop='*shadow_spread',
    shadow_spread='*button_shadow_active'
    )

with gr.Blocks(theme=my_theme) as demo:
    with gr.Column(visible=True, elem_id="pantalla-inicial") as pantalla_inicial:
        gr.HTML(
            """
            <style>
            @import url('https://fonts.googleapis.com/css2?family=Lobster&display=swap');
            @import url('https://fonts.googleapis.com/css2?family=Roboto&display=swap');
            
            h1 {
                font-family: 'Lobster', cursive;
                font-size: 5em !important;
                text-align: center;
                margin: 0;
            }
            
            .gr-button {
                background-color: #4CAF50 !important; 
                color: white !important; 
                border: none;
                padding: 15px 32px;
                text-align: center;
                text-decoration: none;
                display: inline-block;
                font-size: 16px;
                margin: 4px 2px;
                cursor: pointer;
                border-radius: 12px;
            }
            
            .gr-button:hover {
                background-color: #45a049; 
            }
            h2 {
                font-family: 'Lobster', cursive;
                font-size: 3em !important;
                text-align: center;
                margin: 0;
            }
            p.slogan, h4, p, h3 {
                font-family: 'Roboto', sans-serif;
                text-align: center;
            }
            </style>
            <h1>Iremia</h1>
            <h4 style='text-align: center; font-size: 1.5em'>Tu aliado para el bienestar de tu bebé</h4>
            """
        )
        gr.Markdown(
            "<h4 style='text-align: left; font-size: 1.5em;'>¿Qué es Iremia?</h4>"
        )
        gr.Markdown(
            "<p style='text-align: left'>Iremia es un proyecto llevado a cabo por un grupo de estudiantes interesados en el desarrollo de modelos de inteligencia artificial, enfocados específicamente en casos de uso relevantes para ayudar a cuidar a los más pequeños de la casa.</p>"
        )
        gr.Markdown(
            "<h4 style='text-align: left; font-size: 1.5em;'>Nuestra misión</h4>"
        )
        gr.Markdown(
            "<p style='text-align: left'>Sabemos que la paternidad puede suponer un gran desafío. Nuestra misión es brindarles a todos los padres unas herramientas de última tecnología que los ayuden a navegar esos primeros meses de vida tan cruciales en el desarrollo de sus pequeños.</p>"
        )
        gr.Markdown(
            "<h4 style='text-align: left; font-size: 1.5em;'>¿Qué ofrece Iremia?</h4>"
        )
        gr.Markdown(
            "<p style='text-align: left'>Iremia ofrece dos funcionalidades muy interesantes:</p>"
        )
        gr.Markdown(
            "<p style='text-align: left'>Predictor: Con nuestro modelo de inteligencia artificial, somos capaces de predecir por qué tu hijo de menos de 2 años está llorando. Además, tendrás acceso a un asistente personal para consultar cualquier duda que tengas sobre el cuidado de tu pequeño.</p>"
        )
        gr.Markdown(
            "<p style='text-align: left'>Monitor: Nuestro monitor no es como otros que hay en el mercado, ya que es capaz de reconocer si un sonido es un llanto del bebé o no, y si está llorando, predice automáticamente la causa, lo cual te brindará la tranquilidad de saber siempre qué pasa con tu pequeño y te ahorrará tiempo y muchas horas de sueño.</p>"
        )
        with gr.Row():
            with gr.Column():
                gr.Markdown("<h2>Predictor</h2>")
                boton_pagina_1 = gr.Button("Prueba el predictor")
                gr.Markdown("<p>Descubre por qué llora tu bebé y resuelve dudas sobre su cuidado con nuestro Iremia assistant</p>")
            with gr.Column():
                gr.Markdown("<h2>Monitor</h2>")
                boton_pagina_2 = gr.Button("Prueba el monitor")
                gr.Markdown("<p>Un monitor inteligente que detecta si tu hijo está llorando y te indica el motivo antes de que puedas levantarte del sofá</p>")
    with gr.Column(visible=False) as pagina_1:
        with gr.Row():
            with gr.Column():
                gr.Markdown("<h2>Predictor</h2>")
                audio_input = gr.Audio(
                    min_length=1.0,
                    format="wav",
                    label="Baby recorder",
                    type="filepath", # Para no usar numpy y preprocesar siempre igual
                    )
                classify_btn = gr.Button("¿Por qué llora?")
                classification_output = gr.Textbox(label="Tu bebé llora por:")
                classify_btn.click(predict, inputs=audio_input, outputs=classification_output)
            with gr.Column():
                gr.Markdown("<h2>Assistant</h2>")
                system_message = "You are a Chatbot specialized in baby health and care."
                max_tokens = 512
                temperature = 0.7
                top_p = 0.95
                chatbot = gr.ChatInterface(
                    respond, # TODO: Cambiar para que argumentos estén aquí metidos
                    additional_inputs=[
                        gr.State(value=system_message),
                        gr.State(value=max_tokens),
                        gr.State(value=temperature),
                        gr.State(value=top_p)
                    ],
                )
                gr.Markdown("Este chatbot no sustituye a un profesional de la salud. Ante cualquier preocupación o duda, consulta con tu pediatra.")
                boton_volver_inicio_1 = gr.Button("Volver a la pantalla inicial").click(cambiar_pestaña, outputs=[pagina_1, pantalla_inicial])
    with gr.Column(visible=False) as pagina_2:
        gr.Markdown("<h2>Monitor</h2>")
        gr.Markdown("Contenido de la Página 2")
        boton_volver_inicio_2 = gr.Button("Volver a la pantalla inicial").click(cambiar_pestaña, outputs=[pagina_2, pantalla_inicial])
    boton_pagina_1.click(cambiar_pestaña, outputs=[pantalla_inicial, pagina_1])
    boton_pagina_2.click(cambiar_pestaña, outputs=[pantalla_inicial, pagina_2])
demo.launch()