# red (a1, a2, b1, b2)
red = {(0,0),(1,0),(0,1),(1,1)}
# magenta (a4, a5, b4, b5)
magenta = {(3,0),(4,0),(3,1),(4,1)}
# cyan (d1, d2, e1, e2)
cyan = {(0,3),(1,3),(0,4),(1,4)}
# yellow (d4, d5, e4, e5)
yellow = {(3,3),(4,3),(3,4),(4,4)}
# green (a3, b3, c1, c2, c3, c4, c5, d3, e3)
green = {(2,0),(2,1),(0,2),(1,2),(2,2),(3,2),(4,2),(2,3),(2,4)}

#      0   1   2   3   4
#    +---+---+---+---+---+
# 0  | 1 | 0 | 0 | 1 | 1 |	1
#    +---+---+---+---+---+
# 1  | 1 | 1 | 0 | 1 | 0 |	2
#    +---+---+---+---+---+
# 2  | 0 | 1 | 0 | 1 | 0 |	3
#    +---+---+---+---+---+
# 3  | 1 | 1 | 1 | 0 | 0 |	4
#    +---+---+---+---+---+
# 4  | 1 | 1 | 0 | 1 | 1 |	5
#    +---+---+---+---+---+
#      A   B   C   D   E


def cell_from_index(index):
    """
    Convert a linear index to a grid cell (x, y).
    """
    return (index // 5, index % 5)


solved_grid = [
  [1, 0, 0, 1, 1],  # Row 0
  [1, 1, 0, 1, 0],  # Row 1
  [0, 1, 0, 1, 0],  # Row 2
  [1, 1, 1, 0, 0],  # Row 3
  [1, 1, 0, 1, 1]   # Row 4
]

flips = [

  [
      #a1 -> (a1, a2, a4, a5, b4).
      [(0,0), (1,0), (3,0), (4,0), (3,1)],
      #b1 -> (a5, b1, b3, d2, d4).
      [(4,0), (0,1), (2,1), (1,3), (3,3)],
      #c1 -> (a5, c1, c4, c5, d3).
      [(4,0), (0,2), (3,2), (4,2), (2,3)],
      #d1 -> (a1, a4, c1, d1, d2).
      [(0,0), (3,0), (0,2), (0,3), (1,3)],
      #e1 -> (a3, c3, e1, e3, e5).
      [(2,0), (2,2), (0,4), (2,4), (4,4)]
  ],
  [
      #a2 -> (a2, b3. d3, e2, e4).
      [(1,0), (2,1), (2,3), (0,4), (3,4)],
      #b2 -> (b2, b4, c5, e2, e5).
      [(1,1), (3,1), (4,2), (1,4), (4,4)],
      #c2 -> (a3, c2, d4, e1, e3).
      [(2,0), (1,2), (3,3), (0,4), (2,4)],
      #d2 -> (a1, a4, d1, d2, d5).
      [(0,0), (3,0), (0,3), (1,3), (4,3)],
      #e2 -> (b2, c1, c3, d2, e2).
      [(1,1), (0,2), (2,2), (1,3), (1,4)]
  ],
  [
      #a3 -> (a3, c3, c5, e3, e4).
      [(2,0), (2,2), (4,2), (2,4), (3,4)],
      #b3 -> (a2, a4, b3, b5, d3).
      [(1,0), (3,0), (2,1), (4,1), (2,3)],
      #c3 -> (b1, b2, c3, e3, e5).
      [(0,1), (1,1), (2,2), (2,4), (4,4)],
      #d3 -> (a2, a3, c4, d2, d3).
      [(1,0), (2,0), (3,2), (1,3), (2,3)],
      #e3 -> (a1, a3, b2, e2, e3).
      [(0,0), (2,0), (1,1), (1,4), (2,4)],
  ],
  [
      #a4 -> (a4, b1, c3, c4, d4).
      [(3,0), (0,1), (2,2), (3,2), (3,3)],
      #b4 -> (a5, b4, c1, c2, e1).
      [(4,0), (3,1), (0,2), (1,2), (0,4)],
      #c4 -> (a2, b1, c4, e1, e4).
      [(1,0), (0,1), (3,2), (0,4), (3,4)],
      #d4 -> (b5, c4, c5, d4, d5).
      [(4,1), (3,2), (4,2), (3,3), (4,3)],
      #e4 -> (c2, d1, d4, d5, e4).
      [(1,2), (0,3), (3,3), (4,3), (3,4)],
  ],
  [
      #a5 -> (a5, b3, c2, e2, e5).
      [(4, 0), (1, 1), (2, 1), (1, 4), (4, 4)],
      #b5 -> (b4, b5, c2, d1, d3).
      [(3,1), (4,1), (1,2), (0,3), (2,3)],
      #c5 -> (b2, b3, c5, d5, e1).
      [(1,1), (2,1), (4,2), (4,3), (0,4)],
      #d5 -> (b1, b5, d1, d5, e4).
      [(0,1), (4,1), (0,3), (4,3), (3,4)],
      #e5 -> (a1, b4, b5, c1, e5).
      [(0,0), (3,1), (4,1), (0,2), (4,4)]
  ]
]

def flip_grid(grid, flip, pressed_cell):
    """
    Flip the cells in the grid according to the flip pattern.
    """
    if grid[pressed_cell[0]][pressed_cell[1]] == 1:
        return None
    # Create a copy of the grid to avoid modifying the original
    local_grid = [row[:] for row in grid]
    for cell in flip:
        x, y = cell
        local_grid[x][y] = 1 - local_grid[x][y]  # Toggle between 0 and 1
    if is_valid(local_grid):
        return local_grid
    #print_grid(local_grid)
    return None

def is_valid(grid):
    #print_grid(grid)
    return any(grid[x][y] == 1 for x,y in red) and \
           any(grid[x][y] == 1 for x,y in magenta) and \
           any(grid[x][y] == 1 for x,y in cyan) and \
           any(grid[x][y] == 1 for x,y in yellow) and \
           any(grid[x][y] == 1 for x,y in green)

def is_solved(grid):
    """
    Check if the grid is solved
    """
    return grid == solved_grid

possible_prefixes = []

# b1 is first flip
indexes = set([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24])

# todo: for each prefix, generate all possible next prefixes and give them with previous prefix length
# -> can generate the prefix grid and just try the last steps
def generate_prefixes(prefixes):
    for prefix in prefixes:
        for index in indexes - set(prefix):
            yield prefix + [index]

def find_solution():
    """
    Find a solution to the lights out puzzle.
    """
    from random import shuffle
    step = 0
    possible_prefixes = [[1]]  # Start with the first flip (b1)
    while True:
        current_prefixes = []
        for indexes in generate_prefixes(possible_prefixes):
            grid = [[0] * 5 for _ in range(5)]  # Initialize a 5x5 grid with all lights off
            #input(f"Trying path: {indexes}")
            solution = []
            bad_prefix = False
            for index in indexes:
                cell = cell_from_index(index)
                #print(f"Flipping cell {cell}")
                temp_grid = flip_grid(grid, flips[cell[0]][cell[1]], cell)
                if temp_grid is None:
                    #print(f"Invalid flip at {cell}, skipping.")
                    bad_prefix = True
                    break
                #print_grid(temp_grid)
                grid = temp_grid
                solution.append(cell)
                if is_solved(grid):
                    return grid, solution
            if not bad_prefix:
                current_prefixes.append(indexes)
        possible_prefixes = current_prefixes[:]
        print(f"{len(possible_prefixes)} prefixes of length {len(possible_prefixes[0])} found.")
    return None, None

def print_grid(grid):
    """
    Print the grid in a readable format.
    """
    for row in grid:
        print(" ".join(str(cell) for cell in row))
    print()

if __name__ == "__main__":
    # Find a solution
    solved_grid, solution = find_solution()
    print(len(possible_prefixes), "possible prefixes found.")

    if solution:
        print("Solution Found:")
        print("Flips applied:", solution)
        print_grid(solved_grid)
    else:
        print("No solution found.")