import gradio as gr
import torch
from transformers import pipeline
from PIL import Image
import folium
from geopy.geocoders import Nominatim
from geopy.distance import geodesic
import time
import random

# --- 1. CARGA DE MODELOS AVANZADOS (Se ejecuta una sola vez al iniciar) ---
# Esta es la parte más pesada. En un Space de Hugging Face, esto puede tardar un poco en cargar.

print("Cargando modelos de IA. Por favor, espere...")
try:
    # Pipeline para clasificación de imágenes (Usa un Vision Transformer - ViT)
    vision_classifier = pipeline("image-classification", model="google/vit-base-patch16-224")
    
    # Pipeline para generación de texto a texto (Usa un LLM pequeño - T5)
    text_generator = pipeline("text2text-generation", model="google-t5/t5-small")
    
    print("Modelos de IA cargados exitosamente.")
    models_loaded = True
except Exception as e:
    print(f"Error al cargar los modelos: {e}")
    models_loaded = False

# --- 2. FUNCIONES DE AYUDA Y LÓGICA DE LA APLICACIÓN ---

# Inicializar geolocalizador
geolocator = Nominatim(user_agent="ecomind_ai_app")

# Mapeo de etiquetas generales a categorías de residuos específicas
def map_label_to_category(label):
    """Mapea una etiqueta del clasificador de ImageNet a una categoría de residuo."""
    label = label.lower()
    if any(keyword in label for keyword in ["bottle", "can", "plastic bag", "cup"]):
        return "Reciclable (Contenedor Amarillo)", "una botella de plástico o lata"
    if any(keyword in label for keyword in ["banana", "apple", "orange", "food"]):
        return "Orgánico (Contenedor Marrón)", "un residuo orgánico"
    if any(keyword in label for keyword in ["carton", "envelope", "paper towel"]):
        return "Papel y Cartón (Contenedor Azul)", "una caja de cartón o papel"
    if "battery" in label:
        return "Punto Limpio (Residuo Peligroso)", "una batería"
    return "General (Contenedor Gris)", "un residuo general"

# --- 3. FUNCIONES PRINCIPALES VINCULADAS A LA INTERFAZ ---

def classify_and_suggest_waste(image, state):
    """
    Función principal para el clasificador de residuos.
    1. Clasifica la imagen usando el modelo de visión.
    2. Mapea la etiqueta a una categoría.
    3. Usa el LLM para generar una idea de reciclaje.
    4. Actualiza el estado del usuario.
    """
    if not models_loaded or image is None:
        return "Los modelos de IA no están disponibles o no se subió una imagen.", None, state, "Error"

    # Clasificar la imagen
    predictions = vision_classifier(image)
    best_prediction = predictions[0]
    label = best_prediction['label']
    
    # Mapear a nuestra categoría y obtener un descriptor
    category, descriptor = map_label_to_category(label)

    # Generar idea de reciclaje creativa con el LLM
    prompt = f"Genera una idea de reciclaje corta y creativa para {descriptor} en español:"
    try:
        generated_ideas = text_generator(prompt, max_length=40, num_beams=4, early_stopping=True)
        creative_idea = generated_ideas[0]['generated_text']
    except Exception:
        creative_idea = "No se pudo generar una idea creativa en este momento."

    # Actualizar el estado del usuario
    state["recycled_items"] += 1
    
    # Formatear la salida
    result_text = f"""
    **Análisis de IA Completado:**
    - **Clasificación del Modelo:** `{label.title()}` (Confianza: {best_prediction['score']:.2%})
    - **Categoría de Residuo:** **{category}**
    - **Idea Creativa (Generada por LLM):** *{creative_idea}*
    """
    return result_text, gr.update(value=state), f"¡Clasificado! {category}"

