"""Run a trained policy in the PyBullet sim with full curriculum progression.

Usage:
    python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 5 --gui
"""

import argparse
import json
import time
import sys
from pathlib import Path
import numpy as np
from stable_baselines3 import PPO
sys.path.append(str(Path(__file__).resolve().parent.parent))
from ml.env import JackBotEnv

def main():
    parser = argparse.ArgumentParser(description="JackBot Policy Evaluator")
    parser.add_argument("--model", type=str, required=True, help="Path to trained model checkpoint (.zip)")
    parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation episodes to run")
    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
    parser.add_argument("--no-random-command", dest="random_command", action="store_false", help="Lock commands to zero (disable random command sampling)")
    parser.set_defaults(random_command=True)
    parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to output evaluation summary")
    args = parser.parse_args()

    print(f"[Eval] Loading policy model from: {args.model}")
    model = PPO.load(args.model, device="cpu")

    # Instantiate single evaluation environment matching train setup
    env = JackBotEnv(
        use_gui=args.gui,
        random_command=args.random_command,
    )

    episode_rewards = []
    episode_lengths = []
    episode_phases = []

    try:
        for ep in range(args.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}/{args.episodes} [Phase: {initial_phase}] ---")

            while not done:
                # Deterministic prediction matches evaluation standards
                action, _ = model.predict(obs, deterministic=True)
                
                prev_phase = env.curriculum_phase
                obs, reward, terminated, truncated, _ = env.step(action)
                done = terminated or truncated

                total_reward += float(reward)
                steps += 1

                if env.curriculum_phase != prev_phase:
                    print(f"  └─ [Eval Milestone] Reached {env.curriculum_phase.name} at step {steps}!")

                if args.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}"
            )

    finally:
        env.close()

    # Calculate metrics report
    mean_reward = float(np.mean(episode_rewards))
    std_reward = float(np.std(episode_rewards))
    mean_length = float(np.mean(episode_lengths))

    print("\n" + "=" * 60)
    print(f"EVALUATION COMPLETE ({args.episodes} Episodes)")
    print(f"Final Reached Phase:  {env.curriculum_phase.name}")
    print(f"Mean Reward:          {mean_reward:.2f} ± {std_reward:.2f}")
    print(f"Mean Episode Length:  {mean_length:.1f} steps")
    print("=" * 60)

    if args.save_metrics:
        metrics = {
            "model_path": str(args.model),
            "episodes_evaluated": args.episodes,
            "final_curriculum_phase": env.curriculum_phase.name,
            "mean_reward": mean_reward,
            "std_reward": std_reward,
            "mean_episode_length": mean_length,
            "raw_rewards": episode_rewards,
            "episode_phases": episode_phases,
        }
        out_path = Path(args.save_metrics)
        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 report to: {out_path.resolve()}")


if __name__ == "__main__":
    main()