import gradio as gr
import gymnasium as gym
from gymnasium import spaces
import numpy as np
import matplotlib.pyplot as plt
from stable_baselines3 import PPO

# --- 1. DEFINIMOS EL ENTORNO (Igual que en Colab) ---
class DataCenterEnv(gym.Env):
    def __init__(self):
        super(DataCenterEnv, self).__init__()
        self.action_space = spaces.Discrete(5)
        self.observation_space = spaces.Box(low=0, high=100, shape=(1,), dtype=np.float32)
        self.state = 20
        self.max_steps = 100
        self.current_step = 0
        self.temps = []
        self.energies = []

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self.state = 20 + np.random.rand()
        self.current_step = 0
        self.temps = [self.state]
        self.energies = []
        return np.array([self.state], dtype=np.float32), {}

    def step(self, action):
        # Usamos variables globales para simular la carga que el usuario elija en la web
        global USER_CPU_LOAD_MIN, USER_CPU_LOAD_MAX
        cpu_load = np.random.uniform(USER_CPU_LOAD_MIN, USER_CPU_LOAD_MAX)
        
        cooling_effect = action * 2.5  # Potencia industrial
        energy_cost = action ** 2
        
        self.state = self.state + cpu_load - cooling_effect
        self.state = np.clip(self.state, 10, 100)
        self.current_step += 1
        
        reward = 0
        if 37 <= self.state <= 60: reward += 5
        else: reward -= 2
        if self.state > 80: reward -= 50
        reward -= (energy_cost * 0.1)
        
        terminated = False
        if self.current_step >= self.max_steps: terminated = True
        
        self.temps.append(self.state)
        self.energies.append(energy_cost)
            
        return np.array([self.state], dtype=np.float32), reward, terminated, False, {}

# Variables globales para controlar la dificultad desde la web
USER_CPU_LOAD_MIN = 1
USER_CPU_LOAD_MAX = 4

# --- 2. FUNCIÓN DE SIMULACIÓN (Lo que hace la web) ---
def run_simulation(load_min, load_max):
    # Actualizamos las variables globales con lo que el usuario eligió
    global USER_CPU_LOAD_MIN, USER_CPU_LOAD_MAX
    USER_CPU_LOAD_MIN = load_min
    USER_CPU_LOAD_MAX = load_max

    # 1. Crear entorno y Cargar Modelo
    env = DataCenterEnv()
    
    # INTENTA CARGAR TU MODELO. Si no lo encuentra, usa uno aleatorio (para que no falle la demo)
    try:
        model = PPO.load("agente_aire_acondicionado.zip", env=env)
        msg = "✅ Modelo Inteligente Cargado Correctamente."
    except:
        model = None
        msg = "⚠️ ADVERTENCIA: No se encontró 'agente_aire_acondicionado.zip'. Usando acciones aleatorias."

    # 2. Correr la simulación
    obs, _ = env.reset()
    done = False
    
    while not done:
        if model:
            action, _ = model.predict(obs)
        else:
            action = env.action_space.sample() # Fallback aleatorio
            
        obs, reward, done, truncated, info = env.step(action)

    # 3. Generar Gráficos
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # Gráfico Temperatura
    ax1.plot(env.temps, color="red", label="Temperatura")
    ax1.axhline(y=37, color='green', linestyle='--', label="Mínimo (37°C)")
    ax1.axhline(y=60, color='green', linestyle='--', label="Máximo (60°C)")
    ax1.set_title("Control de Temperatura")
    ax1.set_ylabel("°C")
    ax1.set_ylim(10, 90)
    ax1.legend()
    ax1.grid(True)
    
    # Gráfico Energía
    ax2.plot(env.energies, color="blue", label="Energía")
    ax2.set_title("Consumo de Energía")
    ax2.set_ylabel("Watts")
    ax2.grid(True)
    
    plt.tight_layout()
    return fig, msg

# --- 3. INTERFAZ GRÁFICA (Gradio) ---
with gr.Blocks() as demo:
    gr.Markdown("# ❄️ AI Cooling System Optimization Agent")
    gr.Markdown("Este agente de Aprendizaje por Refuerzo (RL) controla el aire acondicionado de un servidor para ahorrar energía sin que se sobrecaliente.")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Configuración de Carga")
            s_min = gr.Slider(1, 10, value=1, label="Carga Mínima de CPU (Calor)")
            s_max = gr.Slider(1, 10, value=4, label="Carga Máxima de CPU (Calor)")
            btn = gr.Button("🚀 Ejecutar Simulación")
            status = gr.Textbox(label="Estado del Agente")
        
        with gr.Column():
            plot = gr.Plot(label="Resultados en Tiempo Real")

    btn.click(fn=run_simulation, inputs=[s_min, s_max], outputs=[plot, status])

# Lanzar la app
demo.launch()