def plan_eco_route(start_loc, end_loc, transport_mode, state):
    """
    Genera un mapa interactivo con una ruta simulada y calcula el ahorro de CO2.
    """
    try:
        location1 = geolocator.geocode(start_loc)
        location2 = geolocator.geocode(end_loc)
        
        if not location1 or not location2:
            return None, "No se pudieron encontrar una o ambas ubicaciones.", state

        coords1 = (location1.latitude, location1.longitude)
        coords2 = (location2.latitude, location2.longitude)

        # Calcular distancia
        distance_km = geodesic(coords1, coords2).km

        # Simulación de emisiones (g CO2 por km)
        emissions = {"Coche": 200, "Transporte Público": 40, "Bicicleta": 0, "Caminar": 0}
        
        car_emissions_g = distance_km * emissions["Coche"]
        mode_emissions_g = distance_km * emissions[transport_mode]
        co2_saved_g = car_emissions_g - mode_emissions_g
        
        # Actualizar el estado
        state["co2_saved_kg"] += co2_saved_g / 1000

        # Crear mapa con Folium
        map_center = [(c1+c2)/2 for c1, c2 in zip(coords1, coords2)]
        m = folium.Map(location=map_center, zoom_start=13)

        # Añadir marcadores
        folium.Marker(coords1, popup=f"Inicio: {start_loc}", icon=folium.Icon(color='green')).add_to(m)
        folium.Marker(coords2, popup=f"Destino: {end_loc}", icon=folium.Icon(color='red')).add_to(m)

        # Dibujar líneas (simuladas)
        folium.PolyLine([coords1, coords2], color="red", weight=2.5, opacity=1, tooltip="Ruta en coche (referencia)").add_to(m)
        # Simular una ruta eco un poco diferente
        eco_midpoint = [map_center[0] + (random.random()-0.5)*0.01, map_center[1] + (random.random()-0.5)*0.01]
        folium.PolyLine([coords1, eco_midpoint, coords2], color="green", weight=3.5, opacity=0.8, tooltip=f"Tu Ruta en {transport_mode}").add_to(m)
        
        map_path = "/tmp/eco_route_map.html"
        m.save(map_path)
        
        summary = f"""
        **Plan de Ruta Ecológica:**
        - **Distancia:** {distance_km:.2f} km
        - **Modo Elegido:** {transport_mode}
        - **Ahorro de CO₂ (vs. Coche):** **{co2_saved_g / 1000:.2f} kg**
        """
        
        return map_path, summary, gr.update(value=state)

    except Exception as e:
        return None, f"Ocurrió un error: {e}", state

def mindfulness_session_generator(duration):
    """
    Genera una sesión de mindfulness guiada por el LLM con un temporizador en tiempo real.
    """
    if not models_loaded:
        yield "El modelo de IA para mindfulness no está disponible."
        return

    prompt = "Genera una instrucción de mindfulness muy corta, calmada y en presente para un ejercicio de respiración en español:"
    try:
        instruction = text_generator(prompt, max_length=50)[0]['generated_text']
    except Exception:
        instruction = "Concéntrate en tu respiración. Siente cómo el aire entra y sale."

    yield f"**Iniciando Sesión de {duration} segundos...**\n\n*{instruction}*"
    time.sleep(2)

    for i in range(duration, 0, -1):
        yield f"*{instruction}*\n\n**Tiempo restante: {i} segundos...**"
        time.sleep(1)
        
    yield "**Sesión completada.**\n\nEspero que te sientas más centrado y en calma."


# --- 4. DEFINICIÓN DE LA INTERFAZ DE USUARIO CON GRADIO ---

