engine.py

#!/usr/bin/env python3
"""
Ants Game Engine
================

A deterministic, turn-based multi-agent RTS for CodeClash, in the spirit of the 2010
Google/Waterloo Ants AI Challenge.

Rules (fog of war, focus-fire combat, food->spawn, hill razing, the squared view/attack/
spawn radii) are reimplemented from the original challenge. Canonical reference / source:
https://github.com/aichallenge/aichallenge  (the ``ants/`` game there). This engine is a
clean-room reimplementation adapted for CodeClash: deterministic tick-synced turns
(vs. the original's real-time TCP), symmetric seeded maps, and simplified food spawning.

Each player commands a swarm of ants on a toroidal (wrap-around) grid. Every turn all
ants move one cell simultaneously. You only see the map within your ants' view radius
(fog of war). Combat is "focus fire": a locally outnumbered ant dies. Food adjacent to
a single player's ants is gathered and spawns a new ant on one of that player's hills.
Move an ant onto an enemy hill to raze it -- razing enemy hills is how you win.

The simulation is deterministic given (seed, bot code): maps, water and food are seeded
(so sims vary but replay exactly), and the whole spectator state is recorded per turn.

Usage:
    python engine.py /path/to/p1/main.py /path/to/p2/main.py -r NUM_GAMES -o OUTPUT_DIR

The bot interface (see README.md and main.py):

    def do_turn(obs: dict) -> list:
        # return a list of [row, col, dir] moves; dir is "N" "S" "E" "W"
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import math
import os
import random
import signal
import sys
from typing import Callable

# --------------------------------------------------------------------------------------
# Game constants
# --------------------------------------------------------------------------------------

ROWS = 32                    # must be even (translational symmetry uses ROWS//2)
COLS = 32                    # must be even
VIEWRADIUS2 = 77             # squared view radius (fog of war)
ATTACKRADIUS2 = 5            # squared attack radius (combat)
SPAWNRADIUS2 = 1             # squared radius for gathering food
MAX_TURNS = 500
TURN_TIMEOUT = 0.30          # seconds a single do_turn call may take before it's dropped
WATER_DENSITY = 0.06         # fraction of cells seeded as water (symmetric)
FOOD_TARGET = 12             # engine tops food up toward this many on the board

DIRS: dict[str, tuple[int, int]] = {"N": (-1, 0), "S": (1, 0), "E": (0, 1), "W": (0, -1)}


# --------------------------------------------------------------------------------------
# Bot loading + sandboxed invocation
# --------------------------------------------------------------------------------------


class TimeoutError_(Exception):
    pass


def _on_alarm(signum, frame):  # noqa: ARG001
    raise TimeoutError_()


_HAS_ALARM = hasattr(signal, "SIGALRM")
if _HAS_ALARM:
    signal.signal(signal.SIGALRM, _on_alarm)

_module_counter = 0


def load_bot(path: str) -> Callable:
    global _module_counter
    _module_counter += 1
    module_name = f"bot_module_{_module_counter}"
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    if not hasattr(module, "do_turn"):
        raise ValueError(f"Bot module {path} must define a do_turn(obs) function")
    return module.do_turn


def call_bot(fn: Callable, obs: dict) -> list:
    """Call a bot's do_turn with a crash/timeout guard. Any failure -> no moves."""
    if _HAS_ALARM:
        signal.setitimer(signal.ITIMER_REAL, TURN_TIMEOUT)
    try:
        moves = fn(obs)
    except Exception:
        return []
    finally:
        if _HAS_ALARM:
            signal.setitimer(signal.ITIMER_REAL, 0)
    return moves if isinstance(moves, list) else []


# --------------------------------------------------------------------------------------
# Geometry helpers (toroidal)
# --------------------------------------------------------------------------------------


def wrap(r: int, c: int) -> tuple[int, int]:
    return (r % ROWS, c % COLS)


def dist2(a: tuple[int, int], b: tuple[int, int]) -> int:
    dr = abs(a[0] - b[0])
    dr = min(dr, ROWS - dr)
    dc = abs(a[1] - b[1])
    dc = min(dc, COLS - dc)
    return dr * dr + dc * dc


def _disk_offsets(r2: int) -> list[tuple[int, int]]:
    rad = int(math.isqrt(r2))
    return [
        (dr, dc)
        for dr in range(-rad, rad + 1)
        for dc in range(-rad, rad + 1)
        if dr * dr + dc * dc <= r2
    ]


