1  from enum import Enum
  2  import collections
  3  import random
  4  import pygame
  5  from pygame import Rect
  6  from pygame.math import Vector2
  7  
  8  class VaultColour(Enum):
  9      RED = 1
 10      GREEN = 2
 11      BLUE = 3
 12      GOLD = 4
 13  
 14  class Schematic:
 15      def __init__(self, name: str, size: Vector2, connectors: [Vector2]):
 16          self.name = name
 17          self.size = size
 18          self.connectors = connectors
 19  
 20      def __repr__(self):
 21          return self.name + " (" + str(self.size) + ")"
 22  
 23      def __hash__(self):
 24          return hash((tuple(self.size), tuple(map(tuple, self.connectors))))
 25  
 26      def can_place(self, topleft: Vector2, connecting_to: Vector2) -> bool:
 27          boundingbox = Rect(topleft, self.size)
 28          for room in placed_rooms:
 29              if boundingbox.colliderect(room.boundingbox):
 30                  return False
 31  
 32          our_connectors = [c + topleft for c in self.connectors]
 33          for location in free_connectors:
 34              for offset in [(x,y) for x in range(-1, 2) for y in range(-1, 2)]:
 35                  if boundingbox.collidepoint(location + offset):
 36                      # Overlaps space adjacent to a connector
 37                      # And is not connecting to it
 38                      if location + offset not in our_connectors:
 39                          return False
 40  
 41          for connector in our_connectors:
 42              if not any(connector.distance_squared_to(other_conn) <= 1 for other_conn in free_connectors + [connecting_to]):
 43                  # Not touching another connector
 44                  for room in placed_rooms:
 45                      for offset in [(x,y) for x in range(-1, 2) for y in range(-1, 2)]:
 46                          if room.boundingbox.collidepoint(connector + offset):
 47                              # And space adjacent to our connector is blocked...
 48                              # which means that we've blocked ourselves off
 49                              return False
 50  
 51          return True
 52  
 53      def place(self, topleft: Vector2, parentrooms: list):
 54          boundingbox = Rect(topleft, self.size)
 55  
 56          connectors_to_connect = [c + topleft for c in self.connectors]
 57          other_connectors_connected = []
 58          for location in free_connectors:
 59              for offset in [(x,y) for x in range(-1, 2) for y in range(-1, 2)]:
 60                  if boundingbox.collidepoint(location + offset):
 61                      # Overlaps space adjacent to a connector
 62                      if location + offset in connectors_to_connect:
 63                          connectors_to_connect.remove(location + offset)  # Connect our end
 64                          other_connectors_connected.append(location)  # Connect their end
 65                          # Connect it!
 66                          
 67          for connected in other_connectors_connected:free_connectors.remove(connected)  # Connector is no longer free
 68  
 69          placed = Room(self, boundingbox, [c + topleft for c in self.connectors], parentrooms)
 70          placed_rooms.append(placed)
 71  
 72          for connector in connectors_to_connect:
 73              free_connectors.append(connector)
 74              upcoming_calculations.append((connector, parentrooms + [placed]))
 75              
 76          return placed
 77  
 78      def rotatedClockwise90(self):
 79          size = Vector2(self.size.y, self.size.x)
 80  
 81          connectors = []
 82          for conn in self.connectors:
 83              rotated = Vector2(conn.y, self.size.x-1-conn.x)
 84              connectors.append(rotated)
 85  
 86          return Schematic(self.name + "-rotated", size, connectors)
 87  
 88  class Room:
 89      def __init__(self, schematic: Schematic, boundingbox: Rect, connectors: [Vector2], parentrooms):
 90          self.schematic = schematic
 91          self.boundingbox = boundingbox
 92          self.connectors = connectors
 93          self.parentrooms = parentrooms
 94          self.colour = None
 95  
 96      def __repr__(self):
 97          return str(self.schematic) + " placed at " + str(self.boundingbox)
 98  
 99  print("Loading schematics...")
