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()