Lingo: Add trap weights option (#2837)

1 files changed, 18 insertions, 5 deletions

@@ -6,7 +6,7 @@ from logging import warning
 from BaseClasses import Item, ItemClassification, Tutorial
 from worlds.AutoWorld import WebWorld, World
 from .datatypes import Room, RoomEntrance
-from .items import ALL_ITEM_TABLE, ITEMS_BY_GROUP, LingoItem
+from .items import ALL_ITEM_TABLE, ITEMS_BY_GROUP, TRAP_ITEMS, LingoItem
 from .locations import ALL_LOCATION_TABLE, LOCATIONS_BY_GROUP
 from .options import LingoOptions
 from .player_logic import LingoPlayerLogic
@@ -91,10 +91,23 @@ class LingoWorld(World):
                     pool.append(self.create_item("Puzzle Skip"))
 
             if traps:
-                traps_list = ["Slowness Trap", "Iceland Trap", "Atbash Trap"]
-
-                for i in range(0, traps):
-                    pool.append(self.create_item(traps_list[i % len(traps_list)]))
+                total_weight = sum(self.options.trap_weights.values())
+
+                if total_weight == 0:
+                    raise Exception("Sum of trap weights must be at least one.")
+
+                trap_counts = {name: int(weight * traps / total_weight)
+                               for name, weight in self.options.trap_weights.items()}
+                
+                trap_difference = traps - sum(trap_counts.values())
+                if trap_difference > 0:
+                    allowed_traps = [name for name in TRAP_ITEMS if self.options.trap_weights[name] > 0]
+                    for i in range(0, trap_difference):
+                        trap_counts[allowed_traps[i % len(allowed_traps)]] += 1
+
+                for name, count in trap_counts.items():
+                    for i in range(0, count):
+                        pool.append(self.create_item(name))
 
         self.multiworld.itempool += pool