"""
ml/evaluate.py - Evaluation routine for trained JackBot PPO policies.
"""
import argparse
import json
import time
from pathlib import Path
from typing import Optional, Dict, List, Any
import numpy as np


def evaluate(
    model_path: str,
    episodes: int = 5,
    use_gui: bool = True,
    num_robots: int = 16,  # Default to 16 to match your trained (352,) observation space
    robot_spacing: float = 0.5,
    start_pose: str = "init_deg",
    random_command: bool = True,
    save_json: Optional[str] = None,
) -> Dict[str, Any]:
    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

    from .env import JackBotEnv

    print(f"[Eval] Loading policy model from: {model_path}")
    # Force device="cpu" to prevent AMD ROCm/hipBLASLt matrix multiplication crashes
    model = PPO.load(model_path, device="cpu")

    # Initialize standard environment (returns single array of shape (352,))
    env = JackBotEnv(
        use_gui=use_gui,
        random_command=random_command,
        num_robots=num_robots,
        robot_spacing=robot_spacing,
        start_pose=start_pose,
    )

    episode_rewards: List[float] = []
    episode_lengths: List[int] = []

    for ep in range(episodes):
        obs, _ = env.reset()
        done = False
        total_reward = 0.0
        steps = 0

        print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} ---")

        while not done:
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, _ = env.step(action)
            
            done = terminated or truncated
            total_reward += float(reward)
            steps += 1

            if use_gui:
                time.sleep(1.0 / 240.0)

        episode_rewards.append(total_reward)
        episode_lengths.append(steps)
        print(f"Episode {ep + 1} Finished: Total Reward = {total_reward:.2f} | Steps = {steps}")

    env.close()

    metrics = {
        "model_path": str(model_path),
        "episodes_evaluated": episodes,
        "num_robots": num_robots,
        "mean_reward": float(np.mean(episode_rewards)),
        "std_reward": float(np.std(episode_rewards)),
        "mean_episode_length": float(np.mean(episode_lengths)),
        "raw_rewards": episode_rewards,
    }

    print("\n" + "=" * 50)
    print(f"EVALUATION COMPLETE ({episodes} Episodes)")
    print(f"Mean Reward:         {metrics['mean_reward']:.2f} ± {metrics['std_reward']:.2f}")
    print(f"Mean Episode Length: {metrics['mean_episode_length']:.1f} steps")
    print("=" * 50)

    if save_json:
        out_path = Path(save_json)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        with open(out_path, "w") as f:
            json.dump(metrics, f, indent=4)
        print(f"[Eval] Saved evaluation metrics to: {out_path.resolve()}")

    return metrics