import operator import numpy as np class SelectedGrids: def __init__(self, grids): self.grids = grids def __iter__(self): return iter(self.grids) def __getitem__(self, item): if isinstance(item, int): return self.grids[item] else: return SelectedGrids(self.grids[item]) def __contains__(self, item): return item in self.grids def __str__(self): # return str([str(grid) for grid in self]) return '[' + ', ' .join([str(grid) for grid in self]) + ']' def __len__(self): return len(self.grids) def __bool__(self): return self.count > 0 # def __getattr__(self, item): # return [grid.__getattribute__(item) for grid in self.grids] @property def location(self): """ Returns: list[tuple]: """ return [grid.location for grid in self.grids] @property def cost(self): """ Returns: list[int]: """ return [grid.cost for grid in self.grids] @property def weight(self): """ Returns: list[int]: """ return [grid.weight for grid in self.grids] @property def count(self): """ Returns: int: """ return len(self.grids) def select(self, **kwargs): """ Args: **kwargs: Attributes of Grid. Returns: SelectedGrids: """ result = [] for grid in self: flag = True for k, v in kwargs.items(): if grid.__getattribute__(k) != v: flag = False if flag: result.append(grid) return SelectedGrids(result) def add(self, grids): """ Args: grids(SelectedGrids): Returns: SelectedGrids: """ return SelectedGrids(self.grids + grids.grids) def delete(self, grids): """ Args: grids(SelectedGrids): Returns: SelectedGrids: """ g = [grid for grid in self.grids if grid not in grids] return SelectedGrids(g) def sort(self, cost=True, weight=True): """ Args: cost (bool): weight (bool): Returns: """ attr = [] if weight: attr.append('weight') if cost: attr.append('cost') if len(attr): grids = sorted(self.grids, key=operator.attrgetter(*attr)) return SelectedGrids(grids) else: return self def sort_by_camera_distance(self, camera): """ Args: camera (tuple): Returns: SelectedGrids: """ location = np.array(self.location) diff = np.sum(np.abs(np.array(location) - camera), axis=1) # grids = [x for _, x in sorted(zip(diff, self.grids))] grids = tuple(np.array(self.grids)[np.argsort(diff)]) return SelectedGrids(grids) class RoadGrids: def __init__(self, grids): """ Args: grids (list): """ self.grids = [] for grid in grids: if isinstance(grid, list): self.grids.append(SelectedGrids(grids=grid)) else: self.grids.append(SelectedGrids(grids=[grid])) def __str__(self): return str(' - '.join([str(grid) for grid in self.grids])) def roadblocks(self): """ Returns: SelectedGrids: """ grids = [] for block in self.grids: if block.count == block.select(is_enemy=True).count: grids += block.grids return SelectedGrids(grids) def potential_roadblocks(self): """ Returns: SelectedGrids: """ grids = [] for block in self.grids: if np.any([grid.is_fleet for grid in block]): continue if np.any([grid.is_cleared for grid in block]): continue if block.count - block.select(is_enemy=True).count == 1: grids += block.select(is_enemy=True).grids return SelectedGrids(grids) def first_roadblock(self): """ Returns: SelectedGrids: """ grids = [] for block in self.grids: if np.any([grid.is_fleet for grid in block]): continue if np.any([grid.is_cleared for grid in block]): continue if block.select(is_enemy=True).count == 1: grids += block.select(is_enemy=True).grids return SelectedGrids(grids)