# This Gradio app provides a simple interface to evaluate a trained drone hovering policy.
# It allows users to select a checkpoint and optionally record a video of the evaluation.
# Note: The `torch` library is not available, so the code has been modified to remove the dependency.

import gradio as gr
import os
import pickle

# Assuming the `evaluate` function from the provided code is available here.
# The `evaluate` function is responsible for running the trained policy in a single environment with visualization.

def evaluate_policy(exp_name, ckpt, record):
    # Load the environment and policy configuration saved during training
    log_dir = f"logs/{exp_name}"
    if not os.path.exists(log_dir):
        raise FileNotFoundError(f"Log directory '{log_dir}' does not exist.  Did you run training?")
    env_cfg, obs_cfg, reward_cfg, command_cfg, train_cfg = pickle.load(open(f"{log_dir}/cfgs.pkl", "rb"))
    
    # For evaluation, we disable reward scaling (pure inference)
    reward_cfg["reward_scales"] = {}
    
    # Always visualize the target during evaluation
    env_cfg["visualize_target"] = True
    
    # Optionally set up a camera for recording
    env_cfg["visualize_camera"] = record
    env_cfg["max_visualize_FPS"] = 60
    
    # Build a single-environment instance with viewer
    env = HoverEnv(num_envs=1, env_cfg=env_cfg, obs_cfg=obs_cfg, reward_cfg=reward_cfg, command_cfg=command_cfg, show_viewer=True)
    runner = OnPolicyRunner(env, train_cfg, log_dir, device=gs.device)
    
    # Load the specified checkpoint
    resume_path = os.path.join(log_dir, f"model_{ckpt}.pt")
    runner.load(resume_path)
    
    # Get the inference policy
    policy = runner.get_inference_policy(device=gs.device)
    
    # Reset the environment
    obs, _ = env.reset()
    
    # Number of simulation steps equal to the episode duration times FPS
    max_sim_step = int(env_cfg["episode_length_s"] * env_cfg["max_visualize_FPS"])
    
    if record and env.cam is not None:
        env.cam.start_recording()
        for _ in range(max_sim_step):
            actions = policy(obs)
            obs, rews, dones, infos = env.step(actions)
            env.cam.render()
        env.cam.stop_recording(save_to_filename="video.mp4", fps=env_cfg["max_visualize_FPS"])
    else:
        for _ in range(max_sim_step):
            actions = policy(obs)
            obs, rews, dones, infos = env.step(actions)
    
    return "Evaluation completed successfully."

# Create a Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Drone Hovering Policy Evaluation")
    
    with gr.Row():
        exp_name = gr.Textbox(label="Experiment Name", placeholder="drone-hovering")
        ckpt = gr.Number(label="Checkpoint Index", value=300, precision=0)
        record = gr.Checkbox(label="Record Video", value=False)
    
    with gr.Row():
        evaluate_btn = gr.Button("Evaluate Policy")
    
    output = gr.Textbox(label="Evaluation Status")
    
    evaluate_btn.click(fn=evaluate_policy, inputs=[exp_name, ckpt, record], outputs=output)

# Launch the interface
demo.launch(show_error=True)