import aocd
import collections
from enum import Enum, nonmember

class Direction(Enum):
    NORTH = (0, -1)
    WEST = (-1, 0)
    SOUTH = (0, 1)
    EAST = (1, 0)

    @nonmember
    def clockwise_from(d):
        return list(Direction)[list(Direction).index(d) - 1]

def facing(coords, direction):
    return tuple(map(sum, zip(coords, direction.value)))

def parse(data):
    """
    Given maze data, return
    (Tuple of guard coords, set of obstacle coords, grid width, grid height)
    """
    guard = None
    obstacles = set()

    for y, row in enumerate(data.split('\n')):
        for x, char in enumerate(row):
            if char == '#':
                obstacles.add((x, y))
            elif char == '^':
                guard = (x, y)
    
    return (guard, obstacles, x, y)

def simulate(guard, obstacles, width, height, return_visited_tiles=False):
    guard_direction = Direction.NORTH
    visited = collections.defaultdict(set)

    while (0 <= guard[0] <= width) and (0 <= guard[1] <= height):
        visited[guard].add(guard_direction)
        while (facing_tile := facing(guard, guard_direction)) in obstacles:
            guard_direction = Direction.clockwise_from(guard_direction)

        # Loop detection, can only be true in modified grids
        if guard_direction in visited[facing_tile]:
            return True

        guard = facing_tile
    
    # No loop found

    if return_visited_tiles:
        return set(k for k,v in visited.items() if v)
    
    return False


def solve(data):
    guard, obstacles, width, height = parse(data)
    
    # Only bother checking positions that the guard crosses
    potential_positions = simulate(guard, obstacles, width, height, return_visited_tiles=True)

    return sum(simulate(guard, obstacles | {p}, width, height) for p in potential_positions)

if __name__ == '__main__':
    puz = aocd.models.Puzzle(year=2024, day=6)

    # data = puz.examples[0].input_data
    data = puz.input_data

    solution = solve(data)

    print(solution)
    aocd.submit(solution, year=2024, day=6, part='b')