3 files changed, 107 insertions, 0 deletions
diff --git a/sbdata/tasks/__init__.py b/sbdata/tasks/__init__.py
new file mode 100644
index 0000000..737c65c
--- /dev/null
+++ b/sbdata/tasks/__init__.py
@@ -0,0 +1,2 @@
+from .fetch_dungeon_loot import fetch_dungeon_loot
+from .calculate_mapping_function import calculate_mapping_function
diff --git a/sbdata/tasks/calculate_mapping_function.py b/sbdata/tasks/calculate_mapping_function.py
new file mode 100644
index 0000000..37185a8
--- /dev/null
+++ b/sbdata/tasks/calculate_mapping_function.py
@@ -0,0 +1,31 @@
+import pathlib
+import random
+import numpy as np
+from sklearn.linear_model import LinearRegression
+
+from sbdata.task import register_task, Arguments
+
+
+@register_task("Calculate Map Coordinate Function")
+def calculate_mapping_function(args: Arguments):
+    csv = args.get_arg("Coordinate CSV", "coords", pathlib.Path)
+    points = [[int(x) for x in y.split(",")] for y in csv.read_text().splitlines()[1:]]
+    xs = [(a[0], a[2]) for a in points]
+    zs = [(a[1], a[3]) for a in points]
+    random.shuffle(xs)
+    random.shuffle(zs)
+    find_best_function_for("X", xs)
+    find_best_function_for("Z", zs)
+
+
+def find_best_function_for(label: str, l: list[tuple[int, int]]):
+    x = np.array([a[0] for a in l]).reshape((-1, 1))
+    y = np.array([a[1] for a in l])
+    model = LinearRegression()
+    model.fit(x, y)
+    print(f'------------')
+    print(f' {label} Coordinate:')
+    print(f"   Score: {model.score(x, y)}")
+    print(f"   Slope: {model.coef_[0]}")
+    print(f"   Intercept: {model.intercept_}")
+    print(f'------------')
diff --git a/sbdata/tasks/fetch_dungeon_loot.py b/sbdata/tasks/fetch_dungeon_loot.py
new file mode 100644
index 0000000..bf7ef64
--- /dev/null
+++ b/sbdata/tasks/fetch_dungeon_loot.py
@@ -0,0 +1,74 @@
+import ast
+import dataclasses
+import json
+import re
+import sys
+import typing
+
+from sbdata.repo import find_item_by_name, Item
+from sbdata.task import register_task, Arguments
+from sbdata.wiki import get_wiki_sources_by_title
+
+
+@dataclasses.dataclass
+class DungeonDrop:
+    item: Item
+    floor: int
+    chest: str
+    cost: int
+    drop_chances: dict[str, str]
+
+    def get_drop_chance(self, has_s_plus: bool, talisman_level: int, boss_luck: int):
+        drop_identifier = "S" + ('+' if has_s_plus else '') + 'ABCD'[talisman_level] + str(len([i for i in [0, 1, 3, 5, 10] if i >= boss_luck]))
+        return self.drop_chances.get(drop_identifier)
+
+
+default_chest_costs: dict[str, dict[int, int]] = dict(
+    Wood={7: 0},
+    Gold={1: 25_000, 2: 50_000, 7: 100_000},
+    Diamond={1: 50_000, 2: 100_000, 7: 250_000},
+    Emerald={1: 100_000, 2: 250_000, 7: 500_000},
+    Obsidian={1: 250_000, 2: 500_000, 7: 1_000_000},
+    Bedrock={4: 4, 7: 2_000_000}
+)
+
+
+@register_task("Fetch Dungeon Loot")
+def fetch_dungeon_loot(args: Arguments):
+    items = []
+    for floor_name, floor in get_wiki_sources_by_title(*[f'Template:Catacombs Floor {f} Loot Master' for f in ['I', 'II', 'III', 'IV', 'V', 'VI', 'VII']]).items():
+        print(f"Page: {floor_name}")
+        for template in floor.filter_templates():
+            if template.name.strip() == 'Dungeon Chest Table/Row':
+                item = None
+                ifloor = None
+                chest = None
+                cost = None
+                drop_chances = {}
+
+                for param in template.params:
+                    attr_name = param.name.nodes[0].strip()
+                    attr_value = param.value.nodes[0].strip()
+                    if attr_name == 'item':
+                        if item is None:
+                            item = find_item_by_name(attr_value)
+                    elif attr_name == 'customlink':
+                        if item is None:
+                            item = find_item_by_name(attr_value.split('#')[-1])
+                    elif attr_name == 'cost':
+                        cost = int(attr_value.replace(',', ''))
+                    elif attr_name == 'chest':
+                        chest = attr_value
+                    elif attr_name == 'floor':
+                        ifloor = int(attr_value)
+                    elif attr_name.startswith("S"):
+                        drop_chances[attr_name] = attr_value
+                if item is None or ifloor is None or chest is None or cost is None:
+                    print('WARNING: Missing data for item: ' + str(template))
+                else:
+                    if cost == 0:
+                        defaults = default_chest_costs[chest]
+                        cost = defaults[min(f for f in defaults.keys() if f >= ifloor)]
+                    items.append(DungeonDrop(item, ifloor, chest, cost, drop_chances))
+    return items
+