Refactor game generation.

No real functional improvements yet, but want to get this submitted to
unblock Skynet changes.

game/__init__.py

@@ -1,2 +1,3 @@
 from .game import Game
-from . import db
+from . import db
+from .version import VERSION

game/event/event.py

@@ -14,7 +14,6 @@ from game.infos.information import Information
 from game.operation.operation import Operation
 from gen.ground_forces.combat_stance import CombatStance
 from theater import ControlPoint
-from theater.start_generator import generate_airbase_defense_group
 
 if TYPE_CHECKING:
     from ..game import Game

game/version.py

@@ -0,0 +1,2 @@
+#: Current version of Liberation.
+VERSION = "2.2.0-preview"

gen/sam/group_generator.py

@@ -25,7 +25,7 @@ class GroupGenerator():
     def generate(self):
         raise NotImplementedError
 
-    def get_generated_group(self):
+    def get_generated_group(self) -> unitgroup.VehicleGroup:
         return self.vg
 
     def add_unit(self, unit_type, name, pos_x, pos_y, heading):

gen/sam/sam_group_generator.py

@@ -1,10 +1,10 @@
 import random
-from typing import List, Type
+from typing import List, Optional, Type
 
-from dcs.unittype import UnitType
+from dcs.unitgroup import VehicleGroup
 from dcs.vehicles import AirDefence
 
-from game import db
+from game import Game, db
 from gen.sam.aaa_bofors import BoforsGenerator
 from gen.sam.aaa_flak import FlakGenerator
 from gen.sam.aaa_zu23_insurgent import ZU23InsurgentGenerator
@@ -33,6 +33,7 @@ from gen.sam.sam_zsu23 import ZSU23Generator
 from gen.sam.sam_zu23 import ZU23Generator
 from gen.sam.sam_zu23_ural import ZU23UralGenerator
 from gen.sam.sam_zu23_ural_insurgent import ZU23UralInsurgentGenerator
