#!/usr/bin/python3
import copy
from time import sleep

# Originally these were just letters (X and O for even and odd snake, and x and
# o for even and odd empty cell), but as long as they're distinct, they do what
# we need. So these include ANSI color codes, which colors the snake red and
# blue.
EMPTY = "_"
ODD = "\033[34m_\033[0m"
EVEN = "\033[31m_\033[0m"
ODD_SNAKE = "\033[34mX\033[0m"
EVEN_SNAKE = "\033[31mX\033[0m"

# Here we define a few helper functions to make our lives easier later.
attempts = 0
def print_grid(grid):
    global attempts
    print(f"Attempt {attempts}:")
    attempts += 1
    for row in range(9):
        print(" ".join([r[row] for r in grid]))
    print()

def opposite_parity(x):
    if x == EVEN:
        return ODD
    else:
        return EVEN

def snake(x):
    if x == EVEN:
        return EVEN_SNAKE
    else:
        return ODD_SNAKE

def has_two_adjacent_snakes(grid, cell):
    """
    has_two_adjacent_snakes checks to see if two or more of the adjacent
    cells are filled with snake. If one is filled, we must be at the head of a
    snake; if two are filled we're along the body. No more than two can be
    filled, because then we'd have broken the snake; it would fork.
    """
    adjacent_snakes = 0
    if cell[1] > 0 and grid[cell[0]][cell[1] - 1] in [ODD_SNAKE, EVEN_SNAKE]: # north
        adjacent_snakes += 1

    if cell[0] < 8 and grid[cell[0] + 1][cell[1]] in [ODD_SNAKE, EVEN_SNAKE]: # east
        adjacent_snakes += 1

    if cell[1] < 8 and grid[cell[0]][cell[1] + 1] in [ODD_SNAKE, EVEN_SNAKE]: # south
        adjacent_snakes += 1
    
    if cell[0] > 0 and grid[cell[0] - 1][cell[1]] in [ODD_SNAKE, EVEN_SNAKE]: # west
        adjacent_snakes += 1
    return adjacent_snakes >= 2

def check_grid_parity(grid):
    """
    check_grid_parity employs a multi-pass approach. For each pass, it counts
    the even and odd cells in each row, column, and box. Then, if it has
    either four even or five odd cells, it knows that all of the others must be
    of the opposite parity, so it fills them in. This, however, may complete
    other rows/columns/boxes, so we'll do another pass after this one is done.
    We continue to do passes until either something breaks (returning False),
    or nothing is left to change.

    This is a lot of repetitive but just-slightly-different-each-time logic, so
    it can't easily be condensed. However, it's not particularly complicated,
    just rather verbose.
    """
    do_another_pass=False
    # check each row
    for row in grid:
        even_count = odd_count = 0
        for cell in row:
            if cell == EVEN or cell == EVEN_SNAKE:
                even_count += 1
            elif cell == ODD or cell == ODD_SNAKE:
                odd_count += 1
        if even_count > 4 or odd_count > 5:
            #print("broken on row")
            return False
        elif even_count == 4:
            for cell_index in range(len(row)):
                if row[cell_index] == EMPTY:
                    #print("filling row with odd")
                    row[cell_index] = ODD
                    do_another_pass=True
        elif odd_count == 5:
            for cell_index in range(len(row)):
                if row[cell_index] == EMPTY:
                    #print("filling row with even")
                    row[cell_index] = EVEN
                    do_another_pass=True

    # check each column
    for col_index in range(9):
        even_count = odd_count = 0
        for row in grid:
            if row[col_index] == EVEN or row[col_index] == EVEN_SNAKE:
                even_count += 1
            elif row[col_index] == ODD or row[col_index] == ODD_SNAKE:
                odd_count += 1
        #print(f"row {col_index}: even: {even_count}, odd: {odd_count}")
        if even_count > 4 or odd_count > 5:
            #print("broken on column")
            return False
        elif even_count == 4:
            for row in grid:
                if row[col_index] == EMPTY:
                    row[col_index] = ODD
                    #print("filling col with odd")
                    do_another_pass=True
        elif odd_count == 5:
            for row in grid:
                if row[col_index] == EMPTY:
                    row[col_index] = EVEN
                    #print("filling col with even")
                    do_another_pass=True

    # check each box
    for box_x in range(3):
        for box_y in range(3):
            even_count = odd_count = 0
            for x in range(3*box_x, 3*box_x+3):
                for y in range(3*box_y, 3*box_y + 3):
                    if grid[x][y] == EVEN or grid[x][y] == EVEN_SNAKE:
                        even_count += 1
                    elif grid[x][y] == ODD or grid[x][y] == ODD_SNAKE:
                        odd_count += 1
            if even_count > 4 or odd_count > 5:
                #print(f"bad parity on box {box_x},{box_y}")
                return False
            elif even_count == 4:
                for x in range(3*box_x, 3*box_x+3):
                    for y in range(3*box_y, 3*box_y + 3):
                        if grid[x][y] == EMPTY:
                            grid[x][y] = ODD
                            #print("filling box with odd")
                            do_another_pass=True
            elif odd_count == 5:
                for x in range(3*box_x, 3*box_x+3):
                    for y in range(3*box_y, 3*box_y + 3):
                        if grid[x][y] == EMPTY:
                            grid[x][y] = EVEN
                            #print("filling box with even")
                            do_another_pass=True

    if do_another_pass:
        return check_grid_parity(grid)
    else:
        return True