100  # Schematics
101  schematics = {
102      Schematic("Boring Hallway 1", Vector2(9, 3), [Vector2(0, 1), Vector2(8, 1)]),
103      Schematic("Boring Turning 1", Vector2(3, 3), [Vector2(0, 1), Vector2(1, 0)]),
104      Schematic("Boring T-Junction 1", Vector2(4, 3), [Vector2(0, 1), Vector2(2, 0), Vector2(2, 2)]),
105      }
106  vault_schematics = {
107      VaultColour.RED: {Schematic("Vault", Vector2(18, 9), [Vector2(0, 4)])},
108      VaultColour.GREEN: {Schematic("Vault", Vector2(18, 9), [Vector2(0, 4)])},
109      VaultColour.BLUE: {Schematic("Vault", Vector2(18, 9), [Vector2(0, 4)])},
110      VaultColour.GOLD: {Schematic("Vault", Vector2(20, 9), [Vector2(0, 4)])}
111      }
112  key_schematics = {
113      VaultColour.RED: {Schematic("Key Room", Vector2(9, 12), [Vector2(4, 0)])},
114      VaultColour.GREEN: {Schematic("Key Room", Vector2(5, 5), [Vector2(0, 2)])},
115      VaultColour.BLUE: set(),
116      VaultColour.GOLD: {Schematic("Key Room", Vector2(20, 8), [Vector2(0, 4)])}
117      }
118  # Permutations
119  for sch in schematics.copy():
120      # Rotation
121      for i in range(3):
122          sch = sch.rotatedClockwise90()
123          schematics.add(sch)
124  for colour, schems in vault_schematics.items():
125      for sch in schems.copy():
126          # Rotation
127          for i in range(3):
128              sch = sch.rotatedClockwise90()
129              vault_schematics[colour].add(sch)
130  for colour, schems in key_schematics.items():
131      for sch in schems.copy():
132          # Rotation
133          for i in range(3):
134              sch = sch.rotatedClockwise90()
135              key_schematics[colour].add(sch)
136  
137  invalid = True
138  while invalid:
139      # Datamodel
140      placed_rooms = []
141      free_connectors = []
142  
143      upcoming_calculations = []
144  
145      vault_depths = {
146          VaultColour.RED: random.randint(7,9),
147          VaultColour.GREEN: random.randint(2,4),
148          VaultColour.BLUE: random.randint(7,9),
149          VaultColour.GOLD: random.randint(10,12)
150          }
151      key_depths = {
152          VaultColour.RED: random.randint(1,2),
153          VaultColour.GREEN: random.randint(5,7),
154          VaultColour.BLUE: None,
155          VaultColour.GOLD: random.randint(7,9)
156          }
157  
158      # Hub room
159      hub_schem = Schematic("Hub Room", Vector2(51,51), [Vector2(0, 26), Vector2(0, 10), Vector2(0, 42),                  Vector2(50, 10), Vector2(50, 42),
160                                                         Vector2(26, 0), Vector2(10, 0), Vector2(42, 0), Vector2(26, 50), Vector2(10, 50), Vector2(42, 50)])
161      hub_schem.place(Vector2(-25, -25), [])
162      # schematics[0].place(Vector2(-25-9, 0))
163      # schematics[0].rotatedClockwise90()
164  
165      # Place rooms
166      print("Calculating...")
167      arrangement = collections.namedtuple("arrangement", "schematic connector topleft")
168      depth = 0
169      MAX_DEPTH = 12
170      while free_connectors and depth < MAX_DEPTH:
171          depth += 1
172          # print("Calculating at depth = " + str(depth) + "...")
173          
174          calculations = upcoming_calculations.copy()
175          for connector, parentrooms in calculations:
176              if connector not in free_connectors: continue
177  
178              free_connectors.remove(connector)
179  
180              # Check for vault/keyroom
181              additional_schems = set()
182              chosen_vault = None
183              chosen_type = None
184              for colour in VaultColour:
185                  if vault_depths[colour] is not None and depth >= vault_depths[colour]:
186                      additional_schems = vault_schematics[colour]
187                      chosen_vault = colour
188                      chosen_type = "vault"
189                      break
190                  elif key_depths[colour] is not None and depth >= key_depths[colour]:
191                      additional_schems = key_schematics[colour]
192                      chosen_vault = colour
193                      chosen_type = "key"
194                      break
195              
196              # Find suitable schematic
197              suitable_arrangements = []
198              for schematic in schematics.union(additional_schems):
199                  for other_con in schematic.connectors:
200                      for offset in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
201                          con_pos = connector + offset
202                          schem_tl = con_pos - other_con
203                          if schematic.can_place(schem_tl, connector): suitable_arrangements.append(arrangement(schematic=schematic, connector=other_con, topleft=schem_tl))
204  
205              if not suitable_arrangements:
206                  # Nothing suitable
207                  # TODO
208                  continue
209              else:
210                  chosen_arrangement = random.choice(suitable_arrangements)
211                  placed_room = chosen_arrangement.schematic.place(chosen_arrangement.topleft, parentrooms)
212  
213                  # Handle vault/keyroom placement
214                  if chosen_arrangement.schematic in additional_schems:
215                      if chosen_type == "vault":
216                          placed_room.colour = chosen_vault
217                          parent = parentrooms[1]
218                          parent.colour = chosen_vault
219                          for pchild in placed_rooms:
220                              if parent in pchild.parentrooms and pchild.colour is None:
221                                  pchild.colour = chosen_vault
222                                  
223                          vault_depths[chosen_vault] = None
224                      elif chosen_type == "key":
225                          placed_room.colour = chosen_vault
226                          key_depths[chosen_vault] = None
227  
228      # Cap off last bits
229      # TODO
230  
231      # Check validity
232      invalid = any((vault_depths[colour] is not None for colour in VaultColour)) or any((key_depths[colour] is not None for colour in VaultColour))
233  
234  # Visualise
235  print("Visualising...")
236  min_x = min(map(lambda r: r.boundingbox.left, placed_rooms))-1
237  min_y = min(map(lambda r: r.boundingbox.top, placed_rooms))-1
238  max_x = max(map(lambda r: r.boundingbox.right, placed_rooms))+1
239  max_y = max(map(lambda r: r.boundingbox.bottom, placed_rooms))+1
240  
241  min_coords = min_x, min_y
242  
243  WIN_SIZE = 700, 700
244  pixels_per_unit = WIN_SIZE[0] / (max_x - min_x), WIN_SIZE[1] / (max_y - min_y)
245  
246  import pygame
247  pygame.init()
248  pygame.display.set_mode(WIN_SIZE)
249  win = pygame.display.get_surface()
250  win.fill((0, 0, 0))
251  
252  # Draw rooms
253  for room in placed_rooms:
254      if room.colour == VaultColour.RED: colour = 255, 0, 0
255      elif room.colour == VaultColour.GREEN: colour = 0, 255, 0
256      elif room.colour == VaultColour.BLUE: colour = 0, 0, 255
257      elif room.colour == VaultColour.GOLD: colour = 255, 255, 0
258      else: colour = 255, 255, 255
259      pygame.draw.rect(win, colour, Rect((room.boundingbox.left - min_x) * pixels_per_unit[0],
260                                                  (room.boundingbox.top - min_y) * pixels_per_unit[1],
261                                                  *(room.boundingbox.size[x] * pixels_per_unit[x] for x in range(2))))
262  
263  # Draw connectors
264  for room in placed_rooms:
265      for connector in room.connectors:
266          r = Rect(0, 0, *(1 * pixels_per_unit[x] for x in range(2)))
267          r.center = tuple((connector[x] - min_coords[x] + 0.5) * pixels_per_unit[x] for x in range(2))
268          pygame.draw.rect(win, (255, 127, 0), r)
269  
270  pygame.display.update()