"""Evaluation: run a policy over a held-out spec suite on fixed seeds.

Seeds are enumerated from a fixed base so every checkpoint and baseline sees
IDENTICAL walls - eval numbers are comparable across the whole project.

    uv run python -m train.evaluate --policy oracle --suite interp --episodes 20
    uv run python -m train.evaluate --checkpoint runs/x/ckpt.pt --suite interp
"""

from __future__ import annotations

import argparse
import json
import statistics

import gymnasium as gym

import atrium_sim  # noqa: F401
from atrium_sim.blueprint import _SUITES

BASE_SEED = 10000

METRIC_KEYS = ("episode_return", "frac_in_tol", "frac_filled", "waste_frac",
               "completed", "score", "mean_abs_dev_mm")


def evaluate(env, policy, suite: str, episodes: int) -> dict:
    """Round-robin the suite's specs on fixed seeds; aggregate terminal metrics."""
    specs = _SUITES[suite]
    per_episode: list[dict] = []
    for i in range(episodes):
        spec = specs[i % len(specs)]
        obs, info = env.reset(seed=BASE_SEED + i, options={"spec": spec})
        done = False
        while not done:
            obs, r, terminated, truncated, info = env.step(policy.act(obs))
            done = terminated or truncated
        per_episode.append(info["metrics"])
    summary = {
        k: {
            "mean": statistics.fmean(ep[k] for ep in per_episode),
            "std": statistics.pstdev([ep[k] for ep in per_episode]),
        }
        for k in METRIC_KEYS
    }
    return {"suite": suite, "episodes": episodes, "metrics": summary}


def load_policy(args, env):
    if args.checkpoint:
        from train.agent import CheckpointPolicy

        return CheckpointPolicy(args.checkpoint)
    from baselines.policy import make_policy

    return make_policy(args.policy, env, seed=0)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--policy", default="random", help="oracle | greedy | random")
    parser.add_argument("--checkpoint", default=None, help="trained agent checkpoint (.pt)")
    parser.add_argument("--suite", default="interp", choices=list(_SUITES))
    parser.add_argument("--episodes", type=int, default=100)
    parser.add_argument("--gif", default=None, help="also record one episode to this path")
    args = parser.parse_args()

    env = gym.make("atrium_sim/BrickLayer-v0",
                   render_mode="rgb_array" if args.gif else None)
    policy = load_policy(args, env)
    result = evaluate(env, policy, args.suite, args.episodes)
    result["policy"] = args.checkpoint or args.policy
    if args.gif:
        from atrium_sim.render.recorder import record_episode

        record_episode(env, policy, args.gif, seed=BASE_SEED)
        result["gif"] = args.gif
    env.close()
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()