evaluate.py

from __future__ import annotations

import argparse
import os
from pathlib import Path

# Disable tqdm's rich integration BEFORE any imports to avoid shutdown errors
# This must be set before tqdm is imported by any module
os.environ["TQDM_DISABLE"] = "1"

import numpy as np

from sb3_contrib import MaskablePPO
from sb3_contrib.common.wrappers import ActionMasker

from gomoku_rl import GomokuEnv, HeuristicOpponent, PolicyGuidedOpponent, RandomOpponent
from gomoku_rl.opponents import predict_masked_action


def mask_fn(env: GomokuEnv):
    return env.action_masks()


def build_opponent(
    name: str,
    seed: int,
    heuristic_search_depth: int = 2,
    heuristic_candidate_radius: int = 2,
    heuristic_max_candidates: int = 10,
    heuristic_early_max_candidates: int = 14,
):
    if name == "random":
        return RandomOpponent(seed=seed)
    if name == "heuristic":
        return HeuristicOpponent(
            seed=seed,
            search_depth=heuristic_search_depth,
            candidate_radius=heuristic_candidate_radius,
            max_candidates=heuristic_max_candidates,
            early_max_candidates=heuristic_early_max_candidates,
        )
    raise ValueError(f"Unsupported opponent: {name}")


def build_agent(
    model: MaskablePPO,
    agent_mode: str,
    seed: int,
    search_depth: int,
    candidate_radius: int,
    max_candidates: int,
    policy_weight: float,
):
    if agent_mode == "model":
        return model
    if agent_mode == "hybrid":
        return PolicyGuidedOpponent(
            model=model,
            seed=seed,
            search_depth=search_depth,
            candidate_radius=candidate_radius,
            max_candidates=max_candidates,
            early_max_candidates=max(max_candidates + 4, max_candidates),
            policy_weight=policy_weight,
        )
    raise ValueError(f"Unsupported agent mode: {agent_mode}")


def play_one_game(model: MaskablePPO, agent, env: ActionMasker, seed: int) -> str:
    env.reset(seed=seed)
    done = False
    info: dict[str, bool] = {}

    while not done:
        board = env.env.board
        if isinstance(agent, MaskablePPO):
            action = predict_masked_action(agent, board, player=1, deterministic=True)
        else:
            action = agent.choose_action(board.copy(), player=1, win_length=env.env.win_length)
        _, _, terminated, truncated, info = env.step(action)
        done = terminated or truncated

    if info.get("agent_won"):
        return "win"
    if info.get("opponent_won"):
        return "loss"
    if info.get("draw"):
        return "draw"
    raise RuntimeError(f"Unexpected terminal info: {info}")


def evaluate(
    model_path: Path,
    opponent_name: str,
    games: int,
    board_size: int,
    win_length: int,
    seed: int,
    agent_mode: str = "model",
    search_depth: int = 3,
    candidate_radius: int = 2,
    max_candidates: int = 12,
    policy_weight: float = 220_000.0,
    opponent_search_depth: int = 2,
    opponent_candidate_radius: int = 2,
    opponent_max_candidates: int = 10,
    opponent_early_max_candidates: int = 14,
) -> dict[str, str | int | float]:
    opponent = build_opponent(
        opponent_name,
        seed,
        heuristic_search_depth=opponent_search_depth,
        heuristic_candidate_radius=opponent_candidate_radius,
        heuristic_max_candidates=opponent_max_candidates,
        heuristic_early_max_candidates=opponent_early_max_candidates,
    )
    env = GomokuEnv(
        board_size=board_size,
        win_length=win_length,
        opponent=opponent,
        opponent_starts_prob=0.5,
        seed=seed,
    )
    wrapped_env = ActionMasker(env, mask_fn)
    model = MaskablePPO.load(model_path)
    agent = build_agent(
        model=model,
        agent_mode=agent_mode,
        seed=seed,
        search_depth=search_depth,
        candidate_radius=candidate_radius,
        max_candidates=max_candidates,
        policy_weight=policy_weight,
    )

    wins = 0
    losses = 0
    draws = 0

    for game_idx in range(games):
        result = play_one_game(model, agent, wrapped_env, seed + game_idx)
        if result == "win":
            wins += 1
        elif result == "loss":
            losses += 1
        else:
            draws += 1

    total = wins + losses + draws
    stats: dict[str, str | int | float] = {
        "model_path": str(model_path),
        "opponent": opponent_name,
        "agent_mode": agent_mode,
        "games": total,
        "wins": wins,
        "losses": losses,
        "draws": draws,
        "win_rate": wins / total,
        "loss_rate": losses / total,
        "draw_rate": draws / total,
    }

    print(f"Model: {model_path}")
    print(f"Agent mode: {agent_mode}")
    print(f"Opponent: {opponent_name}")
    if opponent_name == "heuristic":
        print(
            "Opponent heuristic config: "
            f"depth={opponent_search_depth}, "
            f"radius={opponent_candidate_radius}, "
            f"max_candidates={opponent_max_candidates}, "
            f"early_max_candidates={opponent_early_max_candidates}"
        )
    print(f"Games: {total}")
    print(f"Wins: {wins}")
    print(f"Losses: {losses}")
    print(f"Draws: {draws}")
    print(f"Win rate: {stats['win_rate']:.3%}")
    print(f"Loss rate: {stats['loss_rate']:.3%}")
    print(f"Draw rate: {stats['draw_rate']:.3%}")
    return stats


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Evaluate a trained Gomoku MaskablePPO model.")
    parser.add_argument("--model-path", type=Path, required=True)
    parser.add_argument("--opponent", choices=["random", "heuristic"], required=True)
    parser.add_argument("--agent-mode", choices=["model", "hybrid"], default="model")
    parser.add_argument("--games", type=int, default=200)
    parser.add_argument("--board-size", type=int, default=9)
    parser.add_argument("--win-length", type=int, default=5)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--search-depth", type=int, default=3)
    parser.add_argument("--candidate-radius", type=int, default=2)
    parser.add_argument("--max-candidates", type=int, default=12)
    parser.add_argument("--policy-weight", type=float, default=220_000.0)
    parser.add_argument("--opponent-search-depth", type=int, default=2)
    parser.add_argument("--opponent-candidate-radius", type=int, default=2)
    parser.add_argument("--opponent-max-candidates", type=int, default=10)
    parser.add_argument("--opponent-early-max-candidates", type=int, default=14)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    evaluate(
        model_path=args.model_path,
        opponent_name=args.opponent,
        games=args.games,
        board_size=args.board_size,
        win_length=args.win_length,
        seed=args.seed,
        agent_mode=args.agent_mode,
        search_depth=args.search_depth,
        candidate_radius=args.candidate_radius,
        max_candidates=args.max_candidates,
        policy_weight=args.policy_weight,
        opponent_search_depth=args.opponent_search_depth,
        opponent_candidate_radius=args.opponent_candidate_radius,
        opponent_max_candidates=args.opponent_max_candidates,
        opponent_early_max_candidates=args.opponent_early_max_candidates,
    )


if __name__ == "__main__":
    main()