nodes/ai.py

Part of Tanks of Freedom.

  1"""Simple turn-based AI for the red side.
  2
  3Heuristic priority each unit:
  4  1. Adjacent enemy it is allowed to shoot → attack the weakest.
  5  2. Infantry with neutral or enemy building reachable → move to capture.
  6  3. Otherwise advance toward the enemy HQ along a BFS path, stopping at
  7     the furthest reachable cell within AP budget.
  8
  9The AI yields one action per call to ``step()``. The world drives it
 10incrementally so the player can watch enemy movement.
 11"""
 12
 13from __future__ import annotations
 14
 15from .data import ATTACK_RANGE, PLAYER_BLUE, UNIT_SOLDIER, can_attack_unit_type
 16from .pathfinder import cells_within_range, find_path, reachable_cells
 17
 18
 19class AIController:
 20    """Stateful AI iterator. ``step()`` returns one of:
 21    - ``("attack", attacker, defender)``
 22    - ``("move", unit, path)``: path is a list of cells incl. start
 23    - ``None``: no more actions this turn.
 24    """
 25
 26    def __init__(self, world):
 27        self._world = world
 28        self._pending_units: list = []
 29        self._began_turn = False
 30
 31    def begin_turn(self) -> None:
 32        self._began_turn = True
 33        # Snapshot current AI units (avoid iterating while units may die).
 34        self._pending_units = [u for u in self._world.units if u.owner != PLAYER_BLUE]
 35
 36    def step(self):
 37        """Pick one action. Returns the action or None when done."""
 38        if not self._began_turn:
 39            return None
 40        # Drop dead/destroyed units.
 41        self._pending_units = [u for u in self._pending_units if u.life > 0 and u in self._world.units]
 42        for u in self._pending_units:
 43            if u.ap <= 0:
 44                continue
 45            action = self._plan_for(u)
 46            if action is not None:
 47                return action
 48        # All AI units have acted (or have no AP): turn over.
 49        self._began_turn = False
 50        self._pending_units = []
 51        return None
 52
 53    # ---------------------------------------------------------- planners
 54    def _plan_for(self, unit):
 55        # 1) Attack adjacent enemy if possible.
 56        if unit.can_attack():
 57            target = self._best_target_in_range(unit)
 58            if target is not None:
 59                return ("attack", unit, target)
 60
 61        # 2) Move toward best objective.
 62        path = self._best_move(unit)
 63        if path is not None and len(path) > 1:
 64            return ("move", unit, path)
 65
 66        # 3) Nothing useful. Skip.
 67        unit.ap = 0
 68        return None
 69
 70    def _best_target_in_range(self, unit):
 71        candidates = cells_within_range(unit.cell, ATTACK_RANGE[unit.type])
 72        best = None
 73        best_score = None
 74        for cell in candidates:
 75            other = self._world.unit_at(cell)
 76            if other is None or other.owner == unit.owner:
 77                continue
 78            if not can_attack_unit_type(unit.type, other.type):
 79                continue
 80            # Prefer lowest HP.
 81            score = other.life
 82            if best_score is None or score < best_score:
 83                best = other
 84                best_score = score
 85        return best
 86
 87    def _best_move(self, unit):
 88        # Pick a goal: closest enemy HQ takes priority for non-infantry,
 89        # closest neutral building for infantry, then nearest enemy unit.
 90        goals = []
 91        if unit.type == UNIT_SOLDIER:
 92            for b in self._world.buildings:
 93                if b.owner != unit.owner:
 94                    goals.append((b.cell, 0 if b.is_neutral() else 1))
 95        # Always consider the enemy HQ.
 96        for b in self._world.buildings:
 97            if b.is_hq and b.owner == PLAYER_BLUE:
 98                goals.append((b.cell, 2))
 99        # Plus closest enemy unit (for tank/heli).
100        for u in self._world.units:
101            if u.owner != unit.owner:
102                goals.append((u.cell, 3))
103
104        if not goals:
105            return None
106
107        # Pre-compute reachable so we can clamp the goal cell to within AP.
108        def passable(x, y):
109            return self._world.is_passable_for(unit, x, y)
110
111        def blocked(x, y):
112            return self._world.unit_at((x, y)) is not None and (x, y) != unit.cell
113
114        reach = reachable_cells(unit.cell, unit.ap, passable=passable, blocked=blocked)
115
116        # Sort goals by Manhattan distance to current cell, prefer lower priority value.
117        def goal_key(g):
118            cell, prio = g
119            cx, cy = cell
120            ux, uy = unit.cell
121            return (prio, abs(cx - ux) + abs(cy - uy))
122
123        goals.sort(key=goal_key)
124
125        for goal_cell, _ in goals:
126            # Try to find a real path; truncate to within reach.
127            path = find_path(unit.cell, goal_cell, passable=passable, blocked=blocked)
128            if not path or len(path) == 1:
129                continue
130            # Truncate to the last cell that is in ``reach``.
131            best_idx = 0
132            for i, c in enumerate(path):
133                if c in reach:
134                    best_idx = i
135                else:
136                    break
137            if best_idx == 0:
138                continue
139            return path[: best_idx + 1]
140        return None