From c42f53d64f734b53d84d9642d030cd82c21938b2 Mon Sep 17 00:00:00 2001
From: Fabian Dill <fabian.dill@web.de>
Date: Sat, 18 Dec 2021 13:01:24 +0100
Subject: [PATCH] Factorio: add some more tech tree shapes

---
 worlds/factorio/Options.py |  2 ++
 worlds/factorio/Shapes.py  | 50 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/worlds/factorio/Options.py b/worlds/factorio/Options.py
index 38991102..b8b06cb2 100644
--- a/worlds/factorio/Options.py
+++ b/worlds/factorio/Options.py
@@ -94,6 +94,8 @@ class TechTreeLayout(Choice):
     option_small_funnels = 7
     option_medium_funnels = 8
     option_large_funnels = 9
+    option_trees = 10
+    option_choices = 11
     default = 0
 
 
diff --git a/worlds/factorio/Shapes.py b/worlds/factorio/Shapes.py
index 12e9f0b6..4a622de9 100644
--- a/worlds/factorio/Shapes.py
+++ b/worlds/factorio/Shapes.py
@@ -1,4 +1,5 @@
 from typing import Dict, List, Set
+from collections import deque
 
 from worlds.factorio.Options import TechTreeLayout
 
@@ -189,6 +190,55 @@ def get_shapes(factorio_world) -> Dict[str, List[str]]:
                         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