Files
eolrb/__main__.py
2024-08-24 10:24:27 +03:00

129 lines
4.2 KiB
Python

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_eolrb_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 = 1000
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_eolrb_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),
# "mc": solution.mc,
"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()