VIEW_OFFSETS = _disk_offsets(VIEWRADIUS2)


# --------------------------------------------------------------------------------------
# Simulation
# --------------------------------------------------------------------------------------


class Game:
    def __init__(self, num_players: int, seed: int):
        self.n = num_players
        self.rng = random.Random(seed)
        self.turn = 0
        self.water: set[tuple[int, int]] = set()
        self.hills: dict[tuple[int, int], int] = {}   # (r,c) -> owner (living hills only)
        self.ants: dict[tuple[int, int], int] = {}    # (r,c) -> owner (<=1 ant per cell)
        self.food: set[tuple[int, int]] = set()
        self.razed: dict[int, int] = {i: 0 for i in range(num_players)}
        self.gathered: dict[int, int] = {i: 0 for i in range(num_players)}
        self._build_map()

    def _partner(self, r: int, c: int) -> tuple[int, int]:
        """The symmetric counterpart of a cell (translational half-shift on the torus)."""
        return ((r + ROWS // 2) % ROWS, (c + COLS // 2) % COLS)

    def _build_map(self) -> None:
        # Symmetric water: seed water in the whole grid but mirror every cell to its
        # partner, so the map is identical for both players (fair). Retry until hills
        # are mutually reachable.
        hill0 = None
        for _ in range(40):
            water: set[tuple[int, int]] = set()
            n_water = int(WATER_DENSITY * ROWS * COLS)
            cells = [(r, c) for r in range(ROWS) for c in range(COLS)]
            self.rng.shuffle(cells)
            for r, c in cells:
                if len(water) >= n_water:
                    break
                p = self._partner(r, c)
                if (r, c) == p:  # self-symmetric cell (shouldn't happen for even dims)
                    continue
                water.add((r, c))
                water.add(p)

            # Place player 0's hill on open land with a clear surrounding, partner = p1.
            open_cells = [(r, c) for r in range(ROWS) for c in range(COLS) if (r, c) not in water]
            self.rng.shuffle(open_cells)
            hill0 = None
            for r, c in open_cells:
                if all((wrap(r + dr, c + dc)) not in water for dr in (-1, 0, 1) for dc in (-1, 0, 1)):
                    hill0 = (r, c)
                    break
            if hill0 is None:
                continue

            # Clear any water on the hill cells themselves / their partners.
            hills = {}
            for i in range(self.n):
                # Spread hills evenly along the torus diagonal: player i sits a fraction
                # i/n of the way around the grid from hill0. For 2 players this is exactly
                # hill0 and its partner (the half-shift), so 2-player maps stay symmetric.
                hr, hc = wrap(hill0[0] + i * (ROWS // self.n), hill0[1] + i * (COLS // self.n))
                water.discard((hr, hc))
                hills[(hr, hc)] = i

            if self._connected(water, list(hills.keys())):
                self.water = water
                self.hills = hills
                break
        else:
            # Fallback: no water at all, hills placed deterministically.
            self.water = set()
            self.hills = {(ROWS // 4, COLS // 4): 0, self._partner(ROWS // 4, COLS // 4): 1}

        # One ant on each hill to start.
        for (hr, hc), owner in self.hills.items():
            self.ants[(hr, hc)] = owner

        # Seed initial food and top up.
        self._spawn_food()

    def _connected(self, water: set[tuple[int, int]], points: list[tuple[int, int]]) -> bool:
        start = points[0]
        seen = {start}
        stack = [start]
        while stack:
            r, c = stack.pop()
            for dr, dc in DIRS.values():
                nr, nc = wrap(r + dr, c + dc)
                if (nr, nc) not in water and (nr, nc) not in seen:
                    seen.add((nr, nc))
                    stack.append((nr, nc))
        return all(p in seen for p in points)

    def _spawn_food(self) -> None:
        """Top the board up toward FOOD_TARGET, adding food in symmetric pairs."""
        blocked = self.water | set(self.ants) | set(self.hills) | self.food
        attempts = 0
        while len(self.food) < FOOD_TARGET and attempts < 200:
            attempts += 1
            r = self.rng.randrange(ROWS)
            c = self.rng.randrange(COLS)
            p = self._partner(r, c)
            if (r, c) in blocked or p in blocked or (r, c) == p:
                continue
            self.food.add((r, c))
            self.food.add(p)
            blocked.add((r, c))
            blocked.add(p)

    # -- observation -------------------------------------------------------------------

    def observation(self, me: int) -> dict:
        visible: set[tuple[int, int]] = set()
        for (r, c), owner in self.ants.items():
            if owner == me:
                for dr, dc in VIEW_OFFSETS:
                    visible.add(wrap(r + dr, c + dc))
        return {
            "turn": self.turn,
            "max_turns": MAX_TURNS,
            "rows": ROWS,
            "cols": COLS,
            "viewradius2": VIEWRADIUS2,
            "attackradius2": ATTACKRADIUS2,
            "spawnradius2": SPAWNRADIUS2,
            "you": me,
            "my_ants": [[r, c] for (r, c), o in self.ants.items() if o == me],
            "my_hills": [[r, c] for (r, c), o in self.hills.items() if o == me],
            "enemy_ants": [[r, c, o] for (r, c), o in self.ants.items() if o != me and (r, c) in visible],
            "enemy_hills": [[r, c, o] for (r, c), o in self.hills.items() if o != me and (r, c) in visible],
            "food": [[r, c] for (r, c) in self.food if (r, c) in visible],
            "water": [[r, c] for (r, c) in self.water if (r, c) in visible],
        }

    # -- stepping ----------------------------------------------------------------------

    def _apply_moves(self, orders: dict[int, list]) -> None:
        # First valid direction per ant; ants with no order stay put.
        chosen: dict[tuple[int, int], str] = {}
        for owner, moves in orders.items():
            for mv in moves:
                try:
                    r, c, d = int(mv[0]), int(mv[1]), str(mv[2]).strip().upper()
                except (ValueError, IndexError, TypeError):
                    continue
                pos = (r % ROWS, c % COLS)
                if self.ants.get(pos) == owner and pos not in chosen and d in DIRS:
                    chosen[pos] = d

        # Compute destinations (water blocks -> stay).
        arrivals: dict[tuple[int, int], list[int]] = {}
        for (r, c), owner in self.ants.items():
            d = chosen.get((r, c))
            if d:
                dr, dc = DIRS[d]
                nr, nc = wrap(r + dr, c + dc)
                if (nr, nc) in self.water:
                    nr, nc = r, c
            else:
                nr, nc = r, c
            arrivals.setdefault((nr, nc), []).append(owner)

        # Any cell that 2+ ants arrive on: all of them die (collision).
        self.ants = {cell: owners[0] for cell, owners in arrivals.items() if len(owners) == 1}

    def _resolve_combat(self) -> None:
        antlist = list(self.ants.items())
        atk: dict[tuple[int, int], int] = {}
        for pos, owner in antlist:
            atk[pos] = sum(
                1 for p2, o2 in antlist if o2 != owner and dist2(pos, p2) <= ATTACKRADIUS2
            )
        dead = set()
        for pos, owner in antlist:
            for p2, o2 in antlist:
                if o2 != owner and dist2(pos, p2) <= ATTACKRADIUS2 and atk[p2] <= atk[pos]:
                    dead.add(pos)
                    break
        for pos in dead:
            del self.ants[pos]

    def _raze_hills(self) -> None:
        for hill, owner in list(self.hills.items()):
            occupant = self.ants.get(hill)
            if occupant is not None and occupant != owner:
                self.razed[occupant] += 1
                del self.hills[hill]

    def _gather_food(self) -> None:
        for f in list(self.food):
            owners = {o for pos, o in self.ants.items() if dist2(f, pos) <= SPAWNRADIUS2}
            if len(owners) == 1:
                owner = next(iter(owners))
                self.food.discard(f)
                self.gathered[owner] += 1
                free = [h for h, ho in self.hills.items() if ho == owner and h not in self.ants]
                if free:
                    free.sort()
                    spawn = free[self.rng.randrange(len(free))]
                    self.ants[spawn] = owner
            # 0 owners -> nobody near; 2+ owners -> contested; either way leave the food.

    def step(self, orders: dict[int, list]) -> None:
        self._apply_moves(orders)
        self._resolve_combat()
        self._raze_hills()
        self._gather_food()
        self._spawn_food()

    def players_with_ants(self) -> set[int]:
        return set(self.ants.values())

    def ants_alive(self) -> dict[int, int]:
        counts = {i: 0 for i in range(self.n)}
        for o in self.ants.values():
            counts[o] += 1
        return counts

    def frame(self) -> dict:
        return {
            "t": self.turn,
            "ants": [[r, c, o] for (r, c), o in sorted(self.ants.items())],
            "hills": [[r, c, o] for (r, c), o in sorted(self.hills.items())],
            "food": [[r, c] for (r, c) in sorted(self.food)],
        }


def run_game(bot_paths: list[str], seed: int) -> dict:
    n = len(bot_paths)
    bots: list[Callable | None] = []
    load_errors: dict[int, str] = {}
    for i, path in enumerate(bot_paths):
        try:
            bots.append(load_bot(path))
        except Exception as e:
            bots.append(None)
            load_errors[i] = str(e)

    game = Game(num_players=n, seed=seed)
    frames = [game.frame()]

    for t in range(1, MAX_TURNS + 1):
        game.turn = t
        orders: dict[int, list] = {}
        for owner in game.players_with_ants():
            fn = bots[owner]
            orders[owner] = [] if fn is None else call_bot(fn, game.observation(owner))
        game.step(orders)
        frames.append(game.frame())
        if len(game.players_with_ants()) <= 1:
            break

    alive = game.ants_alive()
    # Rank by (hills razed, ants alive, food gathered); winner is the unique max.
    def key(i: int) -> tuple[int, int, int]:
        return (game.razed[i], alive[i], game.gathered[i])

    best = max((key(i) for i in range(n)), default=(0, 0, 0))
    leaders = [i for i in range(n) if key(i) == best]
    winner = leaders[0] if len(leaders) == 1 else None

    return {
        "winner": winner,
        "load_errors": load_errors,
        "num_turns": game.turn,
        "razed": game.razed,
        "ants_alive": alive,
        "gathered": game.gathered,
        "replay": {
            "rows": ROWS,
            "cols": COLS,
            "num_players": n,
            "water": [[r, c] for (r, c) in sorted(game.water)],
            "frames": frames,
        },
    }


# --------------------------------------------------------------------------------------
# CLI / tournament driver
# --------------------------------------------------------------------------------------


def write_replay(result: dict, game_num: int, bot_paths: list[str], output_dir: str) -> None:
    rp = dict(result["replay"])
    rp["names"] = [
        os.path.basename(os.path.dirname(p)) or f"player{i + 1}" for i, p in enumerate(bot_paths)
    ]
    rp["winner"] = result["winner"]
    with open(os.path.join(output_dir, f"sim_{game_num}.json"), "w") as f:
        json.dump(rp, f)


def main() -> None:
    parser = argparse.ArgumentParser(description="Ants Game Engine")
    parser.add_argument("bots", nargs="+", help="Paths to bot files (main.py)")
    parser.add_argument("-r", "--rounds", type=int, default=10, help="Number of games")
    parser.add_argument("-o", "--output-dir", type=str, default=None, help="Replay output dir")
    parser.add_argument("-s", "--seed", type=int, default=0, help="Base RNG seed")
    args = parser.parse_args()

    bot_paths = args.bots
    n = len(bot_paths)
    names = [
        os.path.basename(os.path.dirname(p)) or f"player{i + 1}" for i, p in enumerate(bot_paths)
    ]

    wins = {i: 0 for i in range(n)}
    draws = 0

    print(f"Running {args.rounds} games between:")
    for i, p in enumerate(bot_paths):
        print(f"  Player {i + 1}: {p}  ({names[i]})")
    print()

    if args.output_dir:
        os.makedirs(args.output_dir, exist_ok=True)

    for g in range(args.rounds):
        result = run_game(bot_paths, seed=args.seed + g)
        if result["winner"] is None:
            draws += 1
            wtxt = "draw"
        else:
            wins[result["winner"]] += 1
            wtxt = f"Player {result['winner'] + 1}"
        razed = "  ".join(f"P{i + 1}:raze={result['razed'][i]},ants={result['ants_alive'][i]}" for i in range(n))
        print(f"Game {g + 1}: {wtxt} (after {result['num_turns']} turns)  [{razed}]")
        if result["load_errors"]:
            for i, err in result["load_errors"].items():
                print(f"    (Player {i + 1} failed to load: {err})")
        if args.output_dir:
            write_replay(result, g, bot_paths, args.output_dir)

    print()
    print("FINAL_RESULTS")
    for i in range(n):
        print(f"Bot_{i + 1}: {wins[i]} games won ({names[i]})")
    print(f"Draws: {draws}")


if __name__ == "__main__":
    main()