"""
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 = 1,
    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}")
    model = PPO.load(model_path)

    # Initialize environment with active command sampling so robot actually walks
    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:
            # Deterministic evaluation (no exploration noise)
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, _ = env.step(action)
            
            done = terminated or truncated
            total_reward += float(reward)
            steps += 1

            # Give PyBullet GUI frame pacing if running visually
            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()

    # Calculate summary statistics
    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)

    # Optional JSON metrics export
    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


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate a trained JackBot policy.")
    parser.add_argument("--model-path", type=str, required=True, help="Path to trained PPO model zip")
    parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation rounds")
    parser.add_argument("--gui", action="store_true", help="Render GUI simulation")
    parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in sim")
    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")
    parser.add_argument("--save-json", type=str, default=None, help="Optional output JSON path for metrics")
    args = parser.parse_args()

    evaluate(
        model_path=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,
        save_json=args.save_json,
    )