import gradio as gr
import torch
import gymnasium as gym
import mpi4py
import imageio
import os
import sys
from huggingface_hub import hf_hub_download

# --- Block to manually load your custom library ---
SPINUP_REPO_PATH = 'spinningup'
if not os.path.exists(SPINUP_REPO_PATH):
    print("Cloning modernized spinningup repository...")
    os.system('git clone https://github.com/monigarr/spinningup.git -b monigarr-dev')

sys.path.insert(0, os.path.abspath(SPINUP_REPO_PATH))
print("Spinning Up added to path.")

from spinup.algos.pytorch.sac.core import MLPActorCritic
# --- End of block ---


# IMPORTANT: Update this to your HF model repo ID
REPO_ID = "monigarr/spinning-up-on-ant-sac"

# Load the trained model from the Hub
model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pt")
if model_path is None or not os.path.exists(model_path):
    raise FileNotFoundError(f"Model file not found at {model_path}")
model = torch.load(model_path, weights_only=False)
print("MoniGarr’s SAC Ant V5 Model Loaded Successfully")


def run_agent(max_steps=500):
    """
    Runs the agent in the Ant-v5 environment and returns a video.
    """
    # Create the environment with "rgb_array" render mode for video capture
    env = gym.make("Ant-v5", render_mode="rgb_array")
    frames = []

    (obs, _) = env.reset()
    
    for _ in range(max_steps):
        # Capture the frame BEFORE taking a step
        frames.append(env.render())
        
        # Get action from the loaded policy
        with torch.no_grad():
            action = model.act(torch.as_tensor(obs, dtype=torch.float32))
        
        # Take a step in the environment
        obs, _, terminated, truncated, _ = env.step(action)
        
        if terminated or truncated:
            break
            
    env.close()
    
    # Save frames as a video
    video_path = "monigarr_ant_v5_demo.gif"
    imageio.mimsave(video_path, frames, fps=30)
    
    return video_path

# Create the Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# OpenAI Spinning Up: SAC Agent on Ant-v5")
    gr.Markdown(
        "This demo loads a pre-trained Soft Actor-Critic (SAC) agent "
        "and visualizes its performance in the `Ant-v5` environment. "
        "Click the button below to generate a new simulation video." \
        "Details: https://github.com/monigarr/spinningup/tree/monigarr-dev"
        "MoniGarr Modernized Spinning Up: https://researchengineer.wordpress.com/2025/07/15/reviving-spinning-up-a-modern-accessible-rl-learning-framework/"
        "Benchmarking SAC Agents with OpoenAI Spinning Up: https://researchengineer.wordpress.com/2025/07/19/spinning-up-sac-ant-results/"
        "What I learned Benchmarking SAC Agents: https://researchengineer.wordpress.com/2025/07/19/what-i-learned-benchmarking-sac-agents-with-openai-spinning-up-modernized/"
    )

    
    run_button = gr.Button("Run Agent Simulation", variant="primary")
    video_output = gr.Video(label="Agent Performance")
    
    run_button.click(
        fn=run_agent,
        inputs=None,
        outputs=video_output
    )

demo.launch()