From 194358622113c3ece65a3ff794a81221ecae24b9 Mon Sep 17 00:00:00 2001
From: Fabian Dill <fabian.dill@web.de>
Date: Sat, 10 Apr 2021 19:34:30 +0200
Subject: [PATCH] Factorio: add medium_diamonds and pyramid tech tree layouts

---
 Options.py                  |  2 ++
 worlds/factorio/__init__.py | 64 +++++++++++++++++++++++++++++++------
 2 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/Options.py b/Options.py
index 5f4009b7..73eb65f4 100644
--- a/Options.py
+++ b/Options.py
@@ -274,6 +274,8 @@ class FreeSamples(Choice):
 class TechTreeLayout(Choice):
     option_single = 0
     option_small_diamonds = 1
+    option_medium_diamonds = 2
+    option_pyramid = 3
     default = 0
 
 class Visibility(Choice):
diff --git a/worlds/factorio/__init__.py b/worlds/factorio/__init__.py
index bd773c83..13f955b3 100644
--- a/worlds/factorio/__init__.py
+++ b/worlds/factorio/__init__.py
@@ -1,5 +1,5 @@
 import logging
-from typing import List, Dict
+from typing import List, Dict, Set
 
 from BaseClasses import Region, Entrance, Location, MultiWorld, Item
 from Options import TechTreeLayout
@@ -37,22 +37,68 @@ def factorio_create_regions(world: MultiWorld, player: int):
 
 
 def get_shapes(world: MultiWorld, player: int) -> Dict[str, List[str]]:
-    prerequisites = {}
-    if world.tech_tree_layout[player].value == TechTreeLayout.option_small_diamonds:
+    prerequisites:Dict[str, Set[str]] = {}
+    layout = world.tech_tree_layout[player].value
+    if layout == TechTreeLayout.option_small_diamonds:
+        slice_size = 4
         tech_names: List[str] = list(set(technology_table)-static_nodes)
         tech_names.sort()
         world.random.shuffle(tech_names)
-        while len(tech_names) > 4:
-            slice = tech_names[:4]
-            tech_names = tech_names[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(technology_table[tech_name].ingredients))
             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]
+            prerequisites[diamond_3] = {diamond_1, diamond_2}
+            prerequisites[diamond_2] = prerequisites[diamond_1] = {diamond_0}
+    elif layout == TechTreeLayout.option_medium_diamonds:
+        slice_size = 9
+        tech_names: List[str] = list(set(technology_table)-static_nodes)
+        tech_names.sort()
+        world.random.shuffle(tech_names)
+        while len(tech_names) > slice_size:
+            slice = tech_names[:slice_size]
+            tech_names = tech_names[slice_size:]
+            slice.sort(key=lambda tech_name: len(technology_table[tech_name].ingredients))
+
+            #     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_pyramid:
+        slice_size = 1
+        tech_names: List[str] = list(set(technology_table)-static_nodes)
+        tech_names.sort()
+        world.random.shuffle(tech_names)
+        tech_names.sort(key=lambda tech_name: len(technology_table[tech_name].ingredients))
+        previous_slice = []
+        while len(tech_names) > slice_size:
+            slice = tech_names[:slice_size]
+            tech_names = tech_names[slice_size:]
+            for i, tech_name in enumerate(previous_slice):
+                prerequisites.setdefault(slice[i], set()).add(tech_name)
+                prerequisites.setdefault(slice[i+1], set()).add(tech_name)
+            previous_slice = slice
+            slice_size += 1
+
     world.tech_tree_layout_prerequisites[player] = prerequisites
     return prerequisites
 
@@ -68,7 +114,7 @@ def set_rules(world: MultiWorld, player: int):
             Rules.set_rule(location, technology.build_rule(allowed_packs, player))
             prequisites = shapes.get(tech_name)
             if prequisites:
-                locations = [world.get_location(requisite, player) for requisite in prequisites]
+                locations = {world.get_location(requisite, player) for requisite in prequisites}
                 Rules.add_rule(location, lambda state,
                                                 locations=locations: all(state.can_reach(loc) for loc in locations))