import random import uuid from units import gas, rat, bomb, mine class UnitManager: def has_weapon_at(self, position): """Check if there's a weapon (bomb, gas, mine) at the given position""" for unit in self.units.values(): if unit.position == position: # Check if it's a weapon type (not a rat or points) if isinstance(unit, (bomb.Timer, bomb.NuclearBomb, gas.Gas, mine.Mine)): return True return False def count_rats(self): count = 0 for unit in self.units.values(): if isinstance(unit, rat.Rat): count += 1 return count def spawn_gas(self, parent_id=None): if self.map.is_wall(self.pointer[0], self.pointer[1]): return if self.ammo["gas"]["count"] <= 0: return self.ammo["gas"]["count"] -= 1 self.render_engine.play_sound("GAS.WAV") self.spawn_unit(gas.Gas, self.pointer, parent_id=parent_id) def spawn_rat(self, position=None): if position is None: position = self.choose_start() # Don't spawn rats on top of weapons if self.has_weapon_at(position): # Try nearby positions for dx, dy in [(0,1), (1,0), (0,-1), (-1,0), (1,1), (-1,-1), (1,-1), (-1,1)]: alt_pos = (position[0] + dx, position[1] + dy) if not self.map.is_wall(alt_pos[0], alt_pos[1]) and not self.has_weapon_at(alt_pos): position = alt_pos break else: # All nearby positions blocked, abort spawn return rat_class = rat.Male if random.random() < 0.5 else rat.Female self.spawn_unit(rat_class, position) def spawn_bomb(self, position): if self.ammo["bomb"]["count"] <= 0: return self.render_engine.play_sound("PUTDOWN.WAV") self.spawn_unit(bomb.Timer, position) self.ammo["bomb"]["count"] -= 1 def spawn_nuclear_bomb(self, position): """Spawn a nuclear bomb at the specified position""" if self.ammo["nuclear"]["count"] <= 0: return if self.map.is_wall(position[0], position[1]): return self.render_engine.play_sound("NUCLEAR.WAV") self.ammo["nuclear"]["count"] -= 1 self.spawn_unit(bomb.NuclearBomb, position) def spawn_mine(self, position): if self.ammo["mine"]["count"] <= 0: return if self.map.is_wall(position[0], position[1]): return self.render_engine.play_sound("PUTDOWN.WAV") self.ammo["mine"]["count"] -= 1 self.spawn_unit(mine.Mine, position, on_bottom=True) def spawn_unit(self, unit, position, on_bottom=False, **kwargs): id = uuid.uuid4() if on_bottom: self.units = {id: unit(self, position, id, **kwargs), **self.units} else: self.units[id] = unit(self, position, id, **kwargs) def choose_start(self): if not hasattr(self, '_valid_positions'): self._valid_positions = [ (x, y) for y in range(1, self.map.height-1) for x in range(1, self.map.width-1) if self.map.matrix[y][x] ] return random.choice(self._valid_positions) def get_unit_by_id(self, id): return self.units.get(id) or None