From 49211dae4f291a741ea7e8955a51c1efa643190c Mon Sep 17 00:00:00 2001
From: Raffson <Raffson@users.noreply.github.com>
Date: Sun, 15 Sep 2024 21:12:34 +0200
Subject: [PATCH] Autoplanner support for ArmedRecon

---
 .../tasks/compound/attackbattlepositions.py   |  4 +++
 game/commander/tasks/primitive/armedrecon.py  | 25 +++++++++++++++++++
 game/commander/theaterstate.py                |  3 +++
 3 files changed, 32 insertions(+)
 create mode 100644 game/commander/tasks/primitive/armedrecon.py

diff --git a/game/commander/tasks/compound/attackbattlepositions.py b/game/commander/tasks/compound/attackbattlepositions.py
index 2286a341..48a61c47 100644
--- a/game/commander/tasks/compound/attackbattlepositions.py
+++ b/game/commander/tasks/compound/attackbattlepositions.py
@@ -1,5 +1,6 @@
 from collections.abc import Iterator
 
+from game.commander.tasks.primitive.armedrecon import PlanArmedRecon
 from game.commander.tasks.primitive.bai import PlanBai
 from game.commander.theaterstate import TheaterState
 from game.htn import CompoundTask, Method
@@ -10,3 +11,6 @@ class AttackBattlePositions(CompoundTask[TheaterState]):
         for battle_positions in state.enemy_battle_positions.values():
             for battle_position in battle_positions.in_priority_order:
                 yield [PlanBai(battle_position)]
+        # Only plan against the 2 most important CPs
+        for cp in state.control_point_priority_queue[:2]:
+            yield [PlanArmedRecon(cp)]
diff --git a/game/commander/tasks/primitive/armedrecon.py b/game/commander/tasks/primitive/armedrecon.py
new file mode 100644
index 00000000..026a17e4
--- /dev/null
+++ b/game/commander/tasks/primitive/armedrecon.py
@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from game.ato.flighttype import FlightType
+from game.commander.tasks.packageplanningtask import PackagePlanningTask
+from game.commander.theaterstate import TheaterState
+from game.theater import ControlPoint
+
+
+@dataclass
+class PlanArmedRecon(PackagePlanningTask[ControlPoint]):
+    def preconditions_met(self, state: TheaterState) -> bool:
+        if self.target not in state.control_point_priority_queue:
+            return False
+        if not self.target_area_preconditions_met(state):
+            return False
+        return super().preconditions_met(state)
+
+    def apply_effects(self, state: TheaterState) -> None:
+        state.control_point_priority_queue.remove(self.target)
+
+    def propose_flights(self) -> None:
+        self.propose_flight(FlightType.ARMED_RECON, self.get_flight_size())
+        self.propose_common_escorts()
diff --git a/game/commander/theaterstate.py b/game/commander/theaterstate.py
index 9e045728..3d57b224 100644
--- a/game/commander/theaterstate.py
+++ b/game/commander/theaterstate.py
@@ -63,6 +63,7 @@ class TheaterState(WorldState["TheaterState"]):
     enemy_barcaps: list[ControlPoint]
     threat_zones: ThreatZones
     vulnerable_control_points: list[ControlPoint]
+    control_point_priority_queue: list[ControlPoint]
 
     def _rebuild_threat_zones(self) -> None:
         """Recreates the theater's threat zones based on the current planned state."""
@@ -137,6 +138,7 @@ class TheaterState(WorldState["TheaterState"]):
             threatening_air_defenses=self.threatening_air_defenses,
             detecting_air_defenses=self.detecting_air_defenses,
             vulnerable_control_points=self.vulnerable_control_points,
+            control_point_priority_queue=self.control_point_priority_queue,
         )
 
     @classmethod
@@ -198,4 +200,5 @@ class TheaterState(WorldState["TheaterState"]):
             enemy_barcaps=list(game.theater.control_points_for(not player)),
             threat_zones=game.threat_zone_for(not player),
             vulnerable_control_points=vulnerable_control_points,
+            control_point_priority_queue=ordered_capturable_points,
         )