129 lines
4.2 KiB
Python
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_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()
|