import accelerate
import gradio as gr
import joblib
import numpy as np
import requests
import torch
import os
import json
import matplotlib.pyplot as plt
from io import BytesIO
from transformers import AutoModelForCausalLM, AutoTokenizer

# Configuración del modelo de lenguaje
MODEL_NAME = "PlanTL-GOB-ES/roberta-base-bne"
device = "cuda" if torch.cuda.is_available() else "cpu"
HF_TOKEN = os.getenv("HF_TOKEN")

if not HF_TOKEN:
    raise ValueError("❌ ERROR: No se encontró HF_TOKEN. Asegúrate de definirlo en las variables de entorno.")

print("🔄 Cargando modelo de lenguaje...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)

# API de Node-RED
NODE_RED_URL = "https://appairedecolmena.es/colmena1/datos"
USERNAME = "user"
PASSWORD = "velutina"

def obtener_datos_colmena():
    """Obtiene los datos más recientes de Node-RED."""
    try:
        respuesta = requests.get(NODE_RED_URL, auth=(USERNAME, PASSWORD), timeout=5)
        if respuesta.status_code == 200:
            datos = respuesta.json()
            if "data" in datos and isinstance(datos["data"], list) and datos["data"]:
                return datos["data"]
            return {"error": "No hay datos recientes en Node-RED."}
        else:
            return {"error": f"Error en la API: {respuesta.status_code}"}
    except Exception as e:
        return {"error": str(e)}

def analizar_tendencia(datos, variable):
    """Analiza la tendencia de una variable específica."""
    valores = [float(d[variable]) for d in datos if variable in d]
    timestamps = [d["timestamp"] for d in datos if variable in d]
    promedio = np.mean(valores)
    maximo = np.max(valores)
    minimo = np.min(valores)

    respuesta = f"📊 **Análisis de {variable}:**\n"
    respuesta += f"- Promedio: {promedio:.2f}\n"
    respuesta += f"- Máximo: {maximo:.2f}\n"
    respuesta += f"- Mínimo: {minimo:.2f}\n"

    return respuesta, valores, timestamps

def generar_grafico(valores, timestamps, variable):
    """Genera un gráfico de la variable analizada."""
    plt.figure(figsize=(10, 5))
    plt.plot(timestamps, valores, marker='o', linestyle='-')
    plt.xlabel("Tiempo")
    plt.ylabel(variable)
    plt.title(f"Tendencia de {variable}")
    plt.xticks(rotation=45)
    plt.grid()
    
    buf = BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)
    return buf

def interpretar_instruccion(mensaje, datos):
    """Interpreta la consulta del usuario y ejecuta el análisis solicitado."""
    mensaje = mensaje.lower()
    
    # Ampliamos el mapeo con más sinónimos
    variable_map = {
        "temperatura": "temperaturaInterior",
        "calor": "temperaturaInterior",
        "frío": "temperaturaInterior",
        "humedad": "humedadInterior",
        "seco": "humedadInterior",
        "co2": "co2",
        "carbono": "co2",
        "peso": "peso",
        "colmena": "peso",
        "ventilación": "ver_ventilador",
        "ventilador": "ver_ventilador",
        "calefactor": "ver_calefactor",
        "vco": "vco",
        "frecuencia": "frecuencia"
    }
    
    # Buscar la variable en el mensaje del usuario
    for key, var in variable_map.items():
        if key in mensaje:
            respuesta, valores, timestamps = analizar_tendencia(datos, var)
            grafico = generar_grafico(valores, timestamps, var)
            return respuesta, grafico
    
    # Si no encuentra una coincidencia, usa el modelo de lenguaje para interpretar la consulta
    print("🤖 No se detectó una variable directamente. Usando IA para interpretar...")
    
    contexto = f"Usuario: {mensaje}\nColmena:"
    inputs = tokenizer(contexto, return_tensors="pt").to(device)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_length=150,  
            do_sample=True,
            top_k=50,
            temperature=0.7,
            pad_token_id=tokenizer.eos_token_id
        )

    respuesta_ia = tokenizer.decode(output[0], skip_special_tokens=True).strip()
    return respuesta_ia, None  # Si la IA responde, no genera imagen

def conversar_con_colmena(mensaje):
    """Genera una respuesta combinando el modelo de lenguaje con análisis de datos."""
    datos = obtener_datos_colmena()
    if "error" in datos:
        return datos["error"], None  # Retorna un string y `None` para evitar error en la imagen

    resultado = interpretar_instruccion(mensaje, datos)
    if isinstance(resultado, tuple):
        texto, grafico = resultado
        if grafico:  
            # Convierte BytesIO en una imagen para Gradio
            from PIL import Image
            img = Image.open(grafico)
            return texto, img  
        else:
            return texto, None  
    else:
        return resultado, None  # Asegurar que siempre haya dos valores de retorno



iface = gr.Interface(
    fn=conversar_con_colmena,
    inputs="text",
    outputs=["text", "image"],
    title="🐝 Chat con la Colmena",
    description="Consulta el estado de la colmena y pide análisis de datos."
)

iface.launch()