mall-solver/engine.py

from typing import List, Set, Tuple
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from typing import List, Any

class Item:
	def __init__(self, name: str) -> None:
		self.name = name

class Recipe:
	def __init__(self, inputs: List[Item], outputs: List[Item], name: str) -> None:
		self.inputs = inputs
		self.outputs = outputs
		self.name = name

class World:
	def __init__(
			self,
			recipes: Set[Recipe],
			items: Set[Item],
			bus: Set[Item],
			size: int
		) -> None:
		self.recipes = recipes
		self.items = items
		self.bus = bus
		self.size = size


def find_valid_recipes(output: Item, world: World) -> List[Recipe]:
	valid_recipes = []
	for recipe in world.recipes:
		if output in recipe.value.outputs:
			valid_recipes.append(recipe.value)
	return valid_recipes


class Assembler:
	def __init__(self, recipe: Recipe, world: World) -> None:
		self.recipe = recipe
		self.input_links = [] # List[Assembler]
		self.world = world

	def link_input(self, assembler: "Assembler") -> None:
		self.input_links.append(assembler)

	def get_unlinked_inputs(self) -> List[Item]:
		return [
			item
			for item in self.recipe.inputs
			if item not in [
				assembler.recipe.outputs
				for assembler in self.input_links
			].flatten()
			and item not in self.world.bus
		]

	def is_solved(self) -> bool:
		return len(self.get_unlinked_inputs()) == 0

class Graph:
	def __init__(self, world: World) -> None:
		self.grid: List[List["Assembler"]] = [
			[
				None for _ in range(world.size)
			] for _ in range(world.size)
		]
		self.world = world

	def create_graph(world: World, recipe: Recipe) -> "Graph":
		graph = Graph(world)
		graph.grid[1][world.size // 2] = recipe
		return graph
	
	def get_coordinates_of_assembler(self, assembler: Assembler) -> (int, int):
		for i in range(len(self.grid)):
			for j in range(len(self.grid)):
				if self.grid[i][j] == assembler:
					return (i, j)
		return None

	def clone(self) -> "Graph":
		new_graph = Graph(len(self.grid))
		for i in range(len(self.grid)):
			for j in range(len(self.grid)):
				if self.grid[i][j] is not None:
					assembler = self.grid[i][j]
					new_assembler = Assembler(assembler.recipe)

		# re-link assemblers in cloned graph
		for i in range(len(self.grid)):
			for j in range(len(self.grid)):
				if self.grid[i][j] is not None:
					assembler = self.grid[i][j]
					new_assembler = Assembler(assembler.recipe)
					for link in assembler.input_links:
						coords = self.get_coordinates_of_assembler(link)
						new_assembler.link_input(new_graph.grid[coords[0]][coords[1]])
					
		return new_graph
	
	def add_assembler(self, recipe: Assembler, x: int, y: int) -> "Graph":
		graph = self.clone()
		graph.grid[x][y] = recipe
		return graph
	
	def is_solved_at(self, x: int, y: int) -> bool:
		if self.grid[x][y] is None:
			return False
		return self.grid[x][y].is_solved()

	def is_solved(self) -> bool:
		for i in range(len(self.grid)):
			for j in range(len(self.grid)):
				if not self.is_solved_at(i, j):
					return False
		return True
	
	def get_adjacent_coordinates(self, x: int, y: int) -> List[Tuple[int, int]]:
		return [
			(x + 1, y) if x < len(self.grid) - 1 else None,
			(x - 1, y) if x > 0 else None,
			(x, y + 1) if y < len(self.grid) - 1 else None,
			(x, y - 1) if y > 0 else None,
			(x + 1, y + 1) if x < len(self.grid) - 1 and y < len(self.grid) - 1 and y % 2 == 0 else None,
			(x + 1, y - 1) if x < len(self.grid) - 1 and y > 0 and y % 2 == 0 else None,
			(x - 1, y + 1) if x > 0 and y < len(self.grid) - 1 and y % 2 == 1 else None,
			(x - 1, y - 1) if x > 0 and y > 0 and y % 2 == 1 else None,
		].filter(lambda x: x is not None)
	
	def are_coordinates_adjacent(self, a: (int, int), b: (int, int)) -> bool:
		adjacent_coordinates = self.get_adjacent_coordinates(a[0], a[1])
		return b in adjacent_coordinates

	def draw(self) -> None:
		plt.figure()

		def convert_coords(p: (int, int)) -> (int, int):
			(x, y) = p
			return (x + 0.5 if y % 2 == 1 else x, y * 0.86602540378)


		for row in self.grid:
			for assembler in row:
				if assembler is not None:
					(x, y) = convert_coords(self.get_coordinates_of_assembler(assembler))
					plt.gca().add_patch(Rectangle((x - 0.5, y - 0.5), 1, 1, fill=True))
					plt.annotate("thing", (x, y), textcoords="offset points", xytext=(0,5), ha='center')

		# # Adding edges (arrows) between nodes
		# plt.arrow(0, 0, 0.9, 0.9, head_width=0.05, head_length=0.1, fc='k', ec='k')
		# plt.arrow(1, 1, 0.9, -0.9, head_width=0.05, head_length=0.1, fc='k', ec='k')

		# Setting the plot limits
		plt.xlim(-1, 11)
		plt.ylim(-1, 11)

		# Show the plot
		plt.show()

	def get_solutions(self) -> List["Graph"]:
		if self.is_solved():
			return []
		
		graphs: List["Graph"] = []

		for x in range(len(self.grid)):
			for y in range(len(self.grid)):
				assembler = self.grid[x][y]
				if assembler == None:
					continue
				print(assembler)
				# for item in assembler.get_unlinked_inputs():
				# 	for recipes in find_valid_recipes(item):
				# 		graph = self.add_assembler(assembler, x, y)
				# 		graphs.append(graph)