import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import random

# Oración de ejemplo y palabras posibles con sus probabilidades originales
oracion_incompleta = "El astronauta miró hacia las estrellas y sintió una profunda"
palabras_posibles = ["emoción", "nostalgia", "curiosidad", "admiración", "soledad"]
probabilidades_originales = np.array([0.35, 0.25, 0.20, 0.15, 0.05])  # Suma = 1

def aplicar_temperatura(probs, temperatura):
    """Aplica la transformación de temperatura a las probabilidades."""
    if temperatura == 0:
        # Cuando la temperatura es 0, simplemente devuelve una distribución one-hot
        # para la palabra más probable
        resultado = np.zeros_like(probs)
        resultado[np.argmax(probs)] = 1.0
        return resultado
    
    # Aplicar la fórmula de temperatura: p_i = exp(log(p_i)/T) / suma(exp(log(p_j)/T))
    log_probs = np.log(probs)
    exp_probs = np.exp(log_probs / temperatura)
    return exp_probs / np.sum(exp_probs)

def crear_grafico(temperatura):
    """Crea un gráfico de barras con las probabilidades ajustadas por temperatura."""
    # Aplicar temperatura a las probabilidades
    probs_ajustadas = aplicar_temperatura(probabilidades_originales, temperatura)
    
    # Crear figura
    fig, ax = plt.subplots(figsize=(10, 6))
    
    # Colores para las barras
    colores = ['#FF6B6B', '#4ECDC4', '#FFD166', '#6B5B95', '#88D8B0']
    
    # Crear gráfico de barras
    barras = ax.bar(palabras_posibles, probs_ajustadas, color=colores)
    
    # Añadir valores en la parte superior de cada barra
    for barra in barras:
        altura = barra.get_height()
        ax.annotate(f'{altura:.3f}',
                    xy=(barra.get_x() + barra.get_width() / 2, altura),
                    xytext=(0, 3),  # 3 puntos de desplazamiento vertical
                    textcoords="offset points",
                    ha='center', va='bottom')
    
    # Configurar el gráfico
    ax.set_title(f'Probabilidades con temperatura = {temperatura:.2f}', fontsize=15)
    ax.set_ylabel('Probabilidad', fontsize=12)
    ax.set_ylim(0, 1.0)
    ax.grid(axis='y', linestyle='--', alpha=0.7)
    
    # Rotar etiquetas del eje x para mejor legibilidad
    plt.xticks(rotation=15)
    
    return fig

def muestrear_palabra(probs):
    """Muestrea una palabra basada en sus probabilidades."""
    return np.random.choice(palabras_posibles, p=probs)

def actualizar_interfaz(temperatura):
    """Actualiza la interfaz con el nuevo valor de temperatura."""
    # Crear gráfico actualizado
    grafico = crear_grafico(temperatura)
    
    # Actualizar probabilidades
    probs_ajustadas = aplicar_temperatura(probabilidades_originales, temperatura)
    
    return grafico, probs_ajustadas

def completar_oracion(probs):
    """Completa la oración muestreando una palabra según las probabilidades."""
    palabra_seleccionada = muestrear_palabra(probs)
    
    # Crear texto HTML para colorear la palabra seleccionada
    oracion_completa = f"{oracion_incompleta} <span style='color:#FF4500; font-weight:bold;'>{palabra_seleccionada}</span>."
    
    return oracion_completa

# Función principal que maneja la interacción en la interfaz
def interfaz_temperatura(temperatura=1.0):
    grafico, probs = actualizar_interfaz(temperatura)
    return grafico, probs

# Función para manejar el botón de muestreo
def muestrear_boton(probs):
    return completar_oracion(probs)

# Crear la interfaz Gradio
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("# 🔥 Demostración de Temperatura en LLMs")
    gr.Markdown("""
    Esta aplicación demuestra cómo la temperatura afecta a las probabilidades de selección en los modelos de lenguaje.
    
    **Temperatura baja** (cercana a 0): El modelo se vuelve más determinista, eligiendo la palabra más probable.
    
    **Temperatura alta** (mayor que 1): El modelo se vuelve más aleatorio, dando más oportunidad a palabras menos probables.
    """)
    
    # Mostrar la oración incompleta
    gr.Markdown(f"## Oración: *{oracion_incompleta}...*")
    
    # Variables de estado ocultas
    probs_state = gr.State([])
    
    with gr.Row():
        with gr.Column(scale=2):
            # Gráfico de probabilidades
            output_plot = gr.Plot(label="Distribución de Probabilidades")
            
            # Control deslizante para la temperatura
            temp_slider = gr.Slider(
                minimum=0.1, 
                maximum=2.0, 
                value=1.0, 
                step=0.1, 
                label="Temperatura",
                info="Desliza para ajustar la temperatura"
            )
        
        with gr.Column(scale=1):
            # Botón para muestrear
            sample_button = gr.Button("Muestrear Palabra", variant="primary")
            
            # Mostrar oración completa
            output_text = gr.HTML(label="Oración Completada")
    
    # Configurar eventos y funciones
    temp_slider.change(
        fn=interfaz_temperatura, 
        inputs=[temp_slider], 
        outputs=[output_plot, probs_state]
    )
    
    sample_button.click(
        fn=muestrear_boton,
        inputs=[probs_state],
        outputs=[output_text]
    )
    
    # Inicializar la interfaz
    demo.load(
        fn=interfaz_temperatura,
        inputs=[temp_slider],
        outputs=[output_plot, probs_state]
    )

# Lanzar la aplicación
if __name__ == "__main__":
    demo.launch()