Calculates an optimal solution for the first layer on a skewb

skewb_first_layer.py

#!/usr/bin/env python3

from enum import Enum
from queue import Queue
from copy import deepcopy

NUM_FACES = 6

class Axis(Enum):
    U = 1
    R = 2
    L = 3
    B = 4

class Skewb:
    """
    Numbering scheme and transformations taken from TNoodle:
        https://github.com/thewca/tnoodle
                    +---------+
                    | 1     2 |
                U > |   0-0   |
                    | 3     4 |
          +---------+---------+---------+---------+
          | 1     2 | 1     2 | 1     2 | 1     2 |
          |   4-0   |   2-0   |   1-0   |   5-0   |
          | 3     4 | 3     4 | 3     4 | 3     4 |
          +---------+---------+---------+---------+
               ^    | 1     2 |
               FL   |   3-0   |
                    | 3     4 |
                    +---------+
    """
    def __init__(self):
        self._state = [[color] * 5 for color in range(NUM_FACES)]

    def rotate(self, axis, count=1):
        for i in range(count % 3):
            if axis == Axis.U:
                self._switch_pieces((0, 0), (1, 0), (5, 0))
                self._switch_pieces((0, 2), (1, 2), (5, 1))
                self._switch_pieces((0, 4), (1, 4), (5, 2))
                self._switch_pieces((0, 1), (1, 1), (5, 3))
                self._switch_pieces((4, 1), (2, 2), (3, 4))
            elif axis == Axis.R:
                self._switch_pieces((2, 0), (3, 0), (1, 0))
                self._switch_pieces((2, 4), (3, 2), (1, 3))
                self._switch_pieces((2, 2), (3, 1), (1, 4))
                self._switch_pieces((2, 3), (3, 4), (1, 1))
                self._switch_pieces((4, 4), (5, 3), (0, 4))
            elif axis == Axis.L:
                self._switch_pieces((4, 0), (5, 0), (3, 0))
                self._switch_pieces((4, 3), (5, 4), (3, 3))
                self._switch_pieces((4, 1), (5, 3), (3, 1))
                self._switch_pieces((4, 4), (5, 2), (3, 4))
                self._switch_pieces((2, 3), (0, 1), (1, 4))
            elif axis == Axis.B:
                self._switch_pieces((1, 0), (3, 0), (5, 0))
                self._switch_pieces((1, 4), (3, 4), (5, 3))
                self._switch_pieces((1, 3), (3, 3), (5, 1))
                self._switch_pieces((1, 2), (3, 2), (5, 4))
                self._switch_pieces((0, 2), (2, 4), (4, 3))
            else:
                assert False

    def _switch_pieces(self, first, second, third):
        (self._state[first[0]][first[1]], self._state[second[0]][second[1]], self._state[third[0]][third[1]]) = (self._state[second[0]][second[1]], self._state[third[0]][third[1]], self._state[first[0]][first[1]])

    def is_solved(self):
        for face in range(NUM_FACES):
            if not self._is_face_solved(face):
                return False
        return True

    def is_layer_solved(self):
        for face in range(NUM_FACES):
            if self._is_layer_solved(face):
                return True
        return False

    def _is_layer_solved(self, face):
        if not self._is_face_solved(face):
            return False
        #TODO: Replace with less hacky solution
        if face == 0:
            return self._state[1][1] == self._state[1][2]
        elif face == 1:
            return self._state[2][2] == self._state[2][4]
        elif face == 2:
            return self._state[1][1] == self._state[1][3]
        elif face == 3:
            return self._state[2][3] == self._state[2][4]
        elif face == 4:
            return self._state[2][1] == self._state[2][3]
        elif face == 5:
            return self._state[1][1] == self._state[1][2]

    def _is_face_solved(self, face):
        face_color = None
        for color in self._state[face]:
            if face_color == None:
                face_color = color
            elif face_color != color:
                return False
        return True

    def get_state(self):
        return self._state

    def set_state(self, state):
        self._state = state

def format_solution(path):
    solution = ""
    for step in path:
        if abs(step[1]) != 1:
            solution += abs(step[1])
        solution += step[0].name
        if step[1] < 0:
            solution += "'"
        solution += " "
    return solution[:-1]

def state_to_tuple(state):
    return tuple(tuple(l) for l in state)

def find_solutions(scramble):
    skewb = Skewb()
    for move in scramble:
        skewb.rotate(*move)

    solution_length = -1

    state = skewb.get_state()
    visited = {state_to_tuple(state)}
    queue = Queue()
    queue.put((state, []))

    while not queue.empty():
        state, path = queue.get()
        skewb.set_state(state)
        if solution_length != -1 and len(path) > solution_length:
            break
        if skewb.is_layer_solved():
            solution_length = len(path)
            yield path
        for move in Axis:
            for count in (1, -1):
                skewb.set_state(deepcopy(state))
                skewb.rotate(move, count)
                new_state = skewb.get_state()
                new_state_tuple = state_to_tuple(new_state)
                if new_state_tuple not in visited:
                    queue.put((new_state, path + [(move, count)]))
                    visited.add(new_state_tuple)

def parse_scramble(scramble_string):
    i = 0
    scramble = []
    while i < len(scramble_string):
        count = 1
        if scramble_string[i] .isdigit():
            count = int(scramble_string[i])
            i += 1
        if scramble_string[i] in Axis.__members__:
            if  i+1 >= len(scramble_string) or scramble_string[i+1] != "'":
                scramble.append((Axis[scramble_string[i]], count))
            else:
                scramble.append((Axis[scramble_string[i]], -count))
                i += 1
        i += 1
    return scramble

if __name__ == "__main__":
    while True:
        try:
            scramble_string = input("Please enter scramble sequence: ").replace(" ", "")
        except (EOFError, KeyboardInterrupt):
            print()
            break
        scramble = parse_scramble(scramble_string)

        for solution in find_solutions(scramble):
            print(format_solution(solution))