with gr.Blocks(theme=gr.themes.Soft(primary_hue="green", secondary_hue="blue"), title="EcoMind AI Avanzado") as demo:
    
    # Definir el estado de la sesión del usuario
    user_state = gr.State({"recycled_items": 0, "co2_saved_kg": 0.0})
    
    gr.Markdown("# 🌍 EcoMind AI: Versión Avanzada")
    gr.Markdown("Un prototipo funcional que utiliza IA de Visión, Lenguaje y Geo-localización para una vida sostenible.")

    with gr.Tabs():
        with gr.TabItem("📊 Mi Impacto"):
            gr.Markdown("## Tu Panel de Impacto Positivo")
            with gr.Row():
                recycled_items_display = gr.Number(
                    label="♻️ Objetos Reciclados Correctamente", 
                    value=0,
                    interactive=False
                )
                co2_saved_display = gr.Number(
                    label="💨 CO₂ Ahorrado (kg)", 
                    value=0.0,
                    interactive=False,
                    precision=2
                )

        with gr.TabItem("🗑️ Clasificador Inteligente"):
            gr.Markdown("### Usa IA para identificar y reciclar residuos correctamente.")
            with gr.Row(equal_height=True):
                with gr.Column(scale=1):
                    image_input = gr.Image(type="pil", label="Sube una foto del residuo", height=300)
                    classify_btn = gr.Button("Analizar con IA", variant="primary")
                with gr.Column(scale=2):
                    classification_label = gr.Label(label="Categoría de Residuo")
                    result_text_output = gr.Markdown(label="Análisis y Sugerencia de la IA")

        with gr.TabItem("🗺️ Ruta Ecológica"):
            gr.Markdown("### Planifica tus viajes y calcula tu ahorro de CO₂.")
            with gr.Row():
                with gr.Column(scale=1):
                    start_input = gr.Textbox(label="📍 Punto de Partida", placeholder="Ej: Sagrada Familia, Barcelona")
                    end_input = gr.Textbox(label="🏁 Punto de Destino", placeholder="Ej: Parque Güell, Barcelona")
                    transport_mode_input = gr.Radio(
                        ["Caminar", "Bicicleta", "Transporte Público", "Coche"],
                        label="Modo de Transporte",
                        value="Bicicleta"
                    )
                    route_btn = gr.Button("Planificar Ruta Ecológica", variant="primary")
                with gr.Column(scale=2):
                    route_summary_output = gr.Markdown(label="Resumen de Ahorro")
                    map_output = gr.HTML(label="Mapa Interactivo de la Ruta")

        with gr.TabItem("🧘 Bienestar Personal"):
            gr.Markdown("### Tómate un momento para ti con una sesión de mindfulness guiada por IA.")
            with gr.Row():
                with gr.Column():
                    duration_slider = gr.Slider(10, 120, value=30, step=5, label="Duración de la sesión (segundos)")
                    mindfulness_btn = gr.Button("Iniciar Sesión de Mindfulness", variant="primary")
                with gr.Column():
                    mindfulness_output = gr.Markdown(label="Guía de Mindfulness")

    # --- 5. LÓGICA DE CONEXIÓN DE COMPONENTES ---

    # Conectar el estado a los paneles de visualización
    def update_dashboard(state):
        return state["recycled_items"], state["co2_saved_kg"]

    # Clasificador
    classify_btn.click(
        fn=classify_and_suggest_waste,
        inputs=[image_input, user_state],
        outputs=[result_text_output, user_state, classification_label]
    ).then(
        fn=update_dashboard,
        inputs=user_state,
        outputs=[recycled_items_display, co2_saved_display]
    )

    # Ruta Ecológica
    route_btn.click(
        fn=plan_eco_route,
        inputs=[start_input, end_input, transport_mode_input, user_state],
        outputs=[map_output, route_summary_output, user_state]
    ).then(
        fn=update_dashboard,
        inputs=user_state,
        outputs=[recycled_items_display, co2_saved_display]
    )

    # Mindfulness (usa 'yield' para salida iterativa)
    mindfulness_btn.click(
        fn=mindfulness_session_generator,
        inputs=duration_slider,
        outputs=mindfulness_output
    )


if __name__ == "__main__":
    if not models_loaded:
        print("\nADVERTENCIA: Los modelos de IA no se cargaron. La aplicación se ejecutará en modo degradado.")
    demo.launch(debug=True)