import torchaudio
import gradio as gr
import soundfile as sf
import tempfile
import os
import io
import librosa
import numpy as np
import pandas as pd
from transformers import ASTFeatureExtractor, AutoModelForAudioClassification, Trainer, Wav2Vec2FeatureExtractor, HubertForSequenceClassification, pipeline
from datasets import Dataset, DatasetDict
import torch.nn.functional as F
import torch
from collections import Counter
from scipy.stats import kurtosis
from huggingface_hub import InferenceClient
import os

access_token_mod_1 = os.getenv('HF_Access_Personal')

# Cargar el procesador y modelo
processor = ASTFeatureExtractor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593")
model = AutoModelForAudioClassification.from_pretrained("Robertomarting/tmp_trainer",token=access_token_mod_1)

def is_white_noise(audio, threshold=0.75):
    kurt = kurtosis(audio)
    return np.abs(kurt) < 0.1 and np.mean(np.abs(audio)) < threshold

def create_audio_dataframe(audio_tuple, target_sr=16000, target_duration=1.0):
    data = []
    target_length = int(target_sr * target_duration)

    wav_buffer = io.BytesIO()
    sf.write(wav_buffer, audio_tuple[1], audio_tuple[0], format='wav')
    
    wav_buffer.seek(0)
    audio_data, sample_rate = sf.read(wav_buffer)

    audio_data = audio_data.astype(np.float32)
    
    if len(audio_data.shape) > 1:
        audio_data = np.mean(audio_data, axis=1) 

    if sample_rate != target_sr:
        audio_data = librosa.resample(audio_data, orig_sr=sample_rate, target_sr=target_sr)

    audio_data, _ = librosa.effects.trim(audio_data)

    if len(audio_data) > target_length:
        for i in range(0, len(audio_data), target_length):
            segment = audio_data[i:i + target_length]
            if len(segment) == target_length:
                if not is_white_noise(segment):
                    data.append({"audio": segment})
    else:
        if not is_white_noise(audio_data):
            data.append({"audio": audio_data})

    df = pd.DataFrame(data)
    return df

def convert_bytes_to_float64(byte_list):
    return [float(i) for i in byte_list]

def preprocess_function(examples):
    audio_arrays = examples["audio"]
    inputs = processor(
        audio_arrays,
        padding=True,
        sampling_rate=processor.sampling_rate,
        max_length=int(processor.sampling_rate * 1),
        truncation=True,
    )
    return inputs

def predict_audio(audio):
    df = create_audio_dataframe(audio)
    df['audio'] = df['audio'].apply(convert_bytes_to_float64)
    
    # Convertir el dataframe a Dataset
    predict_dataset = Dataset.from_pandas(df)
    dataset = DatasetDict({
        'train': predict_dataset
    })

    if '__index_level_0__' in dataset['train'].column_names:
        dataset['train'] = dataset['train'].remove_columns(['__index_level_0__'])

    encoded_dataset = dataset.map(preprocess_function, remove_columns=["audio"], batched=True)

    # Crear el Trainer para la predicción
    trainer = Trainer(
        model=model,  
        eval_dataset=encoded_dataset["train"]
    )

    # Realizar las predicciones
    predictions_output = trainer.predict(encoded_dataset["train"].with_format("torch"))

    # Obtener las predicciones y etiquetas verdaderas
    predictions = predictions_output.predictions
    labels = predictions_output.label_ids

    # Convertir logits a probabilidades
    probabilities = F.softmax(torch.tensor(predictions), dim=-1).numpy()
    predicted_classes = probabilities.argmax(axis=1)

    # Obtener la etiqueta más común
    most_common_predicted_label = Counter(predicted_classes).most_common(1)[0][0]
    
    # Mapear etiquetas numéricas a etiquetas de texto
    replace_dict = {0: 'Hambre', 1: 'Problemas para respirar', 2: 'Dolor', 3: 'Cansancio/Incomodidad'}
    most_common_predicted_label = replace_dict[most_common_predicted_label]

    return most_common_predicted_label

def clear_audio_input(audio):
    return ""

access_token = os.getenv('HF_ACCESS_TOKEN')

client = InferenceClient("mistralai/Mistral-Nemo-Instruct-2407", token=access_token)

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,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

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'
)

# Función para mostrar la página 1
def mostrar_pagina_1():
    return gr.update(visible=False), gr.update(visible=True)

# Función para mostrar la página 2
def mostrar_pagina_2():
    return gr.update(visible=False), gr.update(visible=True)

# Función para regresar a la pantalla inicial
def redirigir_a_pantalla_inicial():
    return gr.update(visible=True), gr.update(visible=False)

### Monitor

processor = Wav2Vec2FeatureExtractor.from_pretrained("ntu-spml/distilhubert")

monitor_model = HubertForSequenceClassification.from_pretrained("A-POR-LOS-8000/distilhubert-finetuned-cry-detector",token=access_token_mod_1)

pipeline_monitor = pipeline(model=monitor_model,feature_extractor=processor)

def predict_monitor(stream, new_chunk):
    sr, y = new_chunk
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))

    if stream is not None:
        stream = np.concatenate([stream, y])
    else:
        stream = y
    return stream, pipeline_monitor(stream) 
    
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() 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;
            }
            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'>El mejor 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(type="numpy", label="Baby recorder")
                classify_btn = gr.Button("¿Por qué llora?")
                classification_output = gr.Textbox(label="Tu bebé llora por:")

                classify_btn.click(predict_audio, inputs=audio_input, outputs=classification_output)
                audio_input.change(fn=clear_audio_input, 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,
                    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")
        boton_volver_inicio_1.click(redirigir_a_pantalla_inicial, inputs=None, outputs=[pantalla_inicial, pagina_1])

    with gr.Column(visible=False) as pagina_2:
        gr.Markdown("<h2>Monitor</h2>")
        gr.Markdown("# Detección en tiempo real del llanto del bebé con Pipeline")

        # Componente de audio en streaming
        audio_input = gr.Audio(source="microphone", streaming=True, format="wav", label="Habla cerca del micrófono")
        
        # Salida del texto donde se muestra la predicción
        output_text = gr.Textbox(label="Resultado de la predicción")
        
        # Vincular la predicción en streaming con el audio
        audio_input.stream(fn=lambda audio: predict_monitor(audio, audio_classifier),
                           inputs=audio_input,
                           outputs=output_text)

        boton_volver_inicio_2 = gr.Button("Volver a la pantalla inicial")
        boton_volver_inicio_2.click(redirigir_a_pantalla_inicial, inputs=None, outputs=[pantalla_inicial, pagina_2])

    boton_pagina_1.click(mostrar_pagina_1, inputs=None, outputs=[pantalla_inicial, pagina_1])
    boton_pagina_2.click(mostrar_pagina_2, inputs=None, outputs=[pantalla_inicial, pagina_2])

demo.launch()