diff --git a/game/db.py b/game/db.py
index a4cc3e4a..15024422 100644
--- a/game/db.py
+++ b/game/db.py
@@ -337,15 +337,15 @@ PRICES = {
     AirDefence.SAM_SA_11_Buk_LN_9A310M1: 30,
     AirDefence.SAM_SA_8_Osa_9A33: 28,
     AirDefence.SAM_SA_15_Tor_9A331: 40,
-    AirDefence.SAM_SA_13_Strela_10M3_9A35M3: 24,
-    AirDefence.SAM_SA_9_Strela_1_9P31: 16,
+    AirDefence.SAM_SA_13_Strela_10M3_9A35M3: 16,
+    AirDefence.SAM_SA_9_Strela_1_9P31: 12,
     AirDefence.SAM_SA_11_Buk_CC_9S470M1: 25,
     AirDefence.SAM_SA_8_Osa_LD_9T217: 22,
     AirDefence.SAM_Patriot_AMG_AN_MRC_137: 35,
     AirDefence.SAM_Patriot_ECS_AN_MSQ_104: 30,
     AirDefence.SPAAA_Gepard: 24,
     AirDefence.SAM_Hawk_PCP: 14,
-    AirDefence.AAA_Vulcan_M163: 12,
+    AirDefence.AAA_Vulcan_M163: 10,
     AirDefence.SAM_Hawk_LN_M192: 8,
     AirDefence.SAM_Chaparral_M48: 16,
     AirDefence.SAM_Linebacker_M6: 18,
@@ -358,7 +358,7 @@ PRICES = {
     AirDefence.Stinger_MANPADS: 6,
     AirDefence.SAM_Stinger_comm_dsr: 4,
     AirDefence.SAM_Stinger_comm: 4,
-    AirDefence.SPAAA_ZSU_23_4_Shilka: 12,
+    AirDefence.SPAAA_ZSU_23_4_Shilka: 10,
     AirDefence.AAA_ZU_23_Closed: 6,
     AirDefence.AAA_ZU_23_Emplacement: 6,
     AirDefence.AAA_ZU_23_on_Ural_375: 8,
@@ -387,19 +387,19 @@ PRICES = {
     AirDefence.SAM_SA_2_LN_SM_90: 8,
     AirDefence.SAM_SA_2_TR_SNR_75_Fan_Song: 12,
     AirDefence.Rapier_FSA_Launcher: 6,
-    AirDefence.Rapier_FSA_Optical_Tracker: 12,
-    AirDefence.Rapier_FSA_Blindfire_Tracker: 16,
+    AirDefence.Rapier_FSA_Optical_Tracker: 6,
+    AirDefence.Rapier_FSA_Blindfire_Tracker: 8,
     AirDefence.HQ_7_Self_Propelled_LN: 20,
     AirDefence.HQ_7_Self_Propelled_STR: 24,
     AirDefence.AAA_8_8cm_Flak_18: 6,
     AirDefence.AAA_Flak_38: 6,
     AirDefence.AAA_8_8cm_Flak_36: 8,
-    AirDefence.AAA_8_8cm_Flak_37: 10,
+    AirDefence.AAA_8_8cm_Flak_37: 9,
     AirDefence.AAA_Flak_Vierling_38:6,
     AirDefence.AAA_Kdo_G_40: 8,
     AirDefence.Flak_Searchlight_37: 4,
     AirDefence.Maschinensatz_33: 10,
-    AirDefence.AAA_8_8cm_Flak_41: 12,
+    AirDefence.AAA_8_8cm_Flak_41: 10,
     AirDefence.AAA_Bofors_40mm: 8,
 
     # FRENCH PACK MOD
diff --git a/gen/sam/aaa_bofors.py b/gen/sam/aaa_bofors.py
index ad515931..528edd8b 100644
--- a/gen/sam/aaa_bofors.py
+++ b/gen/sam/aaa_bofors.py
@@ -11,10 +11,11 @@ class BoforsGenerator(GroupGenerator):
     """
 
     name = "Bofors AAA"
+    price = 75
 
     def generate(self):
-        grid_x = random.randint(2, 4)
-        grid_y = random.randint(2, 4)
+        grid_x = random.randint(2, 3)
+        grid_y = random.randint(2, 3)
 
         spacing = random.randint(10,40)
 
diff --git a/gen/sam/aaa_flak.py b/gen/sam/aaa_flak.py
index 896b48aa..5a0d9121 100644
--- a/gen/sam/aaa_flak.py
+++ b/gen/sam/aaa_flak.py
@@ -12,12 +12,13 @@ class FlakGenerator(GroupGenerator):
     """
 
     name = "Flak Site"
+    price = 135
 
     def generate(self):
-        grid_x = random.randint(2, 4)
-        grid_y = random.randint(2, 4)
+        grid_x = random.randint(2, 3)
+        grid_y = random.randint(2, 3)
 
-        spacing = random.randint(30,60)
+        spacing = random.randint(30, 60)
 
         index = 0
         mixed = random.choice([True, False])
@@ -34,7 +35,7 @@ class FlakGenerator(GroupGenerator):
                     unit_type = random.choice(GFLAK)
 
         # Search lights
-        search_pos = self.get_circular_position(random.randint(2,5), 90)
+        search_pos = self.get_circular_position(random.randint(2,3), 90)
         for index, pos in enumerate(search_pos):
             self.add_unit(AirDefence.Flak_Searchlight_37, "SearchLight#" + str(index), pos[0], pos[1], self.heading)
 
diff --git a/gen/sam/aaa_zu23_insurgent.py b/gen/sam/aaa_zu23_insurgent.py
index f9c53620..ec659756 100644
--- a/gen/sam/aaa_zu23_insurgent.py
+++ b/gen/sam/aaa_zu23_insurgent.py
@@ -11,10 +11,11 @@ class ZU23InsurgentGenerator(GroupGenerator):
     """
 
     name = "Zu-23 Site"
+    price = 56
 
     def generate(self):
-        grid_x = random.randint(2, 4)
-        grid_y = random.randint(2, 4)
+        grid_x = random.randint(2, 3)
+        grid_y = random.randint(2, 3)
 
         spacing = random.randint(10,40)
 
diff --git a/gen/sam/sam_avenger.py b/gen/sam/sam_avenger.py
index 63844f7b..44d3aed9 100644
--- a/gen/sam/sam_avenger.py
+++ b/gen/sam/sam_avenger.py
@@ -11,6 +11,7 @@ class AvengerGenerator(GroupGenerator):
     """
 
     name = "Avenger Group"
+    price = 62
 
     def generate(self):
         num_launchers = random.randint(2, 3)
diff --git a/gen/sam/sam_chaparral.py b/gen/sam/sam_chaparral.py
index 6178df1c..a8d89181 100644
--- a/gen/sam/sam_chaparral.py
+++ b/gen/sam/sam_chaparral.py
@@ -11,6 +11,7 @@ class ChaparralGenerator(GroupGenerator):
     """
 
     name = "Chaparral Group"
+    price = 66
 
     def generate(self):
         num_launchers = random.randint(2, 4)
diff --git a/gen/sam/sam_gepard.py b/gen/sam/sam_gepard.py
index 3ece2cc9..501ed7b7 100644
--- a/gen/sam/sam_gepard.py
+++ b/gen/sam/sam_gepard.py
@@ -11,6 +11,7 @@ class GepardGenerator(GroupGenerator):
     """
 
     name = "Gepard Group"
+    price = 50
 
     def generate(self):
         self.add_unit(AirDefence.SPAAA_Gepard, "SPAAA", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_group_generator.py b/gen/sam/sam_group_generator.py
index 2a27d752..add0c9fa 100644
--- a/gen/sam/sam_group_generator.py
+++ b/gen/sam/sam_group_generator.py
@@ -101,16 +101,22 @@ SAM_PRICES = {
     AirDefence.HQ_7_Self_Propelled_LN: 35
 }
 
+
 def get_faction_possible_sams_units(faction: str) -> List[UnitType]:
     """
     Return the list
     :param faction: Faction to search units for
-    :return:
     """
     return [u for u in db.FACTIONS[faction]["units"] if u in AirDefence.__dict__.values()]
 
-def get_sam_names():
-    pass
+
+def get_faction_possible_sams_generator(faction: str) -> List[UnitType]:
+    """
+    Return the list of possible SAM generator for the given faction
+    :param faction: Faction to search units for
+    """
+    return [SAM_MAP[u] for u in get_faction_possible_sams_units(faction)]
+
 
 def generate_anti_air_group(game, parent_cp, ground_object, faction:str):
     """
diff --git a/gen/sam/sam_hawk.py b/gen/sam/sam_hawk.py
index f5273ff9..89c11bc0 100644
--- a/gen/sam/sam_hawk.py
+++ b/gen/sam/sam_hawk.py
@@ -11,6 +11,7 @@ class HawkGenerator(GroupGenerator):
     """
 
     name = "Hawk Site"
+    price = 115
 
     def generate(self):
         self.add_unit(AirDefence.SAM_Hawk_PCP, "PCP", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_hq7.py b/gen/sam/sam_hq7.py
index b32d0226..f8a531ea 100644
--- a/gen/sam/sam_hq7.py
+++ b/gen/sam/sam_hq7.py
@@ -11,6 +11,7 @@ class HQ7Generator(GroupGenerator):
     """
 
     name = "HQ-7 Site"
+    price = 120
 
     def generate(self):
         self.add_unit(AirDefence.HQ_7_Self_Propelled_STR, "STR", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_linebacker.py b/gen/sam/sam_linebacker.py
index add33acb..946d14ed 100644
--- a/gen/sam/sam_linebacker.py
+++ b/gen/sam/sam_linebacker.py
@@ -11,6 +11,7 @@ class LinebackerGenerator(GroupGenerator):
     """
 
     name = "Linebacker Group"
+    price = 75
 
     def generate(self):
         num_launchers = random.randint(2, 4)
diff --git a/gen/sam/sam_patriot.py b/gen/sam/sam_patriot.py
index f042e3b1..b55dbaea 100644
--- a/gen/sam/sam_patriot.py
+++ b/gen/sam/sam_patriot.py
@@ -11,6 +11,7 @@ class PatriotGenerator(GroupGenerator):
     """
 
     name = "Patriot Battery"
+    price = 240
 
     def generate(self):
         # Command Post
@@ -20,13 +21,13 @@ class PatriotGenerator(GroupGenerator):
         self.add_unit(AirDefence.SAM_Patriot_EPP_III, "EPP", self.position.x, self.position.y + 30, self.heading)
         self.add_unit(AirDefence.SAM_Patriot_STR_AN_MPQ_53, "ICC", self.position.x + 30, self.position.y + 30, self.heading)
 
-        num_launchers = random.randint(2, 4)
+        num_launchers = random.randint(3, 4)
         positions = self.get_circular_position(num_launchers, launcher_distance=120, coverage=360)
         for i, position in enumerate(positions):
             self.add_unit(AirDefence.SAM_Patriot_LN_M901, "LN#" + str(i), position[0], position[1], position[2])
 
         # Short range protection for high value site
-        num_launchers = random.randint(2, 4)
+        num_launchers = random.randint(3, 4)
         positions = self.get_circular_position(num_launchers, launcher_distance=300, coverage=360)
         for i, position in enumerate(positions):
             self.add_unit(AirDefence.AAA_Vulcan_M163, "SPAAA#" + str(i), position[0], position[1], position[2])
diff --git a/gen/sam/sam_rapier.py b/gen/sam/sam_rapier.py
index 0cbffd9e..99b7b205 100644
--- a/gen/sam/sam_rapier.py
+++ b/gen/sam/sam_rapier.py
@@ -11,6 +11,7 @@ class RapierGenerator(GroupGenerator):
     """
 
     name = "Rapier AA Site"
+    price = 50
 
     def generate(self):
         self.add_unit(AirDefence.Rapier_FSA_Blindfire_Tracker, "BT", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_roland.py b/gen/sam/sam_roland.py
index 22f60435..9e31d5fe 100644
--- a/gen/sam/sam_roland.py
+++ b/gen/sam/sam_roland.py
@@ -9,6 +9,7 @@ class RolandGenerator(GroupGenerator):
     """
 
     name = "Roland Site"
+    price = 40
 
     def generate(self):
         self.add_unit(AirDefence.SAM_Roland_ADS, "ADS", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_sa10.py b/gen/sam/sam_sa10.py
index c90d3ae9..ae2102f0 100644
--- a/gen/sam/sam_sa10.py
+++ b/gen/sam/sam_sa10.py
@@ -11,6 +11,7 @@ class SA10Generator(GroupGenerator):
     """
 
     name = "SA-10/S-300PS Battery"
+    price = 450
 
     def generate(self):
         # Command Post
diff --git a/gen/sam/sam_sa11.py b/gen/sam/sam_sa11.py
index c05fe21a..3af6c242 100644
--- a/gen/sam/sam_sa11.py
+++ b/gen/sam/sam_sa11.py
@@ -11,6 +11,7 @@ class SA11Generator(GroupGenerator):
     """
 
     name = "SA-11 Buk Battery"
+    price = 180
 
     def generate(self):
         self.add_unit(AirDefence.SAM_SA_11_Buk_CC_9S470M1, "CC", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_sa13.py b/gen/sam/sam_sa13.py
index 366746d7..8fc069ad 100644
--- a/gen/sam/sam_sa13.py
+++ b/gen/sam/sam_sa13.py
@@ -11,6 +11,7 @@ class SA13Generator(GroupGenerator):
     """
 
     name = "SA-13 Strela Group"
+    price = 50
 
     def generate(self):
         self.add_unit(Unarmed.Transport_UAZ_469, "UAZ", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_sa15.py b/gen/sam/sam_sa15.py
index 0aa9c162..09fda2ee 100644
--- a/gen/sam/sam_sa15.py
+++ b/gen/sam/sam_sa15.py
@@ -9,6 +9,7 @@ class SA15Generator(GroupGenerator):
     """
 
     name = "SA-15 Tor Group"
+    price = 55
 
     def generate(self):
         self.add_unit(AirDefence.SAM_SA_15_Tor_9A331, "ADS", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_sa19.py b/gen/sam/sam_sa19.py
index 10923a5a..c4f710f4 100644
--- a/gen/sam/sam_sa19.py
+++ b/gen/sam/sam_sa19.py
@@ -11,6 +11,7 @@ class SA19Generator(GroupGenerator):
     """
 
     name = "SA-19 Tunguska Group"
+    price = 90
 
     def generate(self):
         num_launchers = random.randint(1, 3)
diff --git a/gen/sam/sam_sa2.py b/gen/sam/sam_sa2.py
index eb7f014b..c108c1e8 100644
--- a/gen/sam/sam_sa2.py
+++ b/gen/sam/sam_sa2.py
@@ -11,6 +11,7 @@ class SA2Generator(GroupGenerator):
     """
 
     name = "SA-2/S-75 Site"
+    price = 74
 
     def generate(self):
         self.add_unit(AirDefence.SAM_SR_P_19, "SR", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_sa3.py b/gen/sam/sam_sa3.py
index 9a73de13..455bab19 100644
--- a/gen/sam/sam_sa3.py
+++ b/gen/sam/sam_sa3.py
@@ -11,6 +11,7 @@ class SA3Generator(GroupGenerator):
     """
 
     name = "SA-3/S-125 Site"
+    price = 80
 
     def generate(self):
         self.add_unit(AirDefence.SAM_SR_P_19, "SR", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_sa6.py b/gen/sam/sam_sa6.py
index f96d3989..7ec2afca 100644
--- a/gen/sam/sam_sa6.py
+++ b/gen/sam/sam_sa6.py
@@ -11,6 +11,7 @@ class SA6Generator(GroupGenerator):
     """
 
     name = "SA-6 Kub Site"
+    price = 102
 
     def generate(self):
         self.add_unit(AirDefence.SAM_SA_6_Kub_STR_9S91, "STR", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_sa8.py b/gen/sam/sam_sa8.py
index a5df4d85..1c09dd2e 100644
--- a/gen/sam/sam_sa8.py
+++ b/gen/sam/sam_sa8.py
@@ -11,6 +11,7 @@ class SA8Generator(GroupGenerator):
     """
 
     name = "SA-8 OSA Site"
+    price = 55
 
     def generate(self):
         self.add_unit(AirDefence.SAM_SA_8_Osa_9A33, "OSA", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_sa9.py b/gen/sam/sam_sa9.py
index 9eeffe3c..d0045bea 100644
--- a/gen/sam/sam_sa9.py
+++ b/gen/sam/sam_sa9.py
@@ -11,6 +11,7 @@ class SA9Generator(GroupGenerator):
     """
 
     name = "SA-9 Group"
+    price = 40
 
     def generate(self):
         self.add_unit(Unarmed.Transport_UAZ_469, "UAZ", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_vulcan.py b/gen/sam/sam_vulcan.py
index 71789e62..77cfc0a2 100644
--- a/gen/sam/sam_vulcan.py
+++ b/gen/sam/sam_vulcan.py
@@ -11,6 +11,7 @@ class VulcanGenerator(GroupGenerator):
     """
 
     name = "Vulcan Group"
+    price = 25
 
     def generate(self):
         self.add_unit(AirDefence.AAA_Vulcan_M163, "SPAAA", self.position.x, self.position.y, self.heading)
diff --git a/gen/sam/sam_zsu23.py b/gen/sam/sam_zsu23.py
index 33d1cde3..7c90cb4d 100644
--- a/gen/sam/sam_zsu23.py
+++ b/gen/sam/sam_zsu23.py
@@ -11,9 +11,10 @@ class ZSU23Generator(GroupGenerator):
     """
 
     name = "ZSU-23 Group"
+    price = 50
 
     def generate(self):
-        num_launchers = random.randint(2, 5)
+        num_launchers = random.randint(4, 5)
 
         positions = self.get_circular_position(num_launchers, launcher_distance=120, coverage=180)
         for i, position in enumerate(positions):
diff --git a/gen/sam/sam_zu23.py b/gen/sam/sam_zu23.py
index 1f4521cd..3134c3a7 100644
--- a/gen/sam/sam_zu23.py
+++ b/gen/sam/sam_zu23.py
@@ -11,10 +11,11 @@ class ZU23Generator(GroupGenerator):
     """
 
     name = "ZU-23 Group"
+    price = 54
 
     def generate(self):
-        grid_x = random.randint(2, 4)
-        grid_y = random.randint(2, 4)
+        grid_x = random.randint(2, 3)
+        grid_y = random.randint(2, 3)
 
         spacing = random.randint(10,40)
 
diff --git a/gen/sam/sam_zu23_ural.py b/gen/sam/sam_zu23_ural.py
index d6b08fac..1eb31b22 100644
--- a/gen/sam/sam_zu23_ural.py
+++ b/gen/sam/sam_zu23_ural.py
@@ -11,6 +11,7 @@ class ZU23UralGenerator(GroupGenerator):
     """
 
     name = "ZU-23 Ural Group"
+    price = 64
 
     def generate(self):
         num_launchers = random.randint(2, 8)
diff --git a/gen/sam/sam_zu23_ural_insurgent.py b/gen/sam/sam_zu23_ural_insurgent.py
index 55c72196..4512cfc7 100644
--- a/gen/sam/sam_zu23_ural_insurgent.py
+++ b/gen/sam/sam_zu23_ural_insurgent.py
@@ -11,6 +11,7 @@ class ZU23UralInsurgentGenerator(GroupGenerator):
     """
 
     name = "ZU-23 Ural Insurgent Group"
+    price = 64
 
     def generate(self):
         num_launchers = random.randint(2, 8)
diff --git a/qt_ui/widgets/QTopPanel.py b/qt_ui/widgets/QTopPanel.py
index 2a77cb23..30725095 100644
--- a/qt_ui/widgets/QTopPanel.py
+++ b/qt_ui/widgets/QTopPanel.py
@@ -21,6 +21,7 @@ class QTopPanel(QFrame):
         self.setMaximumHeight(70)
         self.init_ui()
         GameUpdateSignal.get_instance().gameupdated.connect(self.setGame)
+        GameUpdateSignal.get_instance().budgetupdated.connect(self.budget_update)
 
     def init_ui(self):
 
@@ -101,4 +102,7 @@ class QTopPanel(QFrame):
 
     def proceed(self):
         self.subwindow = QMissionPlanning(self.game)
-        self.subwindow.show()
\ No newline at end of file
+        self.subwindow.show()
+
+    def budget_update(self, game:Game):
+        self.budgetBox.setGame(game)
diff --git a/qt_ui/windows/GameUpdateSignal.py b/qt_ui/windows/GameUpdateSignal.py
index 86546757..dd32dd58 100644
--- a/qt_ui/windows/GameUpdateSignal.py
+++ b/qt_ui/windows/GameUpdateSignal.py
@@ -16,6 +16,7 @@ class GameUpdateSignal(QObject):
 
     instance = None
     gameupdated = Signal(Game)
+    budgetupdated = Signal(Game)
     debriefingReceived = Signal(DebriefingSignal)
 
     def __init__(self):
@@ -25,6 +26,9 @@ class GameUpdateSignal(QObject):
     def updateGame(self, game: Game):
         self.gameupdated.emit(game)
 
+    def updateBudget(self, game: Game):
+        self.budgetupdated.emit(game)
+
     def sendDebriefing(self, game: Game, gameEvent: Event, debriefing: Debriefing):
         sig = DebriefingSignal(game, gameEvent, debriefing)
         self.gameupdated.emit(game)
diff --git a/qt_ui/windows/basemenu/base_defenses/QBaseDefenseGroupInfo.py b/qt_ui/windows/basemenu/base_defenses/QBaseDefenseGroupInfo.py
index 2b2abb67..836c1772 100644
--- a/qt_ui/windows/basemenu/base_defenses/QBaseDefenseGroupInfo.py
+++ b/qt_ui/windows/basemenu/base_defenses/QBaseDefenseGroupInfo.py
@@ -1,3 +1,4 @@
+from PySide2.QtCore import Qt
 from PySide2.QtWidgets import QGridLayout, QLabel, QGroupBox, QPushButton
 
 from qt_ui.uiconstants import VEHICLES_ICONS
@@ -28,21 +29,21 @@ class QBaseDefenseGroupInfo(QGroupBox):
                     unit_dict[u.type] = 1
         i = 0
         for k, v in unit_dict.items():
-            #icon = QLabel()
-            #if k in VEHICLES_ICONS.keys():
-            #    icon.setPixmap(VEHICLES_ICONS[k])
-            #else:
-            #    icon.setText("<b>" + k[:6] + "</b>")
-            #icon.setProperty("style", "icon-plane")
-            #layout.addWidget(icon, i, 0)
-            layout.addWidget(QLabel(str(v) + " x " + "<strong>" + k + "</strong>"), i, 0)
+            icon = QLabel()
+            if k in VEHICLES_ICONS.keys():
+                icon.setPixmap(VEHICLES_ICONS[k])
+            else:
+                icon.setText("<b>" + k[:6] + "</b>")
+            icon.setProperty("style", "icon-plane")
+            layout.addWidget(icon, i, 0)
+            layout.addWidget(QLabel(str(v) + " x " + "<strong>" + k + "</strong>"), i, 1)
             i = i + 1
 
         manage_button = QPushButton("Manage")
         manage_button.setProperty("style", "btn-success")
         manage_button.setMaximumWidth(180)
         manage_button.clicked.connect(self.onManage)
-        layout.addWidget(manage_button, i+1, 0)
+        layout.addWidget(manage_button, i, 0, Qt.AlignLeft)
         self.setLayout(layout)
 
     def onManage(self):
diff --git a/qt_ui/windows/groundobject/QGroundObjectMenu.py b/qt_ui/windows/groundobject/QGroundObjectMenu.py
index 6e67d0c7..6b8c9ed1 100644
--- a/qt_ui/windows/groundobject/QGroundObjectMenu.py
+++ b/qt_ui/windows/groundobject/QGroundObjectMenu.py
@@ -1,12 +1,16 @@
 import logging
 
-from PySide2.QtGui import QCloseEvent
-from PySide2.QtWidgets import QHBoxLayout, QWidget, QDialog, QGridLayout, QLabel, QGroupBox, QVBoxLayout, QPushButton
+from PySide2 import QtCore
+from PySide2.QtGui import QCloseEvent, Qt
+from PySide2.QtWidgets import QHBoxLayout, QDialog, QGridLayout, QLabel, QGroupBox, QVBoxLayout, QPushButton, \
+    QComboBox, QSpinBox, QMessageBox
 from dcs import Point
 
-from game import Game
+from game import Game, db
 from game.data.building_data import FORTIFICATION_BUILDINGS
-from game.db import PRICES, unit_type_of
+from game.db import PRICES, unit_type_of, PinpointStrike
+from gen.defenses.armor_group_generator import generate_armor_group
+from gen.sam.sam_group_generator import get_faction_possible_sams_generator
 from qt_ui.uiconstants import EVENT_ICONS
 from qt_ui.widgets.QBudgetBox import QBudgetBox
 from qt_ui.windows.GameUpdateSignal import GameUpdateSignal
@@ -29,6 +33,7 @@ class QGroundObjectMenu(QDialog):
         self.buildingBox = QGroupBox("Buildings :")
         self.intelLayout = QGridLayout()
         self.buildingsLayout = QGridLayout()
+        self.sell_all_button = None
         self.total_value = 0
         self.init_ui()
 
@@ -46,11 +51,20 @@ class QGroundObjectMenu(QDialog):
             self.mainLayout.addWidget(self.buildingBox)
 
         self.actionLayout = QHBoxLayout()
-        sell_all_button = QPushButton("Disband (+" + str(self.total_value) + "M)")
-        sell_all_button.clicked.connect(self.sell_all)
-        self.actionLayout.addWidget(sell_all_button)
+
+        self.sell_all_button = QPushButton("Disband (+" + str(self.total_value) + "M)")
+        self.sell_all_button.clicked.connect(self.sell_all)
+        self.sell_all_button.setProperty("style", "btn-danger")
+
+        self.buy_replace = QPushButton("Buy/Replace")
+        self.buy_replace.clicked.connect(self.buy_group)
+        self.buy_replace.setProperty("style", "btn-success")
 
         if self.total_value > 0:
+            self.actionLayout.addWidget(self.sell_all_button)
+        self.actionLayout.addWidget(self.buy_replace)
+
+        if self.cp.captured:
             self.mainLayout.addLayout(self.actionLayout)
         self.setLayout(self.mainLayout)
 
@@ -82,6 +96,9 @@ class QGroundObjectMenu(QDialog):
                     repair.clicked.connect(lambda u=u, g=g, p=price: self.repair_unit(g, u, p))
                     self.intelLayout.addWidget(repair, i, 1)
                 i = i + 1
+        stretch = QVBoxLayout()
+        stretch.addStretch()
+        self.intelLayout.addLayout(stretch, i, 0)
 
         self.buildingBox = QGroupBox("Buildings :")
         self.buildingsLayout = QGridLayout()
@@ -97,12 +114,27 @@ class QGroundObjectMenu(QDialog):
     def do_refresh_layout(self):
         try:
             for i in range(self.mainLayout.count()):
-                self.mainLayout.removeItem(self.mainLayout.itemAt(i))
+                item = self.mainLayout.itemAt(i)
+                if item is not None and item.widget() is not None:
+                    item.widget().setParent(None)
+            self.sell_all_button.setParent(None)
+            self.buy_replace.setParent(None)
+            self.actionLayout.setParent(None)
+
             self.doLayout()
-            if len(self.ground_object.groups) > 0:
+            if self.ground_object.dcs_identifier == "AA":
                 self.mainLayout.addWidget(self.intelBox)
             else:
                 self.mainLayout.addWidget(self.buildingBox)
+
+            self.actionLayout = QHBoxLayout()
+            if self.total_value > 0:
+                self.actionLayout.addWidget(self.sell_all_button)
+            self.actionLayout.addWidget(self.buy_replace)
+
+            if self.cp.captured:
+                self.mainLayout.addLayout(self.actionLayout)
+
         except Exception as e:
             print(e)
         self.update_total_value()
@@ -116,6 +148,8 @@ class QGroundObjectMenu(QDialog):
                     total_value = total_value + PRICES[utype]
                 else:
                     total_value = total_value + 1
+        if self.sell_all_button is not None:
+            self.sell_all_button.setText("Disband (+$" + str(self.total_value) + "M)")
         self.total_value = total_value
 
     def repair_unit(self, group, unit, price):
@@ -140,10 +174,163 @@ class QGroundObjectMenu(QDialog):
         self.update_total_value()
         self.game.budget = self.game.budget + self.total_value
         self.ground_object.groups = []
-        GameUpdateSignal.get_instance().updateGame(self.game)
+        self.do_refresh_layout()
+        GameUpdateSignal.get_instance().updateBudget(self.game)
 
     def buy_group(self):
-        pass
+        self.subwindow = QBuyGroupForGroundObjectDialog(self, self.ground_object, self.cp, self.game, self.total_value)
+        self.subwindow.changed.connect(self.do_refresh_layout)
+        self.subwindow.show()
+
 
     def closeEvent(self, closeEvent: QCloseEvent):
         pass
+
+
+class QBuyGroupForGroundObjectDialog(QDialog):
+
+    changed = QtCore.Signal()
+
+    def __init__(self, parent, ground_object: TheaterGroundObject, cp: ControlPoint, game: Game, current_group_value: int):
+        super(QBuyGroupForGroundObjectDialog, self).__init__(parent)
+
+        self.setMinimumWidth(350)
+        self.ground_object = ground_object
+        self.cp = cp
+        self.game = game
+        self.current_group_value = current_group_value
+
+        self.setWindowTitle("Buy units @ " + self.ground_object.obj_name)
+        self.setWindowIcon(EVENT_ICONS["capture"])
+
+        self.buySamButton = QPushButton("Buy")
+        self.buyArmorButton = QPushButton("Buy")
+        self.buySamLayout = QGridLayout()
+        self.buyArmorLayout = QGridLayout()
+        self.amount = QSpinBox()
+        self.buyArmorCombo = QComboBox()
+        self.samCombo = QComboBox()
+        self.buySamBox = QGroupBox("Buy SAM site :")
+        self.buyArmorBox = QGroupBox("Buy defensive position :")
+
+
+
+        self.init_ui()
+
+    def init_ui(self):
+        faction = self.game.player_name
+
+        # Sams
+
+        possible_sams = get_faction_possible_sams_generator(faction)
+        for sam in possible_sams:
+            self.samCombo.addItem(sam.name + " [$" + str(sam.price) + "M]", userData=sam)
+        self.samCombo.currentIndexChanged.connect(self.samComboChanged)
+
+        self.buySamLayout.addWidget(QLabel("Site Type :"), 0, 0, Qt.AlignLeft)
+        self.buySamLayout.addWidget(self.samCombo, 0, 1, alignment=Qt.AlignRight)
+        self.buySamLayout.addWidget(self.buySamButton, 1, 1, alignment=Qt.AlignRight)
+        stretch = QVBoxLayout()
+        stretch.addStretch()
+        self.buySamLayout.addLayout(stretch, 2, 0)
+
+        self.buySamButton.clicked.connect(self.buySam)
+
+        # Armored units
+
+        armored_units = db.find_unittype(PinpointStrike, faction) # Todo : refactor this legacy nonsense
+        for unit in set(armored_units):
+            self.buyArmorCombo.addItem(db.unit_type_name_2(unit) + " [$" + str(db.PRICES[unit]) + "M]", userData=unit)
+        self.buyArmorCombo.currentIndexChanged.connect(self.armorComboChanged)
+
+        self.amount.setMinimum(2)
+        self.amount.setMaximum(8)
+        self.amount.setValue(2)
+        self.amount.valueChanged.connect(self.amountComboChanged)
+
+        self.buyArmorLayout.addWidget(QLabel("Unit type :"), 0, 0, Qt.AlignLeft)
+        self.buyArmorLayout.addWidget(self.buyArmorCombo, 0, 1, alignment=Qt.AlignRight)
+        self.buyArmorLayout.addWidget(QLabel("Group size :"), 1, 0, alignment=Qt.AlignLeft)
+        self.buyArmorLayout.addWidget(self.amount, 1, 1, alignment=Qt.AlignRight)
+        self.buyArmorLayout.addWidget(self.buyArmorButton, 2, 1, alignment=Qt.AlignRight)
+        stretch2 = QVBoxLayout()
+        stretch2.addStretch()
+        self.buyArmorLayout.addLayout(stretch2, 3, 0)
+
+        self.buyArmorButton.clicked.connect(self.buyArmor)
+
+        # Do layout
+        self.buySamBox.setLayout(self.buySamLayout)
+        self.buyArmorBox.setLayout(self.buyArmorLayout)
+
+        self.mainLayout = QHBoxLayout()
+        self.mainLayout.addWidget(self.buySamBox)
+
+        if self.ground_object.airbase_group:
+            self.mainLayout.addWidget(self.buyArmorBox)
+
+        self.setLayout(self.mainLayout)
+
+        try:
+            self.samComboChanged(0)
+            self.armorComboChanged(0)
+        except:
+            pass
+
+    def samComboChanged(self, index):
+        self.buySamButton.setText("Buy [$" + str(self.samCombo.itemData(index).price) + "M] [-$" + str(self.current_group_value) + "M]")
+
+    def armorComboChanged(self, index):
+        self.buyArmorButton.setText("Buy [$" + str(db.PRICES[self.buyArmorCombo.itemData(index)] * self.amount.value()) + "M][-$" + str(self.current_group_value) + "M]")
+
+    def amountComboChanged(self):
+        self.buyArmorButton.setText("Buy [$" + str(db.PRICES[self.buyArmorCombo.itemData(self.buyArmorCombo.currentIndex())] * self.amount.value()) + "M][-$" + str(self.current_group_value) + "M]")
+
+    def buyArmor(self):
+        utype = self.buyArmorCombo.itemData(self.buyArmorCombo.currentIndex())
+        price = db.PRICES[utype] * self.amount.value() - self.current_group_value
+        if price > self.game.budget:
+            self.error_money()
+            self.close()
+            return
+        else:
+            self.game.budget -= price
+
+        # Generate Armor
+        group = generate_armor_group(self.game.player_name, self.game, self.ground_object)
+        self.ground_object.groups = [group]
+
+        GameUpdateSignal.get_instance().updateBudget(self.game)
+
+        self.changed.emit()
+        self.close()
+
+    def buySam(self):
+        sam_generator = self.samCombo.itemData(self.samCombo.currentIndex())
+        price = sam_generator.price - self.current_group_value
+        if price > self.game.budget:
+            self.error_money()
+            return
+        else:
+            self.game.budget -= price
+
+        # Generate SAM
+        generator = sam_generator(self.game, self.ground_object)
+        generator.generate()
+        generated_group = generator.get_generated_group()
+        self.ground_object.groups = [generated_group]
+
+        GameUpdateSignal.get_instance().updateBudget(self.game)
+
+        self.changed.emit()
+        self.close()
+
+    def error_money(self):
+        msg = QMessageBox()
+        msg.setIcon(QMessageBox.Information)
+        msg.setText("Not enough money to buy these units !")
+        msg.setWindowTitle("Not enough money")
+        msg.setStandardButtons(QMessageBox.Ok)
+        msg.setWindowFlags(Qt.WindowStaysOnTopHint)
+        msg.exec_()
+        self.close()
\ No newline at end of file