def has_adjacent_cells_of_parity(grid, cell, parity):
    adjacent_cells_of_parity = 0
    if cell[1] > 0 and grid[cell[0]][cell[1] - 1] == parity: # north
        adjacent_cells_of_parity += 1

    if cell[0] < 8 and grid[cell[0] + 1][cell[1]] == parity: # east
        adjacent_cells_of_parity += 1

    if cell[1] < 8 and grid[cell[0]][cell[1] + 1] == parity: # south
        adjacent_cells_of_parity += 1
    
    if cell[0] > 0 and grid[cell[0] - 1][cell[1]] == parity: # west
        adjacent_cells_of_parity += 1
    return adjacent_cells_of_parity > 0


def fill_empty_cells_around(grid, cell, value):
    """
    fill_empty_cells_around is applied to the "neck" (that is, the cell next to
    the head) of the snake. Supposing we have a snake like this, where E and O
    represent even and odd snake cells, and e and o non-snake cells:
        e o e _
        E O E O E
        e o e _
    then fill_empty_cells_around will be used with the above as `grid`, the
    last O in the row as `cell`, and "o" as `value`.
    """

    if cell[1] > 0 and grid[cell[0]][cell[1] - 1] == opposite_parity(value): # north
        raise ValueError
    elif cell[1] > 0 and grid[cell[0]][cell[1] - 1] == EMPTY: # north
        grid[cell[0]][cell[1] - 1] = value

    if cell[0] < 8 and grid[cell[0] + 1][cell[1]] == opposite_parity(value): # east
        raise ValueError
    elif cell[0] < 8 and grid[cell[0] + 1][cell[1]] == EMPTY: # east
        grid[cell[0] + 1][cell[1]] = value

    if cell[1] < 8 and grid[cell[0]][cell[1] + 1] == opposite_parity(value): # south
        raise ValueError
    elif cell[1] < 8 and grid[cell[0]][cell[1] + 1] == EMPTY: # south
        grid[cell[0]][cell[1] + 1] = value
    
    if cell[0] > 0 and grid[cell[0] - 1][cell[1]] == opposite_parity(value): # west
        raise ValueError
    elif cell[0] > 0 and grid[cell[0] - 1][cell[1]] == EMPTY: # west
        grid[cell[0] - 1][cell[1]] = value


