"""Evaluate a trained PPO policy (or the Strömbom baseline) on the env.

Reports success rate and time-to-pen across a fixed seed grid for each
flock size 1..MAX_SHEEP. Used to produce the M5 quantitative comparison
table mentioned in plan.md.

Usage::

    python -m training.eval --policy training/runs/latest/best
    python -m training.eval --policy strombom
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
from statistics import mean, stdev

_HERE = os.path.dirname(os.path.abspath(__file__))
_PROJECT_ROOT = os.path.normpath(os.path.join(_HERE, ".."))
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)

import numpy as np

from herding.world.geometry import MAX_SHEEP, PEN_ENTRY
from herding.control.sequential import compute_action as sequential_action
from herding.control.strombom import compute_action as strombom_action
from training.herding_env import HerdingEnv


def rollout(env: HerdingEnv, predict_fn, max_steps: int) -> dict:
    obs, _ = env.reset()
    success = False
    for t in range(max_steps):
        action = predict_fn(env, obs)
        obs, _r, terminated, truncated, info = env.step(action)
        if terminated or truncated:
            success = bool(info.get("is_success", False))
            return {"success": success, "steps": info.get("steps", t + 1),
                    "n_penned": info.get("n_penned", 0)}
    return {"success": False, "steps": max_steps, "n_penned": int(env.sheep_penned.sum())}


def make_analytic_predictor(action_fn):
    def _predict(env, _obs):
        # Use whatever perception the env exposes — tracker output in
        # LiDAR mode, ground truth in privileged mode. This makes
        # evaluation honest: the analytic teacher sees what the
        # deployed controller would see.
        positions = env.perceived_positions()
        vx, vy, _mode = action_fn((env.dog_x, env.dog_y), positions, PEN_ENTRY)
        return np.array([vx, vy], dtype=np.float32)
    return _predict


# Backwards-compat alias.
def make_strombom_predictor():
    return make_analytic_predictor(strombom_action)


def make_policy_predictor(model, vecnorm):
    def _predict(_env, obs):
        if vecnorm is not None:
            obs_b = vecnorm.normalize_obs(np.asarray(obs, dtype=np.float32).reshape(1, -1))
        else:
            obs_b = np.asarray(obs, dtype=np.float32).reshape(1, -1)
        action, _ = model.predict(obs_b, deterministic=True)
        return action[0]
    return _predict


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--policy", required=True,
                        help="Either 'strombom' or path to an SB3 run directory.")
    parser.add_argument("--n-seeds", type=int, default=10)
    parser.add_argument("--max-steps", type=int, default=5000)
    parser.add_argument("--max-flock", type=int, default=MAX_SHEEP)
    # 1.0 = deployment distribution (sheep anywhere in field).
    # Lower values use the training-curriculum spawn band (sheep near gate).
    parser.add_argument("--difficulty", type=float, default=1.0)
    args = parser.parse_args()

    frame_stack = 1  # default; analytic predictors don't use stacked obs
    if args.policy == "strombom":
        predict = make_analytic_predictor(strombom_action)
    elif args.policy == "sequential":
        predict = make_analytic_predictor(sequential_action)
    else:
        from stable_baselines3 import PPO
        run = Path(args.policy)
        # Resolve to a zip: directory of checkpoints, or a direct zip path.
        if run.is_file():
            zip_path = run
        else:
            for name in ("best_model.zip", "policy.zip", "final.zip"):
                if (run / name).exists():
                    zip_path = run / name
                    break
            else:
                raise FileNotFoundError(
                    f"No checkpoint found in {run} (tried best_model.zip, "
                    f"policy.zip, final.zip)"
                )
        model = PPO.load(str(zip_path), device="auto")
        # Auto-detect frame stacking from the policy's expected obs dim,
        # so eval runs with whatever stacking the policy was trained on.
        from herding.obs import OBS_DIM as _SINGLE
        policy_obs_dim = int(model.observation_space.shape[0])
        if policy_obs_dim % _SINGLE == 0 and policy_obs_dim // _SINGLE >= 1:
            frame_stack = policy_obs_dim // _SINGLE
            if frame_stack > 1:
                print(f"[eval] policy expects frame_stack={frame_stack}")
        vecnorm = None
        vn_path = run / "vecnormalize.pkl"
        if not vn_path.exists() and run.parent.name != "best":
            vn_path = run.parent / "vecnormalize.pkl"
        if vn_path.exists():
            import pickle
            with open(vn_path, "rb") as f:
                vecnorm = pickle.load(f)
            vecnorm.training = False
            vecnorm.norm_reward = False
        predict = make_policy_predictor(model, vecnorm)

    print(f"{'n_sheep':>8} {'success%':>10} {'mean_steps':>12} {'mean_penned':>12}")
    print("-" * 46)
    for n in range(1, args.max_flock + 1):
        successes, steps, penned = [], [], []
        for seed in range(args.n_seeds):
            env = HerdingEnv(n_sheep=n, max_steps=args.max_steps,
                             difficulty=args.difficulty, seed=seed,
                             frame_stack=frame_stack)
            r = rollout(env, predict, args.max_steps)
            successes.append(int(r["success"]))
            steps.append(r["steps"])
            penned.append(r["n_penned"])
        sr = 100.0 * mean(successes)
        ms = mean(steps)
        mp = mean(penned)
        print(f"{n:>8d} {sr:>9.1f}% {ms:>12.0f} {mp:>12.2f}")


if __name__ == "__main__":
    main()