from collections import defaultdict
import itertools
import copy

from dataclasses import dataclass
from cube import LSECube, condense_algorithm, reverse_algorithm
from typing import Dict, List, Tuple


@dataclass
class EOLRBSolution:
    alg: List[str]
    pre_auf: str
    post_auf: str
    post_ams: str


def lse_brute_force_generator(max_length: int):
    list1 = ["M", "M'", "M2"]
    list2 = ["U", "U'", "U2"]

    def valid_combination(comb):
        # Ensure no adjacent elements are from the same list
        if any(
            (comb[i] in list1 and comb[i + 1] in list1)
            or (comb[i] in list2 and comb[i + 1] in list2)
            for i in range(len(comb) - 1)
        ):
            return False

        # # Ensure the combination does not end with a U move
        # if comb and comb[-1] in list2:
        #     return False

        return True

    def generate_combinations():
        # Generate combinations of increasing lengths
        for length in range(1, max_length + 1):
            for comb in itertools.product(list1 + list2, repeat=length):
                if valid_combination(comb):
                    yield list(comb)

    return generate_combinations()


#
# prune table
#


def create_eolrb_prune_table(prune_depth: int):
    prune_table = defaultdict(list)
    cube = LSECube()
    prune_table[cube.eolrb_hash()] = [[]]

    generator = lse_brute_force_generator(prune_depth)

    pruned = 0

    # TODO: for every solved state (pre auf), we can apply the generator's moves and save to prune table?
    for alg in generator:
        cube.reset()
        cube.alg(" ".join(alg))
        prune_table[cube.eolrb_hash()].append(reverse_algorithm(alg))
        pruned += 1
        if pruned % 1000 == 0:
            print(pruned)

    return prune_table


def extract_pre_auf(alg: List[str]) -> Tuple[str, List[str]]:
    if alg[0].startswith("U"):
        pre_auf = alg.pop(0)
        return pre_auf, alg
    else:
        return "", alg


def solve_eolrb(cube: LSECube, prune_table: Dict, solve: int):
    for moves in lse_brute_force_generator(solve):
        c = copy.deepcopy(cube)
        c.alg(" ".join(moves))
        if c.eolrb_hash() in prune_table:
            for prune_solution in prune_table[c.eolrb_hash()]:
                # TODO: split this function, very ugly copies and very long
                prune_cube = copy.deepcopy(c)
                prune_cube.alg(" ".join(prune_solution))
                alg = condense_algorithm(moves + prune_solution)
                # TODO: this is very weird
                for _ in range(5):
                    alg = condense_algorithm(alg)

                if len(alg) > 2:
                    pre_auf = ""
                    if alg[0].startswith("U"):
                        pre_auf = alg.pop(0)
                    post_auf = ""
                    if alg[-1].startswith("U"):
                        post_auf = alg.pop(-1)
                    post_ams = prune_cube.fix_m_slice()
                    yield EOLRBSolution(alg, pre_auf, post_auf, post_ams)
        # for auf in ["U", "U2", "U'", ""]:
        #     c = copy.deepcopy(cube)
        #     auf_alg = f"{' '.join(alg)} {auf}"
        #     c.alg(auf_alg)
        #     if eolrb_solved(c):
        #         yield alg