import argparse
import numpy as np

from .env import JackBotEnv


def evaluate(
    model_path: str,
    episodes: int = 5,
    use_gui: bool = False,
    num_robots: int = 1,
    robot_spacing: float = 0.5,
    start_pose: str = "init_deg",
):
    try:
        from stable_baselines3 import PPO
    except ImportError as exc:
        raise ImportError(
            "stable-baselines3 is required for evaluation. Install with: pip install stable-baselines3"
        ) from exc

    env = JackBotEnv(
        use_gui=use_gui,
        random_command=False,
        num_robots=num_robots,
        robot_spacing=robot_spacing,
        start_pose=start_pose,
    )
    model = PPO.load(model_path)

    for episode in range(episodes):
        reset_res = env.reset()
        # handle Gym / Gymnasium compatibility: reset may return (obs, info)
        if isinstance(reset_res, tuple) and len(reset_res) == 2:
            obs, _ = reset_res
        else:
            obs = reset_res

        done = False
        episode_reward = 0.0

        while not done:
            # pass only the observation to the policy
            action, _ = model.predict(obs, deterministic=True)

            step_res = env.step(action)
            # Gymnasium-style: (obs, reward, terminated, truncated, info)
            if isinstance(step_res, tuple) and len(step_res) == 5:
                obs, reward, terminated, truncated, info = step_res
                done = bool(terminated or truncated)
            else:
                # legacy Gym: (obs, reward, done, info)
                obs, reward, done, info = step_res

            episode_reward += float(reward)

        print(f"Episode {episode + 1}: reward={episode_reward:.2f}")

    env.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate a trained JackBot policy.")
    parser.add_argument("--model-path", type=str, required=True)
    parser.add_argument("--episodes", type=int, default=5)
    parser.add_argument("--gui", action="store_true")
    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the environment")
    parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
    args = parser.parse_args()
    evaluate(
        args.model_path,
        episodes=args.episodes,
        use_gui=args.gui,
        num_robots=args.num_robots,
        robot_spacing=args.robot_spacing,
        start_pose=args.start_pose,
    )