#!/usr/bin/env python3
"""
evaluate.py

Evaluates a trained MaskablePPO policy over multiple episodes and
prints summary statistics:

- Episode rewards
- Episode lengths (steps)
- Aggregate stats (mean, median, std, min, max)

Requires the trained model:
    models/model_policy.zip
"""

import os
import sys
from statistics import mean, median, stdev
from typing import List

# Ensure local utdg_env import even if not pip-installed
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

import numpy as np
import gymnasium as gym

from sb3_contrib import MaskablePPO
from sb3_contrib.common.wrappers import ActionMasker
from stable_baselines3.common.vec_env import DummyVecEnv

from utdg_env.utils.hydra_loader import load_config, pretty_print_cfg
import utdg_env


# ---------------------------------------------------------
# Action mask callback
# ---------------------------------------------------------
def mask_fn(env):
    """Return the boolean action mask from the deepest wrapped env."""
    while hasattr(env, "env"):
        env = env.env
    return env.get_action_mask()


# ---------------------------------------------------------
# Evaluation function
# ---------------------------------------------------------
def evaluate(model, env, num_episodes: int = 10, deterministic: bool = True):
    """Run evaluation episodes and return episode rewards + lengths."""

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

    for ep in range(num_episodes):

        obs = env.reset()
        done = False
        total_reward = 0.0
        steps = 0

        print(f"\n--- Episode {ep + 1}/{num_episodes} ---")

        while not done:
            action, _ = model.predict(obs, deterministic=deterministic)
            obs, reward, done, info = env.step(action)

            total_reward += reward[0]
            steps += 1

            print(f"[step {steps}] action={action[0]}, reward={reward[0]}, done={done[0]}")

        episode_rewards.append(total_reward)
        episode_lengths.append(steps)

        print(f"[episode result] total_reward={total_reward:.2f}, steps={steps}")

    return episode_rewards, episode_lengths


# ---------------------------------------------------------
# Main
# ---------------------------------------------------------
def main():

    # -------- Load config --------
    cfg = load_config()
    print("\n========== Loaded Config ==========")
    print(pretty_print_cfg(cfg))
    print("===================================\n")

    # -------- Create vectorized masked env --------
    env = DummyVecEnv([
        lambda: ActionMasker(gym.make("UTDGEnv-v0", sb3_mode=False), mask_fn)
    ])

    # -------- Load model --------
    model_path = os.path.join("models", "model_policy.zip")
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model not found at: {model_path}")

    print(f"[eval] Loading model from {model_path} ...")
    model = MaskablePPO.load(model_path, env=env)
    print("[eval] Model loaded.\n")

    # -------- Run evaluation --------
    num_episodes = 10
    rewards, lengths = evaluate(model, env, num_episodes=num_episodes, deterministic=True)

    # -------- Print aggregate stats --------
    print("\n========== Evaluation Summary ==========")
    print(f"Episodes: {num_episodes}")
    print(f"Rewards: {rewards}")
    print(f"Steps per episode: {lengths}")
    print("----------------------------------------")
    print(f"Mean reward:   {mean(rewards):.2f}")
    print(f"Median reward: {median(rewards):.2f}")
    if len(rewards) > 1:
        print(f"Std dev:       {stdev(rewards):.2f}")
    print(f"Min reward:    {min(rewards):.2f}")
    print(f"Max reward:    {max(rewards):.2f}")
    print("========================================\n")

    # -------- Cleanup --------
    env.close()
    print("[eval] Done.")


# ---------------------------------------------------------
if __name__ == "__main__":
    main()