+from theater import TheaterGroundObject
 
 SAM_MAP = {
     "HawkGenerator": HawkGenerator,
@@ -106,13 +107,14 @@ def get_faction_possible_sams_generator(faction: str) -> List[Type[GroupGenerato
     """
     return [SAM_MAP[s] for s in db.FACTIONS[faction].sams if s in SAM_MAP.keys()]
 
-def generate_anti_air_group(game, parent_cp, ground_object, faction:str):
+def generate_anti_air_group(game: Game, ground_object: TheaterGroundObject,
+                            faction: str) -> Optional[VehicleGroup]:
     """
     This generate a SAM group
-    :param parentCp: The parent control point
-    :param ground_object: The ground object which will own the sam group
-    :param country: Owner country
-    :return: Nothing, but put the group reference inside the ground object
+    :param game: The Game.
+    :param ground_object: The ground object which will own the sam group.
+    :param faction: Owner faction.
+    :return: The generated group, or None if one could not be generated.
     """
     possible_sams_generators = get_faction_possible_sams_generator(faction)
     if len(possible_sams_generators) > 0:
@@ -123,7 +125,8 @@ def generate_anti_air_group(game, parent_cp, ground_object, faction:str):
     return None
 
 
-def generate_shorad_group(game, parent_cp, ground_object, faction_name: str):
+def generate_shorad_group(game: Game, ground_object: TheaterGroundObject,
+                          faction_name: str) -> Optional[VehicleGroup]:
     faction = db.FACTIONS[faction_name]
 
     if len(faction.shorads) > 0:
@@ -132,9 +135,4 @@ def generate_shorad_group(game, parent_cp, ground_object, faction_name: str):
         generator.generate()
         return generator.get_generated_group()
     else:
-        return generate_anti_air_group(game, parent_cp, ground_object, faction_name)
-
-
-
-
-
+        return generate_anti_air_group(game, ground_object, faction_name)

qt_ui/main.py

@@ -7,7 +7,7 @@ from PySide2 import QtWidgets
 from PySide2.QtGui import QPixmap
 from PySide2.QtWidgets import QApplication, QSplashScreen
 
-from game import persistency
+from game import persistency, VERSION
 from qt_ui import (
     liberation_install,
     liberation_theme,
@@ -20,7 +20,7 @@ from qt_ui.windows.preferences.QLiberationFirstStartWindow import \
     QLiberationFirstStartWindow
 
 # Logging setup
-logging_config.init_logging(uiconstants.VERSION_STRING)
+logging_config.init_logging(VERSION)
 
 if __name__ == "__main__":
 

qt_ui/uiconstants.py

@@ -6,8 +6,6 @@ from PySide2.QtGui import QColor, QFont, QPixmap
 from theater.theatergroundobject import CATEGORY_MAP
 from .liberation_theme import get_theme_icons
 
-VERSION_STRING = "2.2.0-preview"
-
 URLS : Dict[str, str] = {
     "Manual": "https://github.com/khopa/dcs_liberation/wiki",
     "Repository": "https://github.com/khopa/dcs_liberation",

qt_ui/windows/QLiberationWindow.py

@@ -1,7 +1,6 @@
 import logging
-import sys
 import webbrowser
-from typing import Optional, Union
+from typing import Optional
 
 from PySide2.QtCore import Qt
 from PySide2.QtGui import QCloseEvent, QIcon
@@ -10,14 +9,14 @@ from PySide2.QtWidgets import (
     QActionGroup, QDesktopWidget,
     QFileDialog,
     QMainWindow,
-    QMenu, QMessageBox,
+    QMessageBox,
     QSplitter,
     QVBoxLayout,
     QWidget,
 )
 
 import qt_ui.uiconstants as CONST
-from game import Game, persistency
+from game import Game, persistency, VERSION
 from qt_ui.dialogs import Dialog
 from qt_ui.displayoptions import DisplayGroup, DisplayOptions, DisplayRule
 from qt_ui.models import GameModel
@@ -47,7 +46,7 @@ class QLiberationWindow(QMainWindow):
         self.setGame(persistency.restore_game())
 
         self.setGeometry(300, 100, 270, 100)
-        self.setWindowTitle("DCS Liberation - v" + CONST.VERSION_STRING)
+        self.setWindowTitle(f"DCS Liberation - v{VERSION}")
         self.setWindowIcon(QIcon("./resources/icon.png"))
         self.statusBar().showMessage('Ready')
 
@@ -225,7 +224,7 @@ class QLiberationWindow(QMainWindow):
             self.liberation_map.setGame(game)
 
     def showAboutDialog(self):
-        text = "<h3>DCS Liberation " + CONST.VERSION_STRING + "</h3>" + \
+        text = "<h3>DCS Liberation " + VERSION + "</h3>" + \
                "<b>Source code :</b> https://github.com/khopa/dcs_liberation" + \
                "<h4>Authors</h4>" + \
                "<p>DCS Liberation was originally developed by <b>shdwp</b>, DCS Liberation 2.0 is a partial rewrite based on this work by <b>Khopa</b>." \

qt_ui/windows/newgame/QNewGameWizard.py

@@ -1,6 +1,5 @@
 from __future__ import unicode_literals
 
-import datetime
 import logging
 from typing import List, Optional
 
@@ -10,15 +9,14 @@ from PySide2.QtWidgets import QVBoxLayout
 from dcs.task import CAP, CAS
 
 import qt_ui.uiconstants as CONST
-from game import Game, db
+from game import db
 from game.settings import Settings
-from gen import namegen
 from qt_ui.windows.newgame.QCampaignList import (
     Campaign,
     QCampaignList,
     load_campaigns,
 )
-from theater import ConflictTheater, start_generator
+from theater.start_generator import GameGenerator
 
 
 class NewGameWizard(QtWidgets.QWizard):
@@ -76,39 +74,13 @@ class NewGameWizard(QtWidgets.QWizard):
         settings.do_not_generate_player_navy = no_player_navy
         settings.do_not_generate_enemy_navy = no_enemy_navy
 
-        self.generatedGame = self.start_new_game(player_name, enemy_name, conflictTheater, midGame, multiplier,
-                                                 timePeriod, settings, starting_money)
+        generator = GameGenerator(player_name, enemy_name, conflictTheater,
+                                  settings, timePeriod, starting_money,
+                                  multiplier, midGame)
+        self.generatedGame = generator.generate()
 
         super(NewGameWizard, self).accept()
 
-    def start_new_game(self, player_name: str, enemy_name: str, conflictTheater: ConflictTheater,
-                       midgame: bool, multiplier: float, period: datetime, settings:Settings, starting_money: int):
-
-        # Reset name generator
-        namegen.reset()
-        start_generator.prepare_theater(conflictTheater, settings, midgame)
-
-        print("-- Starting New Game Generator")
-        print("Enemy name : " + enemy_name)
-        print("Player name : " + player_name)
-        print("Midgame : " + str(midgame))
-        start_generator.generate_initial_units(conflictTheater, enemy_name, True, multiplier)
-
-        print("-- Initial units generated")
-        game = Game(player_name=player_name,
-                    enemy_name=enemy_name,
-                    theater=conflictTheater,
-                    start_date=period,
-                    settings=settings)
-
-        print("-- Game Object generated")
-        start_generator.generate_groundobjects(conflictTheater, game)
-        game.budget = starting_money
-        game.settings.multiplier = multiplier
-        game.settings.sams = True
-        game.settings.version = CONST.VERSION_STRING
-        return game
-
 
 class IntroPage(QtWidgets.QWizardPage):
     def __init__(self, parent=None):

theater/controlpoint.py

@@ -35,7 +35,6 @@ class ControlPoint(MissionTarget):
 
     position = None  # type: Point
     name = None  # type: str
-    allow_sea_units = True
 
     captured = False
     has_frontline = True
@@ -47,10 +46,9 @@ class ControlPoint(MissionTarget):
                  at: db.StartingPosition, radials: List[int], size: int,
                  importance: float, has_frontline=True,
                  cptype=ControlPointType.AIRBASE):
+        super().__init__(" ".join(re.split(r" |-", name)[:2]), position)
         self.id = id
-        self.name = " ".join(re.split(r" |-", name)[:2])
         self.full_name = name
-        self.position: Point = position
         self.at = at
         self.ground_objects: List[TheaterGroundObject] = []
 
@@ -228,10 +226,8 @@ class ControlPoint(MissionTarget):
 
         # Handle cyclic dependency.
         from .start_generator import generate_airbase_defense_group
-        airbase_def_id = 0
-        for ground_object in self.ground_objects:
+        for idx, ground_object in enumerate(self.ground_objects):
             ground_object.groups = []
             if ground_object.airbase_group and faction_name != "":
-                generate_airbase_defense_group(airbase_def_id, ground_object,
-                                               faction_name, game, self)
-                airbase_def_id = airbase_def_id + 1
+                generate_airbase_defense_group(idx, ground_object,
+                                               faction_name, game)

theater/frontline.py

@@ -4,11 +4,27 @@ from typing import Tuple
 from dcs.mapping import Point
 from . import ControlPoint, MissionTarget
 
-
 # TODO: Dedup by moving everything to using this class.
 FRONTLINE_MIN_CP_DISTANCE = 5000
 
 
+def compute_position(control_point_a: ControlPoint,
+                     control_point_b: ControlPoint) -> Point:
+    a = control_point_a.position
+    b = control_point_b.position
+    attack_heading = a.heading_between_point(b)
+    attack_distance = a.distance_to_point(b)
+    middle_point = a.point_from_heading(attack_heading, attack_distance / 2)
+
+    strength_delta = float(control_point_a.base.strength -
+                           control_point_b.base.strength)
+    position = middle_point.point_from_heading(attack_heading,
+                                               strength_delta *
+                                               attack_distance / 2 -
+                                               FRONTLINE_MIN_CP_DISTANCE)
+    return position
+
+
 class FrontLine(MissionTarget):
     """Defines a front line location between two control points.
 
@@ -17,6 +33,8 @@ class FrontLine(MissionTarget):
 
     def __init__(self, control_point_a: ControlPoint,
                  control_point_b: ControlPoint) -> None:
+        super().__init__(f"Front line {control_point_a}/{control_point_b}",
+                         compute_position(control_point_a, control_point_b))
         self.control_point_a = control_point_a
         self.control_point_b = control_point_b
 
@@ -24,22 +42,3 @@ class FrontLine(MissionTarget):
     def control_points(self) -> Tuple[ControlPoint, ControlPoint]:
         """Returns a tuple of the two control points."""
         return self.control_point_a, self.control_point_b
-
-    @property
-    def name(self) -> str:
-        a = self.control_point_a.name
-        b = self.control_point_b.name
-        return f"Front line {a}/{b}"
-
-    @property
-    def position(self) -> Point:
-        a = self.control_point_a.position
-        b = self.control_point_b.position
-        attack_heading = a.heading_between_point(b)
-        attack_distance = a.distance_to_point(b)
-        middle_point = a.point_from_heading(attack_heading, attack_distance / 2)
-
-        strength_delta = (self.control_point_a.base.strength - self.control_point_b.base.strength) / 1.0
-        position = middle_point.point_from_heading(attack_heading,
-                                                   strength_delta * attack_distance / 2 - FRONTLINE_MIN_CP_DISTANCE)
-        return position

theater/missiontarget.py

@@ -1,22 +1,18 @@
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
 from dcs.mapping import Point
 
 
-class MissionTarget(ABC):
-    # TODO: These should just be required objects to the constructor
-    # The TheatherGroundObject class is difficult to modify because it's
-    # generated data that's pickled ahead of time.
-    @property
-    @abstractmethod
-    def name(self) -> str:
-        """The name of the mission target."""
+class MissionTarget:
+    def __init__(self, name: str, position: Point) -> None:
+        """Initializes a mission target.
 
-    @property
-    @abstractmethod
-    def position(self) -> Point:
-        """The location of the mission target."""
+        Args:
+            name: The name of the mission target.
+            position: The location of the mission target.
+        """
+        self.name = name
+        self.position = position
 
     def distance_to(self, other: MissionTarget) -> int:
         """Computes the distance to the given mission target."""

theater/start_generator.py

@@ -1,16 +1,19 @@
+from __future__ import annotations
+
 import logging
 import math
 import pickle
 import random
-import typing
+from typing import Any, Dict, List, Optional
 
 from dcs.mapping import Point
 from dcs.task import CAP, CAS, PinpointStrike
 from dcs.vehicles import AirDefence
 
-from game import db
-from game.data.building_data import DEFAULT_AVAILABLE_BUILDINGS
+from game import Game, db
+from game.factions.faction import Faction
 from game.settings import Settings
+from game.version import VERSION
 from gen import namegen
 from gen.defenses.armor_group_generator import generate_armor_group
 from gen.fleet.ship_group_generator import (
@@ -30,6 +33,13 @@ from theater import (
     TheaterGroundObject,
 )
 from theater.conflicttheater import IMPORTANCE_HIGH, IMPORTANCE_LOW
+from theater.theatergroundobject import (
+    SamGroundObject, BuildingGroundObject, CarrierGroundObject,
+    LhaGroundObject,
+    MissileSiteGroundObject, ShipGroundObject,
+)
+
+GroundObjectTemplates = Dict[str, Dict[str, Any]]
 
 UNIT_VARIETY = 6
 UNIT_AMOUNT_FACTOR = 16
@@ -43,215 +53,387 @@ COUNT_BY_TASK = {
 }
 
 
-def generate_initial_units(theater: ConflictTheater, enemy_country: str, sams: bool, multiplier: float):
-    for cp in theater.enemy_points():
-        if cp.captured:
-            continue
+class GameGenerator:
+    def __init__(self, player: str, enemy: str, theater: ConflictTheater,
+                 settings: Settings, start_date, starting_budget: int,
+                 multiplier: float, midgame: bool) -> None:
+        self.player = player
+        self.enemy = enemy
+        self.theater = theater
+        self.settings = settings
+        self.start_date = start_date
+        self.starting_budget = starting_budget
+        self.multiplier = multiplier
+        self.midgame = midgame
 
+    def generate(self) -> Game:
+        # Reset name generator
+        namegen.reset()
+        self.prepare_theater()
+        self.populate_red_airbases()
+
+        game = Game(player_name=self.player,
+                    enemy_name=self.enemy,
+                    theater=self.theater,
+                    start_date=self.start_date,
+                    settings=self.settings)
+
+        GroundObjectGenerator(game).generate()
+        game.budget = self.starting_budget
+        game.settings.multiplier = self.multiplier
+        game.settings.sams = True
+        game.settings.version = VERSION
+        return game
+
+    def prepare_theater(self) -> None:
+        to_remove = []
+
+        # Auto-capture half the bases if midgame.
+        if self.midgame:
+            control_points = self.theater.controlpoints
+            for control_point in control_points[:len(control_points) // 2]:
+                control_point.captured = True
+
+        # Remove carrier and lha, invert situation if needed
+        for cp in self.theater.controlpoints:
+            no_carrier = self.settings.do_not_generate_carrier
+            no_lha = self.settings.do_not_generate_lha
+            if cp.cptype is ControlPointType.AIRCRAFT_CARRIER_GROUP and \
+                    no_carrier:
+                to_remove.append(cp)
+            elif cp.cptype is ControlPointType.LHA_GROUP and no_lha:
+                to_remove.append(cp)
+
+            if self.settings.inverted:
+                cp.captured = cp.captured_invert
+
+        # do remove
+        for cp in to_remove:
+            self.theater.controlpoints.remove(cp)
+
+        # TODO: Fix this. This captures all bases for blue.
+        # reapply midgame inverted if needed
+        if self.midgame and self.settings.inverted:
+            for i, cp in enumerate(reversed(self.theater.controlpoints)):
+                if i > len(self.theater.controlpoints):
+                    break
+                else:
+                    cp.captured = True
+
+    def populate_red_airbases(self) -> None:
+        for control_point in self.theater.enemy_points():
+            if control_point.captured:
+                continue
+            self.populate_red_airbase(control_point)
+
+    def populate_red_airbase(self, control_point: ControlPoint) -> None:
         # Force reset cp on generation
-        cp.base.aircraft = {}
-        cp.base.armor = {}
-        cp.base.aa = {}
-        cp.base.commision_points = {}
-        cp.base.strength = 1
+        control_point.base.aircraft = {}
+        control_point.base.armor = {}
+        control_point.base.aa = {}
+        control_point.base.commision_points = {}
+        control_point.base.strength = 1
 
         for task in [PinpointStrike, CAP, CAS, AirDefence]:
-            assert cp.importance <= IMPORTANCE_HIGH, "invalid importance {}".format(cp.importance)
-            assert cp.importance >= IMPORTANCE_LOW, "invalid importance {}".format(cp.importance)
+            if IMPORTANCE_HIGH <= control_point.importance <= IMPORTANCE_LOW:
+                raise ValueError(
+                    f"CP importance must be between {IMPORTANCE_LOW} and "
+                    f"{IMPORTANCE_HIGH}, is {control_point.importance}")
 
-            importance_factor = (cp.importance - IMPORTANCE_LOW) / (IMPORTANCE_HIGH - IMPORTANCE_LOW)
-            variety = int(UNIT_VARIETY)
-            unittypes = db.choose_units(task, importance_factor, variety, enemy_country)
-
-            if not sams and task == AirDefence:
-                unittypes = [x for x in db.find_unittype(AirDefence, enemy_country) if x not in db.SAM_BAN]
-
-            count_log = math.log(cp.importance + 0.01, UNIT_COUNT_IMPORTANCE_LOG)
-            count = max(COUNT_BY_TASK[task] * multiplier * (1+count_log), 1)
-
-            if len(unittypes) > 0:
-                count_per_type = max(int(float(count) / len(unittypes)), 1)
-                for unit_type in unittypes:
-                    logging.info("{} - {} {}".format(cp.name, db.unit_type_name(unit_type), count_per_type))
-                    cp.base.commision_units({unit_type: count_per_type})
-
-
-def generate_groundobjects(theater: ConflictTheater, game):
-    with open("resources/groundobject_templates.p", "rb") as f:
-        tpls = pickle.load(f)
-
-    group_id = 0
-    cp_to_remove = []
-    for cp in theater.controlpoints:
-        group_id = generate_cp_ground_points(cp, theater, game, group_id, tpls)
-
-        # CP
-        if cp.captured:
-            faction_name = game.player_name
-        else:
-            faction_name = game.enemy_name
-        faction = db.FACTIONS[faction_name]
-
-        if cp.cptype == ControlPointType.AIRCRAFT_CARRIER_GROUP:
-            # Create ground object group
-            group_id = game.next_group_id()
-            g = TheaterGroundObject("CARRIER")
-            g.group_id = group_id
-            g.object_id = 0
-            g.cp_id = cp.id
-            g.airbase_group = True
-            g.dcs_identifier = "CARRIER"
-            g.sea_object = True
-            g.obj_name = namegen.random_objective_name()
-            g.heading = 0
-            g.position = Point(cp.position.x, cp.position.y)
-            group = generate_carrier_group(faction_name, game, g)
-            g.groups = []
-            if group is not None:
-                g.groups.append(group)
-            cp.ground_objects.append(g)
-            # Set new name :
-            if len(faction.carrier_names) > 0:
-                cp.name = random.choice(faction.carrier_names)
-            else:
-                cp_to_remove.append(cp)
-        elif cp.cptype == ControlPointType.LHA_GROUP:
-            # Create ground object group
-            group_id = game.next_group_id()
-            g = TheaterGroundObject("LHA")
-            g.group_id = group_id
-            g.object_id = 0
-            g.cp_id = cp.id
-            g.airbase_group = True
-            g.dcs_identifier = "LHA"
-            g.sea_object = True
-            g.obj_name = namegen.random_objective_name()
-            g.heading = 0
-            g.position = Point(cp.position.x, cp.position.y)
-            group = generate_lha_group(faction_name, game, g)
-            g.groups = []
-            if group is not None:
-                g.groups.append(group)
-            cp.ground_objects.append(g)
-            # Set new name :
-            if len(faction.helicopter_carrier_names) > 0:
-                cp.name = random.choice(faction.helicopter_carrier_names)
-            else:
-                cp_to_remove.append(cp)
-        else:
-
-            for i in range(random.randint(3, 6)):
-
-                logging.info("GENERATE BASE DEFENSE")
-                point = find_location(True, cp.position, theater, 800, 3200, [], True)
-                logging.info(point)
-
-                if point is None:
-                    logging.info("Couldn't find point for {} base defense".format(cp))
-                    continue
-
-                group_id = game.next_group_id()
-
-                g = TheaterGroundObject("aa")
-                g.group_id = group_id
-                g.object_id = 0
-                g.cp_id = cp.id
-                g.airbase_group = True
-                g.dcs_identifier = "AA"
-                g.sea_object = False
-                g.obj_name = namegen.random_objective_name()
-                g.heading = 0
-                g.position = Point(point.x, point.y)
-
-                generate_airbase_defense_group(i, g, faction_name, game, cp)
-                cp.ground_objects.append(g)
-
-            logging.info("---------------------------")
-            logging.info("CP Generation : " + cp.name)
-            for ground_object in cp.ground_objects:
-                logging.info(ground_object.groups)
-
-        # Generate navy groups
-        if len(faction.navy_generators) > 0 and cp.allow_sea_units:
-
-            if cp.captured and game.settings.do_not_generate_player_navy:
-                continue
-            if not cp.captured and game.settings.do_not_generate_enemy_navy:
+            importance_factor = ((control_point.importance - IMPORTANCE_LOW) /
+                                 (IMPORTANCE_HIGH - IMPORTANCE_LOW))
+            # noinspection PyTypeChecker
+            unit_types = db.choose_units(task, importance_factor, UNIT_VARIETY,
+                                         self.enemy)
+            if not unit_types:
                 continue
 
-            for i in range(faction.navy_group_count):
+            count_log = math.log(control_point.importance + 0.01,
+                                 UNIT_COUNT_IMPORTANCE_LOG)
+            count = max(
+                COUNT_BY_TASK[task] * self.multiplier * (1 + count_log), 1
+            )
 
-                point = find_location(False, cp.position, theater, 5000, 40000, [], False)
-
-                if point is None:
-                    logging.info("Couldn't find point for {} ships".format(cp))
-                    continue
-
-                group_id = game.next_group_id()
-
-                g = TheaterGroundObject("aa")
-                g.group_id = group_id
-                g.object_id = 0
-                g.cp_id = cp.id
-                g.airbase_group = False
-                g.dcs_identifier = "AA"
-                g.sea_object = True
-                g.obj_name = namegen.random_objective_name()
-                g.heading = 0
-                g.position = Point(point.x, point.y)
-
-                group = generate_ship_group(game, g, faction_name)
-                g.groups = []
-                if group is not None:
-                    g.groups.append(group)
-                    cp.ground_objects.append(g)
-
-        if len(faction.missiles) > 0:
-
-            for i in range(faction.missiles_group_count):
-
-                point = find_location(True, cp.position, theater, 2500, 40000, [], False)
-
-                if point is None:
-                    logging.info("Couldn't find point for {} missiles".format(cp))
-                    continue
-
-                group_id = game.next_group_id()
-
-                g = TheaterGroundObject("aa")
-                g.group_id = group_id
-                g.object_id = 0
-                g.cp_id = cp.id
-                g.airbase_group = False
-                g.dcs_identifier = "AA"
-                g.sea_object = False
-                g.obj_name = namegen.random_objective_name()
-                g.heading = 0
-                g.position = Point(point.x, point.y)
-
-                group = generate_missile_group(game, g, faction_name)
-                g.groups = []
-                if group is not None:
-                    g.groups.append(group)
-                    cp.ground_objects.append(g)
-
-    for cp in cp_to_remove:
-        theater.controlpoints.remove(cp)
+            count_per_type = max(int(float(count) / len(unit_types)), 1)
+            for unit_type in unit_types:
+                control_point.base.commision_units({unit_type: count_per_type})
 
 
+class ControlPointGroundObjectGenerator:
+    def __init__(self, game: Game, control_point: ControlPoint,
+                 templates: GroundObjectTemplates) -> None:
+        self.game = game
+        self.control_point = control_point
+        self.templates = templates
 
-def generate_airbase_defense_group(airbase_defense_group_id, ground_obj:TheaterGroundObject, faction, game, cp):
+    @property
+    def faction_name(self) -> str:
+        if self.control_point.captured:
+            return self.game.player_name
+        else:
+            return self.game.enemy_name
 
-    logging.info("GENERATE AIR DEFENSE GROUP")
-    logging.info(faction)
-    logging.info(airbase_defense_group_id)
+    @property
+    def faction(self) -> Faction:
+        return db.FACTIONS[self.faction_name]
 
+    def generate(self) -> bool:
+        self.control_point.ground_objects = []
+        self.generate_ground_points()
+        if self.faction.navy_generators:
+            # Even airbases can generate navies if they are close enough to the
+            # water. This is not controlled by the control point definition, but
+            # rather by whether or not the generator can find a valid position
+            # for the ship.
+            self.generate_navy()
+
+        if self.faction.missiles:
+            # TODO: Presumably only for airbases?
+            self.generate_missile_sites()
+
+        return True
+
+    def generate_ground_points(self) -> None:
+        """Generate ground objects and AA sites for the control point."""
+
+        if self.control_point.is_global:
+            return
+
+        # TODO: Should probably perform this check later.
+        # Just because we don't have factories for the faction doesn't mean we
+        # shouldn't generate AA.
+        available_categories = self.faction.building_set
+        if not available_categories:
+            return
+
+        # Always generate at least one AA point.
+        self.generate_aa_site()
+
+        # And between 2 and 7 other objectives.
+        amount = random.randrange(2, 7)
+        for i in range(amount):
+            # 1 in 4 additional objectives are AA.
+            if random.randint(0, 3) == 0:
+                self.generate_aa_site()
+            else:
+                category = random.choice(available_categories)
+                self.generate_ground_point(category)
+
+    def generate_ground_point(self, category: str) -> None:
+        obj_name = namegen.random_objective_name()
+        template = random.choice(list(self.templates[category].values()))
+        point = find_location(category != "oil",
+                              self.control_point.position,
+                              self.game.theater, 10000, 40000,
+                              self.control_point.ground_objects)
+
+        if point is None:
+            logging.error(
+                f"Could not find point for {obj_name} at {self.control_point}")
+            return
+
+        object_id = 0
+        group_id = self.game.next_group_id()
+
+        # TODO: Create only one TGO per objective, each with multiple units.
+        for unit in template:
+            object_id += 1
+
+            template_point = Point(unit["offset"].x, unit["offset"].y)
+            g = BuildingGroundObject(
+                obj_name, category, group_id, object_id, point + template_point,
+                unit["heading"], self.control_point, unit["type"])
+
+            self.control_point.ground_objects.append(g)
+
+    def generate_aa_site(self) -> None:
+        obj_name = namegen.random_objective_name()
+        position = find_location(True, self.control_point.position,
+                                 self.game.theater, 10000, 40000,
+                                 self.control_point.ground_objects)
+
+        if position is None:
+            logging.error(
+                f"Could not find point for {obj_name} at {self.control_point}")
+            return
+
+        group_id = self.game.next_group_id()
+
+        g = SamGroundObject(namegen.random_objective_name(), group_id,
+                            position, self.control_point, for_airbase=False)
+        group = generate_anti_air_group(self.game, g, self.faction_name)
+        if group is not None:
+            g.groups = [group]
+        self.control_point.ground_objects.append(g)
+
+    def generate_navy(self) -> None:
+        skip_player_navy = self.game.settings.do_not_generate_player_navy
+        if self.control_point.captured and skip_player_navy:
+            return
+
+        skip_enemy_navy = self.game.settings.do_not_generate_enemy_navy
+        if not self.control_point.captured and skip_enemy_navy:
+            return
+
+        for _ in range(self.faction.navy_group_count):
+            if not self.generate_ship():
+                break
+
+    def generate_ship(self) -> None:
+        point = find_location(False, self.control_point.position,
+                              self.game.theater, 5000, 40000, [], False)
+        if point is None:
+            logging.error(
+                f"Could not find point for {self.control_point}'s navy")
+            return
+
+        group_id = self.game.next_group_id()
+
+        g = ShipGroundObject(namegen.random_objective_name(), group_id, point,
+                             self.control_point)
+
+        group = generate_ship_group(self.game, g, self.faction_name)
+        g.groups = []
+        if group is not None:
+            g.groups.append(group)
+            self.control_point.ground_objects.append(g)
+
+    def generate_missile_sites(self) -> None:
+        for i in range(self.faction.missiles_group_count):
+            self.generate_missile_site()
+
+    def generate_missile_site(self) -> None:
+        point = find_location(True, self.control_point.position,
+                              self.game.theater, 2500, 40000, [], False)
+        if point is None:
+            logging.info(
+                f"Could not find point for {self.control_point} missile site")
+            return
+
+        group_id = self.game.next_group_id()
+
+        g = MissileSiteGroundObject(namegen.random_objective_name(), group_id,
+                                    point, self.control_point)
+        group = generate_missile_group(self.game, g, self.faction_name)
+        g.groups = []
+        if group is not None:
+            g.groups.append(group)
+            self.control_point.ground_objects.append(g)
+        return
+
+
+class CarrierGroundObjectGenerator(ControlPointGroundObjectGenerator):
+    def generate(self) -> bool:
+        if not super().generate():
+            return False
+
+        carrier_names = self.faction.carrier_names
+        if not carrier_names:
+            logging.info(
+                f"Skipping generation of {self.control_point.name} because "
+                f"{self.faction_name} has no carriers")
+            return False
+
+        # Create ground object group
+        group_id = self.game.next_group_id()
+        g = CarrierGroundObject(namegen.random_objective_name(), group_id,
+                                self.control_point)
+        group = generate_carrier_group(self.faction_name, self.game, g)
+        g.groups = []
+        if group is not None:
+            g.groups.append(group)
+        self.control_point.ground_objects.append(g)
+        self.control_point.name = random.choice(carrier_names)
+        return True
+
+
+class LhaGroundObjectGenerator(ControlPointGroundObjectGenerator):
+    def generate(self) -> bool:
+        if not super().generate():
+            return False
+
+        lha_names = self.faction.helicopter_carrier_names
+        if not lha_names:
+            logging.info(
+                f"Skipping generation of {self.control_point.name} because "
+                f"{self.faction_name} has no LHAs")
+            return False
+
+        # Create ground object group
+        group_id = self.game.next_group_id()
+        g = LhaGroundObject(namegen.random_objective_name(), group_id,
+                            self.control_point)
+        group = generate_lha_group(self.faction_name, self.game, g)
+        g.groups = []
+        if group is not None:
+            g.groups.append(group)
+        self.control_point.ground_objects.append(g)
+        self.control_point.name = random.choice(lha_names)
+        return True
+
+
+class AirbaseGroundObjectGenerator(ControlPointGroundObjectGenerator):
+    def generate(self) -> bool:
+        if not super().generate():
+            return False
+
+        for i in range(random.randint(3, 6)):
+            self.generate_sam(i)
+        return True
+
+    def generate_sam(self, index: int) -> None:
+        position = find_location(True, self.control_point.position,
+                                 self.game.theater, 800, 3200, [], True)
+        if position is None:
+            logging.error("Could not find position for "
+                          f"{self.control_point} base defense")
+            return
+
+        group_id = self.game.next_group_id()
+
+        g = SamGroundObject(namegen.random_objective_name(), group_id,
+                            position, self.control_point, for_airbase=True)
+
+        generate_airbase_defense_group(index, g, self.faction_name, self.game)
+        self.control_point.ground_objects.append(g)
+
+
+class GroundObjectGenerator:
+    def __init__(self, game: Game) -> None:
+        self.game = game
+        with open("resources/groundobject_templates.p", "rb") as f:
+            self.templates: GroundObjectTemplates = pickle.load(f)
+
+    def generate(self) -> None:
+        # Copied so we can remove items from the original list without breaking
+        # the iterator.
+        control_points = list(self.game.theater.controlpoints)
+        for control_point in control_points:
+            if not self.generate_for_control_point(control_point):
+                self.game.theater.controlpoints.remove(control_point)
+
+    def generate_for_control_point(self, control_point: ControlPoint) -> bool:
+        if control_point.cptype == ControlPointType.AIRCRAFT_CARRIER_GROUP:
+            generator = CarrierGroundObjectGenerator(self.game, control_point,
+                                                     self.templates)
+        elif control_point.cptype == ControlPointType.LHA_GROUP:
+            generator = LhaGroundObjectGenerator(self.game, control_point,
+                                                 self.templates)
+        else:
+            generator = AirbaseGroundObjectGenerator(self.game, control_point,
+                                                     self.templates)
+        return generator.generate()
+
+
+def generate_airbase_defense_group(airbase_defense_group_id: int,
+                                   ground_obj: TheaterGroundObject,
+                                   faction: str, game: Game) -> None:
     if airbase_defense_group_id == 0:
         group = generate_armor_group(faction, game, ground_obj)
     elif airbase_defense_group_id == 1 and random.randint(0, 1) == 0:
-        group = generate_anti_air_group(game, cp, ground_obj, faction)
+        group = generate_anti_air_group(game, ground_obj, faction)
     elif random.randint(0, 2) == 1:
-        group = generate_shorad_group(game, cp, ground_obj, faction)
+        group = generate_shorad_group(game, ground_obj, faction)
     else:
         group = generate_armor_group(faction, game, ground_obj)
 
@@ -260,22 +442,31 @@ def generate_airbase_defense_group(airbase_defense_group_id, ground_obj:TheaterG
         ground_obj.groups.append(group)
 
 
-def find_location(on_ground, near, theater, min, max, others, is_base_defense=False) -> typing.Optional[Point]:
+# TODO: https://stackoverflow.com/a/19482012/632035
+# A lot of the time spent on mission generation is spent in this function since
+# just randomly guess up to 1800 times and often fail. This is particularly
+# problematic while trying to find placement for navies in Nevada.
+def find_location(on_ground: bool, near: Point, theater: ConflictTheater,
+                  min_range: int, max_range: int,
+                  others: List[TheaterGroundObject],
+                  is_base_defense: bool = False) -> Optional[Point]:
     """
     Find a valid ground object location
-    :param on_ground: Whether it should be on ground or on sea (True = on ground)
+    :param on_ground: Whether it should be on ground or on sea (True = on
+    ground)
     :param near: Point
     :param theater: Theater object
-    :param min: Minimal range from point
-    :param max: Max range from point
+    :param min_range: Minimal range from point
+    :param max_range: Max range from point
     :param others: Other already existing ground objects
+    :param is_base_defense: True if the location is for base defense.
     :return:
     """
     point = None
     for _ in range(300):
 
         # Check if on land or sea
-        p = near.random_point_within(max, min)
+        p = near.random_point_within(max_range, min_range)
         if on_ground and theater.is_on_land(p):
             point = p
         elif not on_ground and theater.is_in_sea(p):
@@ -298,7 +489,8 @@ def find_location(on_ground, near, theater, min, max, others, is_base_defense=Fa
 
         if point:
             for other in theater.controlpoints:
-                if is_base_defense: break
+                if is_base_defense:
+                    break
                 if other.position != near:
                     if point is None:
                         break
@@ -313,114 +505,3 @@ def find_location(on_ground, near, theater, min, max, others, is_base_defense=Fa
         if point:
             return point
     return None
-
-
-def generate_cp_ground_points(cp: ControlPoint, theater, game, group_id, templates):
-    """
-    Generate inital ground objects and AA site for given control point
-    :param cp: Control point to initialize
-    :param theater: Theater
-    :param game: Game object
-    :param group_id: Group id
-    :param templates: Ground object templates
-    :return: True if something was generated
-    """
-    # Reset cp ground objects
-    cp.ground_objects = []
-
-    if cp.is_global:
-        return False
-
-    if cp.captured:
-        faction = game.player_name
-    else:
-        faction = game.enemy_name
-    faction_data = db.FACTIONS[faction]
-
-    available_categories = faction_data.building_set
-
-    if len(available_categories) == 0:
-        return False
-
-    amount = random.randrange(3, 8)
-    for i in range(0, amount):
-
-        obj_name = namegen.random_objective_name()
-
-        if i >= amount - 1:
-            tpl_category = "aa"
-        else:
-            if random.randint(0, 3) == 0:
-                tpl_category = "aa"
-            else:
-                tpl_category = random.choice(available_categories)
-
-        tpl = random.choice(list(templates[tpl_category].values()))
-        point = find_location(tpl_category != "oil", cp.position, theater, 10000, 40000, cp.ground_objects)
-
-        if point is None:
-            logging.info("Couldn't find point for {}".format(cp))
-            continue
-
-        object_id = 0
-        group_id = game.next_group_id()
-
-        logging.info("generated {} for {}".format(tpl_category, cp))
-
-        for object in tpl:
-            object_id += 1
-
-            g = TheaterGroundObject(tpl_category)
-            g.group_id = group_id
-            g.object_id = object_id
-            g.cp_id = cp.id
-            g.airbase_group = False
-            g.obj_name = obj_name
-
-            g.dcs_identifier = object["type"]
-            g.heading = object["heading"]
-            g.sea_object = False
-            g.position = Point(point.x + object["offset"].x, point.y + object["offset"].y)
-
-            if g.dcs_identifier == "AA":
-                g.groups = []
-                group = generate_anti_air_group(game, cp, g, faction)
-                if group is not None:
-                    g.groups.append(group)
-
-            cp.ground_objects.append(g)
-    return group_id
-
-
-def prepare_theater(theater: ConflictTheater, settings:Settings, midgame):
-
-    to_remove = []
-
-    # autocapture half the base if midgame
-    if midgame:
-        for i in range(0, int(len(theater.controlpoints) / 2)):
-            theater.controlpoints[i].captured = True
-
-    # Remove carrier and lha, invert situation if needed
-    for cp in theater.controlpoints:
-        if cp.cptype is ControlPointType.AIRCRAFT_CARRIER_GROUP and settings.do_not_generate_carrier:
-            to_remove.append(cp)
-        elif cp.cptype is ControlPointType.LHA_GROUP and settings.do_not_generate_lha:
-            to_remove.append(cp)
-
-        if settings.inverted:
-            cp.captured = cp.captured_invert
-
-    # do remove
-    for cp in to_remove:
-        theater.controlpoints.remove(cp)
-
-    # reapply midgame inverted if needed
-    if midgame and settings.inverted:
-        for i, cp in enumerate(reversed(theater.controlpoints)):
-            if i > len(theater.controlpoints):
-                break
-            else:
-                cp.captured = True
-
-

theater/theatergroundobject.py

@@ -1,4 +1,5 @@
-import uuid
+from __future__ import annotations
+
 from typing import List, TYPE_CHECKING
 
 from dcs.mapping import Point
@@ -68,21 +69,21 @@ CATEGORY_MAP = {
 
 
 class TheaterGroundObject(MissionTarget):
-    cp_id = 0
-    group_id = 0
-    object_id = 0
-    dcs_identifier = None  # type: str
-    is_dead = False
-    airbase_group = False
-    heading = 0
-    position = None  # type: Point
-    groups: List[Group] = []
-    obj_name = ""
-    sea_object = False
-    uuid = uuid.uuid1()
 
-    def __init__(self, category: str):
+    def __init__(self, name: str, category: str, group_id: int, object_id: int,
+                 position: Point, heading: int, cp_id: int, dcs_identifier: str,
+                 airbase_group: bool, sea_object: bool) -> None:
+        super().__init__(name, position)
         self.category = category
+        self.group_id = group_id
+        self.object_id = object_id
+        self.heading = heading
+        self.cp_id = cp_id
+        self.dcs_identifier = dcs_identifier
+        self.airbase_group = airbase_group
+        self.sea_object = sea_object
+        self.is_dead = False
+        self.groups: List[Group] = []
 
     @property
     def string_identifier(self):
@@ -96,20 +97,124 @@ class TheaterGroundObject(MissionTarget):
     def name_abbrev(self) -> str:
         return ABBREV_NAME[self.category]
 
-    def __str__(self):
+    def __str__(self) -> str:
         return NAME_BY_CATEGORY[self.category]
 
-    def matches_string_identifier(self, id):
-        return self.string_identifier == id
+    def matches_string_identifier(self, identifier):
+        return self.string_identifier == identifier
 
     @property
-    def name(self) -> str:
-        return self.obj_name
+    def obj_name(self) -> str:
+        return self.name
 
-    def parent_control_point(
-            self, theater: "ConflictTheater") -> "ControlPoint":
+    def parent_control_point(self, theater: ConflictTheater) -> ControlPoint:
         """Searches the theater for the parent control point."""
         for cp in theater.controlpoints:
             if cp.id == self.cp_id:
                 return cp
         raise RuntimeError("Could not find matching control point in theater")
+
+
+class BuildingGroundObject(TheaterGroundObject):
+    def __init__(self, name: str, category: str, group_id: int, object_id: int,
+                 position: Point, heading: int, control_point: ControlPoint,
+                 dcs_identifier: str) -> None:
+        super().__init__(
+            name=name,
+            category=category,
+            group_id=group_id,
+            object_id=object_id,
+            position=position,
+            heading=heading,
+            cp_id=control_point.id,
+            dcs_identifier=dcs_identifier,
+            airbase_group=False,
+            sea_object=False
+        )
+
+
+# TODO: Why is this both a CP and a TGO?
+class CarrierGroundObject(TheaterGroundObject):
+    def __init__(self, name: str, group_id: int,
+                 control_point: ControlPoint) -> None:
+        super().__init__(
+            name=name,
+            category="CARRIER",
+            group_id=group_id,
+            object_id=0,
+            position=control_point.position,
+            heading=0,
+            cp_id=control_point.id,
+            dcs_identifier="CARRIER",
+            airbase_group=True,
+            sea_object=True
+        )
+
+
+# TODO: Why is this both a CP and a TGO?
+class LhaGroundObject(TheaterGroundObject):
+    def __init__(self, name: str, group_id: int,
+                 control_point: ControlPoint) -> None:
+        super().__init__(
+            name=name,
+            category="LHA",
+            group_id=group_id,
+            object_id=0,
+            position=control_point.position,
+            heading=0,
+            cp_id=control_point.id,
+            dcs_identifier="LHA",
+            airbase_group=True,
+            sea_object=True
+        )
+
+
+class MissileSiteGroundObject(TheaterGroundObject):
+    def __init__(self, name: str, group_id: int, position: Point,
+                 control_point: ControlPoint) -> None:
+        super().__init__(
+            name=name,
+            category="aa",
+            group_id=group_id,
+            object_id=0,
+            position=position,
+            heading=0,
+            cp_id=control_point.id,
+            dcs_identifier="AA",
+            airbase_group=False,
+            sea_object=False
+        )
+
+
+class SamGroundObject(TheaterGroundObject):
+    def __init__(self, name: str, group_id: int, position: Point,
+                 control_point: ControlPoint, for_airbase: bool) -> None:
+        super().__init__(
+            name=name,
+            category="aa",
+            group_id=group_id,
+            object_id=0,
+            position=position,
+            heading=0,
+            cp_id=control_point.id,
+            dcs_identifier="AA",
+            airbase_group=for_airbase,
+            sea_object=False
+        )
+
+
+class ShipGroundObject(TheaterGroundObject):
+    def __init__(self, name: str, group_id: int, position: Point,
+                 control_point: ControlPoint) -> None:
+        super().__init__(
+            name=name,
+            category="aa",
+            group_id=group_id,
+            object_id=0,
+            position=position,
+            heading=0,
+            cp_id=control_point.id,
+            dcs_identifier="AA",
+            airbase_group=False,
+            sea_object=True
+        )