"""
ml/evaluate.py - Curriculum-aware 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,
    robot_spacing: float = 0.5,
    start_pose: str = "init_deg",
    random_command: bool = True,  # Default to True so curriculum commands are sampled
    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 environment with random commands enabled for curriculum progression
    env = JackBotEnv(
        use_gui=use_gui,
        random_command=random_command,
        robot_spacing=robot_spacing,
        start_pose=start_pose,
    )

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

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

        initial_phase = env.curriculum_phase.name
        print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} [Phase: {initial_phase}] ---")

        while not done:
            action, _ = model.predict(obs, deterministic=True)
            
            # Store phase before step to detect live phase transitions
            prev_phase = env.curriculum_phase
            
            obs, reward, terminated, truncated, _ = env.step(action)
            
            # Done on either physical failure (terminated) or phase step limit (truncated)
            done = terminated or truncated
            total_reward += float(reward)
            steps += 1

            # Log live phase transition if unlocked during this step
            if env.curriculum_phase != prev_phase:
                print(f"  └─ [Eval Milestone] Curriculum advanced to {env.curriculum_phase.name} at episode step {steps}!")

            if use_gui:
                time.sleep(1.0 / 240.0)

        status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)"
        final_phase = env.curriculum_phase.name
        
        episode_rewards.append(total_reward)
        episode_lengths.append(steps)
        episode_phases.append(final_phase)

        print(
            f"Episode {ep + 1} Finished [{status_str}]: "
            f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}"
        )

    env.close()

    metrics = {
        "model_path": str(model_path),
        "episodes_evaluated": episodes,
        "final_curriculum_phase": env.curriculum_phase.name,
        "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,
        "episode_phases": episode_phases,
    }

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

    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