"""
ml/run_eval_training.py - Benchmark reward system across all curriculum phases.
"""
import time
import numpy as np
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).resolve().parent.parent))
from ml.env import JackBotEnv, CurriculumPhase

def evaluate_kinematics(episode_length: int = 1000):
    env = JackBotEnv(
        use_gui=True, 
        random_command=False, 
        max_episode_steps=episode_length,
        robot_mode="kinematics"
    )

    print("\n" + "=" * 70)
    print(" RUNNING MULTI-PHASE REWARD BENCHMARK (KINEMATICS MODE)")
    print("=" * 70 + "\n")

    phase_configs = [
        (CurriculumPhase.STAND_ONLY,           "STAND",              np.array([1.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)),
    ]

    for phase_enum, label, cmd in phase_configs:
        obs, _ = env.reset()
        
        env.curriculum_phase = phase_enum
        env.command = cmd.copy()
        
        done = False
        total_reward = 0.0
        step_count = 0

        while not done:
            dummy_action = np.zeros(18, dtype=np.float32)
            obs, reward, terminated, truncated, _ = env.step(dummy_action)
            total_reward += reward
            step_count += 1
            done = terminated or truncated

            time.sleep(1.0 / 60.0)

        comp_averages = env.get_reward_component_averages()

        print(f"\n--- Episode Stage: [{phase_enum.name}] ({label}) ---")
        print(f"Command Applied: vx={cmd[0]:.2f}, vy={cmd[1]:.2f}, yaw={cmd[2]:.2f}")
        print(f"Total Episode Reward: {total_reward:.4f}")
        print("Component Step Averages:")
        for name, value in comp_averages.items():
            print(f"  • {name:<20}: {value:+.5f}")

        metrics = env.get_current_robot_metrics()
        dist = metrics[0]["distance_from_start"] if metrics else 0.0
        speed = metrics[0]["speed"] if metrics else 0.0
        avg_reward = total_reward / max(1, step_count)

        print(f"  ├─ Average Reward / Step: {avg_reward:.4f}")
        print(f"  ├─ Distance Travelled: {dist:.2f} m")
        print(f"  ├─ Actual Avg Speed: {speed:.2f} m/s")
        print(f"  └─ Steps Survived: {step_count} / {episode_length}")
        print("-" * 70)

    env.close()

if __name__ == "__main__":
    evaluate_kinematics()