import gradio as gr
import gymnasium as gym
import panda_gym
import imageio
import os
from stable_baselines3 import A2C
from huggingface_hub import hf_hub_download

# 1. Download and load your trained model
repo_id = "sanju-1007/PandaReachDense-PandaReachDense-v3"
filename = "a2c-PandaReachDense-v3.zip" # Updated to match your exact repo file
custom_objects = {
    "learning_rate": 0.0,
    "lr_schedule": lambda _: 0.0,
    "clip_range": lambda _: 0.0,
}

checkpoint = hf_hub_download(repo_id=repo_id, filename=filename)
model = A2C.load(checkpoint, custom_objects=custom_objects, device="cpu")

def simulate_agent(seed, num_episodes=3):
    """Runs multiple live simulations and strings them into one smooth MP4 video."""
    video_path = "live_rollout.mp4"
    
    # Initialize environment for RGB rendering
    env = gym.make("PandaReachDense-v3", render_mode="rgb_array")
    frames = []
    
    # Run 3 back-to-back episodes so the video is a good length (5-8 seconds)
    for ep in range(num_episodes):
        # We add 'ep' to the seed so the target moves to a new spot each time!
        obs, info = env.reset(seed=int(seed) + ep)
        
        for _ in range(100):
            # The agent predicts the next movement
            action, _states = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            
            # Capture the visual frame
            frames.append(env.render())
            
            # If the robot successfully hits the target...
            if terminated or truncated:
                # Add 20 'padding' frames (about 0.6 seconds) so the viewer 
                # can actually see the robot holding its final position.
                for _ in range(20):
                    frames.append(env.render())
                break # Move on to the next episode
            
    env.close()
    
    # Compile frames to a smooth 30 FPS video
    imageio.mimsave(video_path, frames, fps=30)
    return video_path

# 2. Build the Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 Live A2C Agent: PandaReachDense-v3")
    gr.Markdown("Change the seed to randomize the target position and watch the trained agent calculate the trajectory in real-time.")
    
    with gr.Row():
        with gr.Column():
            seed_input = gr.Slider(minimum=0, maximum=1000, step=1, label="Environment Random Seed", value=42)
            run_button = gr.Button("▶️ Run Simulation Live", variant="primary")
        
        with gr.Column():
            video_output = gr.Video(label="Agent Rollout")
            
    run_button.click(fn=simulate_agent, inputs=seed_input, outputs=video_output)

# Launch the app
demo.launch()