from typing import Dict, List, Set from collections import deque from worlds.factorio.Options import TechTreeLayout funnel_layers = {TechTreeLayout.option_small_funnels: 3, TechTreeLayout.option_medium_funnels: 4, TechTreeLayout.option_large_funnels: 5} funnel_slice_sizes = {TechTreeLayout.option_small_funnels: 6, TechTreeLayout.option_medium_funnels: 10, TechTreeLayout.option_large_funnels: 15} def get_shapes(factorio_world) -> Dict[str, List[str]]: world = factorio_world.world player = factorio_world.player prerequisites: Dict[str, Set[str]] = {} layout = world.tech_tree_layout[player].value custom_technologies = factorio_world.custom_technologies tech_names: List[str] = list(set(custom_technologies) - world.worlds[player].static_nodes) tech_names.sort() world.random.shuffle(tech_names) if layout == TechTreeLayout.option_single: pass elif layout == TechTreeLayout.option_small_diamonds: slice_size = 4 while len(tech_names) > slice_size: slice = tech_names[:slice_size] tech_names = tech_names[slice_size:] slice.sort(key=lambda tech_name: len(custom_technologies[tech_name].get_prior_technologies())) diamond_0, diamond_1, diamond_2, diamond_3 = slice # 0 | # 1 2 | # 3 V prerequisites[diamond_3] = {diamond_1, diamond_2} prerequisites[diamond_2] = prerequisites[diamond_1] = {diamond_0} elif layout == TechTreeLayout.option_medium_diamonds: slice_size = 9 while len(tech_names) > slice_size: slice = tech_names[:slice_size] tech_names = tech_names[slice_size:] slice.sort(key=lambda tech_name: len(custom_technologies[tech_name].get_prior_technologies())) # 0 | # 1 2 | # 3 4 5 | # 6 7 | # 8 V prerequisites[slice[1]] = {slice[0]} prerequisites[slice[2]] = {slice[0]} prerequisites[slice[3]] = {slice[1]} prerequisites[slice[4]] = {slice[1], slice[2]} prerequisites[slice[5]] = {slice[2]} prerequisites[slice[6]] = {slice[3], slice[4]} prerequisites[slice[7]] = {slice[4], slice[5]} prerequisites[slice[8]] = {slice[6], slice[7]} elif layout == TechTreeLayout.option_large_diamonds: slice_size = 16 while len(tech_names) > slice_size: slice = tech_names[:slice_size] tech_names = tech_names[slice_size:] slice.sort(key=lambda tech_name: len(custom_technologies[tech_name].get_prior_technologies())) # 0 | # 1 2 | # 3 4 5 | # 6 7 8 9 | # 10 11 12 | # 13 14 | # 15 | prerequisites[slice[1]] = {slice[0]} prerequisites[slice[2]] = {slice[0]} prerequisites[slice[3]] = {slice[1]} prerequisites[slice[4]] = {slice[1], slice[2]} prerequisites[slice[5]] = {slice[2]} prerequisites[slice[6]] = {slice[3]} prerequisites[slice[7]] = {slice[3], slice[4]} prerequisites[slice[8]] = {slice[4], slice[5]} prerequisites[slice[9]] = {slice[5]} prerequisites[slice[10]] = {slice[6], slice[7]} prerequisites[slice[11]] = {slice[7], slice[8]} prerequisites[slice[12]] = {slice[8], slice[9]} prerequisites[slice[13]] = {slice[10], slice[11]} prerequisites[slice[14]] = {slice[11], slice[12]} prerequisites[slice[15]] = {slice[13], slice[14]} elif layout == TechTreeLayout.option_small_pyramids: slice_size = 6 while len(tech_names) > slice_size: slice = tech_names[:slice_size] tech_names = tech_names[slice_size:] slice.sort(key=lambda tech_name: len(custom_technologies[tech_name].get_prior_technologies())) # 0 | # 1 2 | # 3 4 5 | prerequisites[slice[1]] = {slice[0]} prerequisites[slice[2]] = {slice[0]} prerequisites[slice[3]] = {slice[1]} prerequisites[slice[4]] = {slice[1], slice[2]} prerequisites[slice[5]] = {slice[2]} elif layout == TechTreeLayout.option_medium_pyramids: slice_size = 10 while len(tech_names) > slice_size: slice = tech_names[:slice_size] tech_names = tech_names[slice_size:] slice.sort(key=lambda tech_name: len(custom_technologies[tech_name].get_prior_technologies())) # 0 | # 1 2 | # 3 4 5 | # 6 7 8 9 | prerequisites[slice[1]] = {slice[0]} prerequisites[slice[2]] = {slice[0]} prerequisites[slice[3]] = {slice[1]} prerequisites[slice[4]] = {slice[1], slice[2]} prerequisites[slice[5]] = {slice[2]} prerequisites[slice[6]] = {slice[3]} prerequisites[slice[7]] = {slice[3], slice[4]} prerequisites[slice[8]] = {slice[4], slice[5]} prerequisites[slice[9]] = {slice[5]} elif layout == TechTreeLayout.option_large_pyramids: slice_size = 15 while len(tech_names) > slice_size: slice = tech_names[:slice_size] tech_names = tech_names[slice_size:] slice.sort(key=lambda tech_name: len(custom_technologies[tech_name].get_prior_technologies())) # 0 | # 1 2 | # 3 4 5 | # 6 7 8 9 | # 10 11 12 13 14 | prerequisites[slice[1]] = {slice[0]} prerequisites[slice[2]] = {slice[0]} prerequisites[slice[3]] = {slice[1]} prerequisites[slice[4]] = {slice[1], slice[2]} prerequisites[slice[5]] = {slice[2]} prerequisites[slice[6]] = {slice[3]} prerequisites[slice[7]] = {slice[3], slice[4]} prerequisites[slice[8]] = {slice[4], slice[5]} prerequisites[slice[9]] = {slice[5]} prerequisites[slice[10]] = {slice[6]} prerequisites[slice[11]] = {slice[6], slice[7]} prerequisites[slice[12]] = {slice[7], slice[8]} prerequisites[slice[13]] = {slice[8], slice[9]} prerequisites[slice[14]] = {slice[9]} elif layout in funnel_layers: slice_size = funnel_slice_sizes[layout] world.random.shuffle(tech_names) while len(tech_names) > slice_size: tech_names = tech_names[slice_size:] current_tech_names = tech_names[:slice_size] layer_size = funnel_layers[layout] previous_slice = [] current_tech_names.sort(key=lambda tech_name: len(custom_technologies[tech_name].get_prior_technologies())) for layer in range(funnel_layers[layout]): slice = current_tech_names[:layer_size] current_tech_names = current_tech_names[layer_size:] if previous_slice: for i, tech_name in enumerate(slice): prerequisites.setdefault(tech_name, set()).update(previous_slice[i:i+2]) previous_slice = slice layer_size -= 1 elif layout == TechTreeLayout.option_trees: # 0 | # 1 2 | # 3 | # 4 5 6 7 | # 8 | # 9 10 11 12 13 14 | # 15 | # 16 | slice_size = 17 while len(tech_names) > slice_size: slice = tech_names[:slice_size] tech_names = tech_names[slice_size:] slice.sort(key=lambda tech_name: len(custom_technologies[tech_name].get_prior_technologies())) prerequisites[slice[1]] = {slice[0]} prerequisites[slice[2]] = {slice[0]} prerequisites[slice[3]] = {slice[1], slice[2]} prerequisites[slice[4]] = {slice[3]} prerequisites[slice[5]] = {slice[3]} prerequisites[slice[6]] = {slice[3]} prerequisites[slice[7]] = {slice[3]} prerequisites[slice[8]] = {slice[4], slice[5], slice[6], slice[7]} prerequisites[slice[9]] = {slice[8]} prerequisites[slice[10]] = {slice[8]} prerequisites[slice[11]] = {slice[8]} prerequisites[slice[12]] = {slice[8]} prerequisites[slice[13]] = {slice[8]} prerequisites[slice[14]] = {slice[8]} prerequisites[slice[15]] = {slice[9], slice[10], slice[11], slice[12], slice[13], slice[14]} prerequisites[slice[16]] = {slice[15]} elif layout == TechTreeLayout.option_choices: tech_names.sort(key=lambda tech_name: len(custom_technologies[tech_name].get_prior_technologies())) current_choices = deque([tech_names[0]]) tech_names = tech_names[1:] while len(tech_names) > 1: source = current_choices.pop() choices = tech_names[:2] tech_names = tech_names[2:] for choice in choices: prerequisites[choice] = {source} current_choices.extendleft(choices) else: raise NotImplementedError(f"Layout {layout} is not implemented.") world.tech_tree_layout_prerequisites[player] = prerequisites return prerequisites