def recursive_dfs(grid, head, neck, last_parity):

    # disable this if you don't want to see the solving visualization. Then you
    # can also disable the sleep(10) about 20 lines below, as we'll only be
    # printing the solutions.
    print_grid(grid)

    # check if we broke the grid:
    if not check_grid_parity(grid):
        return # try the next case

    # give the terminal a chance to keep up. On a 60hz display, this should
    # leave the grid on display for 1-2 frames.
    sleep(0.03)

    # base case: have we solved it (reached every box at least once)
    snakes_everywhere = True
    for box_x in range(3):
        for box_y in range(3):
            snake_in_grid = False
            for x in range(3*box_x, 3*box_x+3):
                for y in range(3*box_y, 3*box_y + 3):
                    if grid[x][y] in [ODD_SNAKE, EVEN_SNAKE]:
                        snake_in_grid = True
            if not snake_in_grid:
                snakes_everywhere = False
    if (snakes_everywhere and
        not has_adjacent_cells_of_parity(grid, head, opposite_parity(last_parity)) and
        not has_adjacent_cells_of_parity(grid, neck, last_parity)):
        print("Solved!")
        print_grid(grid)
        sleep(10) # give us a chance to look at what we found
        # note that we don't return here; we can try to see if there's a longer
        # snake that's also a solution.

    # if the cell is empty, then it doesn't have a parity, and therefore a snake
    # does not pass by in an adjacent cell, so it's a valid cell to try to enter
    #
    # For each of the directions (ordered north, east, south, and west), we make
    # sure we're not along that wall (if we're along the north wall, of course
    # we can't go north); we verify that either cell to the $direction of the
    # current head is empty (no constraints, so we can enter), or that it's the
    # same color as the one we're entering _and_ doesn't have two adjacent
    # snakes (that is, another besides the one we're entering from).
    #
    # If we have a neck cell (which we always do except on the very first cell
    # of the snake), we fill the cells around the neck with the appropriate
    # color. Note that we're filling the cells around the neck and not the
    # head. Although it does take one step longer to fail than would otherwise
    # be necessary, this saves us a lot of headache later: Say our head is even.
    # Then all of the cells around it except the neck cell (if applicable) and
    # the cell that will be the head next iteration (if the snake is not
    # complete) will also be even, to make the snake unambiguous. However, then
    # next iteration we will need to continue the snake into one of those even
    # cells---which one? Conversely, if we simply apply it to the neck after the
    # location of the new head has been decided, we avoid all this difficulty.
    #
    # We repeat this process for each possible direction.
    if head[1] > 0 and (grid[head[0]][head[1] - 1] == EMPTY or
        (grid[head[0]][head[1] - 1] == opposite_parity(last_parity) and
            not has_two_adjacent_snakes(grid, (head[0], head[1] - 1)))): # north
        new_grid = copy.deepcopy(grid)
        new_grid[head[0]][head[1] - 1] = snake(opposite_parity(last_parity))
        if neck is not None:
            try:
                fill_empty_cells_around(new_grid, neck, opposite_parity(last_parity))
            except ValueError:  # Cell is required to be both even and odd---abort!
                return # the neck is broken so this _can't_ be the head
        recursive_dfs(new_grid, (head[0], head[1] - 1), head, opposite_parity(last_parity))

    if head[0] < 8 and (grid[head[0] + 1][head[1]] == EMPTY or
        (grid[head[0] + 1][head[1]] == opposite_parity(last_parity) and
            not has_two_adjacent_snakes(grid, (head[0] + 1, head[1])))): # east
        new_grid = copy.deepcopy(grid)
        new_grid[head[0] + 1][head[1]] = snake(opposite_parity(last_parity))
        if neck is not None:
            try:
                fill_empty_cells_around(new_grid, neck, opposite_parity(last_parity))
            except ValueError:  # Cell is required to be both even and odd---abort!
                return
        recursive_dfs(new_grid, (head[0] + 1, head[1]), head, opposite_parity(last_parity))

    if head[0] > 0 and (grid[head[0] - 1][head[1]] == EMPTY or
        (grid[head[0] - 1][head[1]] == opposite_parity(last_parity) and
            not has_two_adjacent_snakes(grid, (head[0] - 1, head[1])))): # west
        new_grid = copy.deepcopy(grid)
        new_grid[head[0] - 1][head[1]] = snake(opposite_parity(last_parity))
        if neck is not None:
            try:
                fill_empty_cells_around(new_grid, neck, opposite_parity(last_parity))
            except ValueError:  # Cell is required to be both even and odd---abort!
                return
        recursive_dfs(new_grid, (head[0] - 1, head[1]), head, opposite_parity(last_parity))

    if head[1] < 8 and (grid[head[0]][head[1] + 1] == EMPTY or
        (grid[head[0]][head[1] + 1] == opposite_parity(last_parity) and
            not has_two_adjacent_snakes(grid, (head[0], head[1] + 1)))): # south
        new_grid = copy.deepcopy(grid)
        new_grid[head[0]][head[1] + 1] = snake(opposite_parity(last_parity))
        if neck is not None:
            try:
                fill_empty_cells_around(new_grid, neck, opposite_parity(last_parity))
            except ValueError:  # Cell is required to be both even and odd---abort!
                return
        recursive_dfs(new_grid, (head[0], head[1] + 1), head, opposite_parity(last_parity))

# We attempt to start in each cell in the grid, though to cut our search space
# and time in half, we constrain start_y <= start_x. Any snakes that would be
# found with start_y > start_x will be a mirror of one that will be found here.
# This still leaves quite a few duplicates by reflection and rotation, but those
# will be harder to weed out by simply adjusting the ranges here.
for start_x in range(9):
    for start_y in range(start_x + 1):
        empty_grid = []
        for i in range(9):
            empty_grid.append([EMPTY] * 9)
        empty_grid[start_x][start_y] = ODD_SNAKE
        recursive_dfs(empty_grid, (start_x, start_y), None, ODD)