"""
ml/run_eval.py - Phase-by-Phase Policy Evaluator for JackBot
Evaluates a trained model across all curriculum phases with fixed command vectors.
"""

import argparse
import json
import time
import sys
from pathlib import Path
import numpy as np
from stable_baselines3 import PPO

# Ensure project root is in sys.path
sys.path.append(str(Path(__file__).resolve().parent.parent))
from ml.env import JackBotEnv, CurriculumPhase


def main():
    parser = argparse.ArgumentParser(description="JackBot Phase-by-Phase Policy Evaluator")
    parser.add_argument("--model", type=str, required=True, help="Path to trained model checkpoint (.zip)")
    parser.add_argument("--episodes-per-phase", type=int, default=1, help="Number of test episodes per phase")
    parser.add_argument("--max-steps-per-episode", type=int, default=600, help="Max simulation steps per episode")
    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
    parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to save evaluation summary")
    args = parser.parse_args()

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

    # Instantiate environment with random_command disabled so our test command stays locked
    env = JackBotEnv(
        use_gui=args.gui,
        random_command=False,
        max_episode_steps=args.max_steps_per_episode,
        robot_mode="direct"
    )

    # Multi-Phase Configurations Suite
    phase_configs = [
        #(CurriculumPhase.STAND_ONLY,          "STAND",                 np.array([0.0, 0.0, 0.0], dtype=np.float32)),
        (CurriculumPhase.FORWARD,             "FORWARD GAIT",           np.array([1.0, 0.0, 0.0], dtype=np.float32)),
        (CurriculumPhase.TURN_AND_DIRECTION,  "FORWARD + YAW TURN",     np.array([0.5, 0.0, 0.4], dtype=np.float32)),
        (CurriculumPhase.OMNI_DIRECTION,      "STRIDE LATERAL",         np.array([0.5, 0.5, 0.0], dtype=np.float32)),
        (CurriculumPhase.FULL_COMMAND,        "FULL OMNI COMBINATION",  np.array([0.3, 0.2, 0.3], dtype=np.float32)),
    ]

    phase_summary = []

    try:
        print("\n" + "=" * 75)
        print(" STARTING MULTI-PHASE EVALUATION SUITE")
        print("=" * 75)

        for phase_enum, phase_name, test_cmd in phase_configs:
            print(f"\n▶ Testing Phase [{phase_enum.value}]: {phase_enum.name} ({phase_name})")
            print(f"  Target Command Vector [vx, vy, omega]: {test_cmd.tolist()}")

            ep_rewards = []
            ep_steps = []
            ep_distances = []

            for ep in range(args.episodes_per_phase):
                obs, _ = env.reset()

                # Force environment into active curriculum phase and lock command BEFORE getting obs
                env.curriculum_phase = phase_enum
                env.command = test_cmd.copy()
                cmd_vx, cmd_vy, cmd_omega = test_cmd
                env.robot.robot_state = (
                    "walking"
                    if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01)
                    else "idle"
                )
                env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]

                # Get correct observation with test_cmd attached
                obs = env._get_obs()

                done = False
                total_reward = 0.0
                steps = 0

                while not done:
                    # Predict deterministic action from policy
                    action, _ = model.predict(obs, deterministic=True)
                    
                    # Step environment
                    obs, reward, terminated, truncated, _ = env.step(action)
                    done = terminated or truncated

                    total_reward += float(reward)
                    steps += 1

                    if args.gui:
                        time.sleep(1.0 / 60.0)

                status_str = "FAILED (Collapsed)" if terminated else "SUCCESS (Completed)"
                dist = float(env.max_distance_from_start)
                print(
                    f"   └─ Ep {ep + 1}/{args.episodes_per_phase}: {status_str:<19} | "
                    f"Steps: {steps:<4} | Reward: {total_reward:+.2f} | Max Dist: {dist:.2f}m"
                )

                ep_rewards.append(total_reward)
                ep_steps.append(steps)
                ep_distances.append(dist)

            phase_summary.append({
                "phase_id": phase_enum.value,
                "phase_name": phase_enum.name,
                "label": phase_name,
                "command": test_cmd.tolist(),
                "mean_reward": float(np.mean(ep_rewards)),
                "mean_steps": float(np.mean(ep_steps)),
                "mean_distance": float(np.mean(ep_distances)),
            })

    finally:
        env.close()

    # Print Summary Table
    print("\n" + "=" * 80)
    print(" EVALUATION SUMMARY REPORT")
    print("=" * 80)
    print(f"{'Phase ID & Name':<25} | {'Label':<22} | {'Reward':<8} | {'Steps':<6} | {'Max Dist':<8}")
    print("-" * 80)
    for res in phase_summary:
        phase_str = f"[{res['phase_id']}] {res['phase_name']}"
        print(
            f"{phase_str:<25} | "
            f"{res['label']:<22} | "
            f"{res['mean_reward']:<+8.2f} | "
            f"{res['mean_steps']:<6.0f} | "
            f"{res['mean_distance']:<8.2f}m"
        )
    print("=" * 80)

    # Save JSON Report if requested
    if args.save_metrics:
        out_path = Path(args.save_metrics)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        with open(out_path, "w") as f:
            json.dump({"model_path": str(args.model), "summary": phase_summary}, f, indent=4)
        print(f"\n[Eval] Saved report to: {out_path.resolve()}")


if __name__ == "__main__":
    main()