initial commit

__main__.py

@@ -0,0 +1,128 @@
+from collections import defaultdict
+import json
+import os
+import pickle
+import sys
+from eolr import (
+    eolrb_states_generator,
+    EOLROrientation,
+    EOLRPermutation,
+    EOLR_PERMUTATIONS,
+)
+import humanize
+from dataclasses import dataclass
+from solver import EOLRBSolution, solve_eolrb, create_prune_table
+from scorer import (
+    FingerTrickWithRegrip,
+    load_config,
+    generate_finger_tricks,
+    build_pretty_string_from_finger_tricks_with_regrips,
+)
+import heapq
+from typing import List, Any, Generator, Dict
+
+SOLUTIONS_TO_EVAL = 40
+SOLUTIONS_TO_SHOW = 3
+
+PRUNE = 10
+SOLVE = 7
+
+
+@dataclass
+class EOLRBSolutionWithFingertricks:
+    solution: EOLRBSolution
+    finger_tricks: List[FingerTrickWithRegrip]
+
+
+def get_first_n_from_generator(gen: Generator[Any, None, Any], n: int) -> List[Any]:
+    results: List[Any] = []
+    try:
+        while len(results) < n:
+            item = next(gen)
+            results.append(item)
+    except StopIteration:
+        pass
+    return results
+
+
+def load_or_generate_prune_table(file_path: str, prune_size: int) -> Dict:
+    if os.path.isfile(file_path):
+        print("opening prune table from", file_path)
+        with open(file_path, "rb") as file:
+            return pickle.load(file)
+    else:
+        print("generating prune table...")
+        prune_table = create_prune_table(prune_size)
+        print("saving prune table to", file_path)
+        with open(file_path, "wb") as file:
+            pickle.dump(prune_table, file)
+        return prune_table
+
+
+def main():
+    config = load_config("./moves.yaml")
+
+    prune_table = load_or_generate_prune_table("prune_table.pkl", PRUNE)
+    print("prune table size:", humanize.naturalsize(sys.getsizeof(prune_table)))
+
+    # TODO: make this function a method of an object that contains a list of finger tricks
+    score_func = lambda solution: sum(ft.score() for ft in solution.finger_tricks)
+    data = defaultdict(list)
+    eolrb_states, n_states = eolrb_states_generator(
+        list(EOLROrientation), EOLR_PERMUTATIONS
+    )
+    # eolrb_states, n_states = eolrb_states_generator(
+    #     [EOLROrientation.Solved], [EOLRPermutation(UR="UR", UL="UL")]
+    # )
+    for i, (eolrb_cube, ori, perm, pre_auf) in enumerate(eolrb_states):
+        print(
+            f"generating algs for {ori.name} (UR in {perm.UR}, UL in {perm.UL}, pre AUF {pre_auf}) ({i}/{n_states})"
+        )
+        alg_generator = solve_eolrb(eolrb_cube, prune_table, PRUNE)
+        solutions = get_first_n_from_generator(alg_generator, SOLUTIONS_TO_EVAL)
+        solutions_with_finger_tricks = [
+            EOLRBSolutionWithFingertricks(
+                solution, generate_finger_tricks(config, solution.alg)
+            )
+            for solution in solutions
+        ]
+        best_solutions = heapq.nsmallest(
+            min(SOLUTIONS_TO_SHOW, len(solutions_with_finger_tricks)),
+            solutions_with_finger_tricks,
+            key=score_func,
+        )
+        for solution_with_finger_tricks in best_solutions:
+            solution = solution_with_finger_tricks.solution
+            finger_tricks = solution_with_finger_tricks.finger_tricks
+
+            pretty_alg = []
+            if solution.pre_auf:
+                pretty_alg.append(solution.pre_auf)
+            pretty_alg.extend(solution.alg)
+            if solution.post_auf:
+                pretty_alg.append(solution.post_auf)
+            if solution.mc:
+                pretty_alg.append("M")
+            pretty_finger_tricks = build_pretty_string_from_finger_tricks_with_regrips(
+                finger_tricks
+            )
+
+            score = score_func(solution_with_finger_tricks)
+            alg_name = f"{ori.name} {perm.UR} {perm.UL} {pre_auf}"
+            data[alg_name].append(
+                {
+                    "alg": " ".join(pretty_alg),
+                    "stm": len(solution.alg),
+                    "score": score,
+                    "raw_alg": " ".join(solution.alg),
+                    "pre_auf": solution.pre_auf,
+                    "post_auf": solution.post_auf,
+                    "finger_tricks": pretty_finger_tricks,
+                }
+            )
+            with open("output.json", "w") as f:
+                f.write(json.dumps(data))
+
+
+if __name__ == "__main__":
+    main()