From fcb1d8e104c72eafb60d1ae99fa18e88fa0e608a Mon Sep 17 00:00:00 2001
From: Dan Albert <dan@gingerhq.net>
Date: Fri, 1 Jan 2021 14:41:02 -0800
Subject: [PATCH] Retreat ground forces from captured bases.

If the captured base has no connection to other friendly objectives (the
base was encircled) then the enemy equipment will be captured and sold.
Otherwise the units will retreat toward connected bases.

Partial fix for https://github.com/Khopa/dcs_liberation/issues/693.
Still need to fix this for air units.
---
 changelog.md                 |  1 +
 game/theater/base.py         | 11 ++++++++
 game/theater/controlpoint.py | 55 ++++++++++++++++++++++++++++++++++--
 3 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/changelog.md b/changelog.md
index 33ca6919..9fb2b7d1 100644
--- a/changelog.md
+++ b/changelog.md
@@ -14,6 +14,7 @@ Saves from 2.3 are not compatible with 2.4.
 * **[Mission Generator]** Multiple groups are created for complex SAM sites (SAMs with additional point defense or SHORADS), improving Skynet behavior.
 * **[Skynet]** Point defenses are now configured to remain on to protect the site they accompany.
 * **[Balance]** Opfor now gains income using the same rules as the player, significantly increasing their income relative to the player for most campaigns.
+* **[Balance]** Ground units now retreat from captured bases when they are connected to a friendly base. Units with no retreat path will be captured and sold.
 * **[Economy]** FOBs generate only $10M per turn (previously $20M like airbases).
 * **[Economy]** Carriers and off-map spawns generate no income (previously $20M like airbases).
 * **[UI]** Multi-SAM objectives now show threat and detection rings per group.
diff --git a/game/theater/base.py b/game/theater/base.py
index 9f0ebe8f..88ede52f 100644
--- a/game/theater/base.py
+++ b/game/theater/base.py
@@ -9,6 +9,7 @@ from dcs.unittype import FlyingType, UnitType, VehicleType
 from dcs.vehicles import AirDefence, Armor
 
 from game import db
+from game.db import PRICES
 from gen.ground_forces.ai_ground_planner_db import TYPE_SHORAD
 
 STRENGTH_AA_ASSEMBLE_MIN = 0.2
@@ -37,6 +38,16 @@ class Base:
     def total_armor(self) -> int:
         return sum(self.armor.values())
 
+    @property
+    def total_armor_value(self) -> int:
+        total = 0
+        for unit_type, count in self.armor.items():
+            try:
+                total += PRICES[unit_type] * count
+            except KeyError:
+                logging.exception(f"No price found for {unit_type.id}")
+        return total
+
     @property
     def total_frontline_aa(self) -> int:
         return sum([v for k, v in self.armor.items() if k in TYPE_SHORAD])
diff --git a/game/theater/controlpoint.py b/game/theater/controlpoint.py
index f0bb53a3..57a0de52 100644
--- a/game/theater/controlpoint.py
+++ b/game/theater/controlpoint.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import heapq
 import itertools
 import logging
 import random
@@ -7,7 +8,8 @@ import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Dict, Iterator, List, Optional, TYPE_CHECKING, Type
+from functools import total_ordering
+from typing import Any, Dict, Iterator, List, Optional, TYPE_CHECKING, Type
 
 from dcs.mapping import Point
 from dcs.ships import (
@@ -192,6 +194,28 @@ class RunwayStatus:
         return f"Runway repairing, {turns_remaining} turns remaining"
 
 
+@total_ordering
+class GroundUnitDestination:
+    def __init__(self, control_point: ControlPoint) -> None:
+        self.control_point = control_point
+
+    @property
+    def total_value(self) -> float:
+        return self.control_point.base.total_armor_value
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, GroundUnitDestination):
+            raise TypeError
+
+        return self.total_value == other.total_value
+
+    def __lt__(self, other: Any) -> bool:
+        if not isinstance(other, GroundUnitDestination):
+            raise TypeError
+
+        return self.total_value < other.total_value
+
+
 class ControlPoint(MissionTarget, ABC):
 
     position = None  # type: Point
@@ -365,9 +389,37 @@ class ControlPoint(MissionTarget, ABC):
                     base_defense.position)
         self.base_defenses = []
 
+    def capture_equipment(self, game: Game) -> None:
+        total = self.base.total_armor_value
+        self.base.armor.clear()
+        game.adjust_budget(total, player=not self.captured)
+        game.message(
+            f"{self.name} is not connected to any friendly points. Ground "
+            f"vehicles have been captured and sold for ${total}M.")
+
+    def retreat_ground_units(self, game: Game):
+        # When there are multiple valid destinations, deliver units to whichever
+        # base is least defended first. The closest approximation of unit
+        # strength we have is price
+        destinations = [GroundUnitDestination(cp)
+                        for cp in self.connected_points
+                        if cp.captured == self.captured]
+        if not destinations:
+            self.capture_equipment(game)
+            return
+
+        heapq.heapify(destinations)
+        destination = heapq.heappop(destinations)
+        while self.base.armor:
+            unit_type, count = self.base.armor.popitem()
+            for _ in range(count):
+                destination.control_point.base.commision_units({unit_type: 1})
+                destination = heapq.heappushpop(destinations, destination)
+
     # TODO: Should be Airbase specific.
     def capture(self, game: Game, for_player: bool) -> None:
         self.pending_unit_deliveries.refund_all(game)
+        self.retreat_ground_units(game)
 
         if for_player:
             self.captured = True
@@ -377,7 +429,6 @@ class ControlPoint(MissionTarget, ABC):
         self.base.set_strength_to_minimum()
 
         self.base.aircraft = {}
-        self.base.armor = {}
 
         self.clear_base_defenses()
         from .start_generator import BaseDefenseGenerator