#!/usr/bin/env python3
import collections

class Direction:
    def __init__(self, pos_change_x, pos_change_y):
        self.pos_change_x = pos_change_x
        self.pos_change_y = pos_change_y
    
    def __eq__(self, other):
        if not isinstance(other, Direction):
            return False
        
        return self.pos_change_x == other.pos_change_x and self.pos_change_y == other.pos_change_y

    def __hash__(self):
        return hash((self.pos_change_x, self.pos_change_y))
    
    def __repr__(self):
        return f"Direction({self.pos_change_x}, {self.pos_change_y})"

class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        
    def __eq__(self, other): 
        if not isinstance(other, Point):
            return False

        return self.x == other.x and self.y == other.y
    
    def __hash__(self):
        return hash((self.x, self.y))
    
    def __repr__(self):
        return f"Point({self.x}, {self.y})"

class Guard:
    def __init__(self, pos, direction):
        self.pos = pos
        self.direction = direction
    
    def __repr__(self):
        return f"Guard({self.pos}, {self.direction})"
    
    def move(self, obstacles):
        if self.direction == UP:
            if (Point(self.pos.x, self.pos.y - 1) in obstacles):
                self.direction = RIGHT
                return Point(self.pos.x, self.pos.y - 1)
            else:
                self.pos.y -= 1
        elif self.direction == DOWN:
            if (Point(self.pos.x, self.pos.y + 1) in obstacles):
                self.direction = LEFT
                return Point(self.pos.x, self.pos.y + 1)
            else:
                self.pos.y += 1
        elif self.direction == LEFT:
            if (Point(self.pos.x - 1, self.pos.y) in obstacles):
                self.direction = UP
                return Point(self.pos.x - 1, self.pos.y)
            else:
                self.pos.x -= 1
        elif self.direction == RIGHT:
            if (Point(self.pos.x + 1, self.pos.y) in obstacles):
                self.direction = DOWN
                return Point(self.pos.x + 1, self.pos.y)
            else:
                self.pos.x += 1

with open("./2024/inputs/day06.txt") as f:
    m = [l.strip() for l in f.readlines()]

UP = Direction(0, -1)
DOWN = Direction(0, 1)
LEFT = Direction(-1, 0)
RIGHT = Direction(1, 0)

guard_starting_pos = None
guard_starting_direction = None
guard = None
obstacle_pos = set()
visited_pos = set()

for y in range(len(m)):
    for x in range(len(m[y])):
        if m[y][x] in ["^", "v", "<", ">"]:
            guard_pos = Point(x, y)

            if m[y][x] == "^":
                guard_direction = UP
            elif m[y][x] == "v":
                guard_direction = DOWN
            elif m[y][x] == "<":
                guard_direction = LEFT
            elif m[y][x] == ">":
                guard_direction = RIGHT
            
            guard = Guard(guard_pos, guard_direction)
            
        elif m[y][x] == "#":
            obstacle_pos.add(Point(x, y))

guard_starting_pos = Point(guard.pos.x, guard.pos.y)
guard_starting_direction = Direction(guard.direction.pos_change_x, guard.direction.pos_change_y)

# if the guard is still in the map, move him and register his position in the
# visited_pos set if he hasn't been there before. to move him, check if he has
# an obstacle right in front of him: if he does he moves 90 degrees to the right,
# if he doesn't, he takes a step forward.
while guard.pos.x >= 0 and guard.pos.x < len(m[0]) and guard.pos.y >= 0 and guard.pos.y < len(m):
    if guard.pos not in visited_pos:
        visited_pos.add(guard.pos)
    
    guard.move(obstacle_pos)
    
print(len(visited_pos))

# now we have to find the number of positions where an obstacle can be added so that
# the guard is forced in a loop and never exits the map.
# we can do this by saving a set of the last 4 obstacles hit, and checking if a new set of 4 obstacles is
# the same 4 obstacles. if it is, we have found a loop.
loop_obstacle_pos = set()

for y in range(len(m)):
    for x in range(len(m[y])):
        if Point(x, y) in obstacle_pos or Point(x, y) == guard_starting_pos:
            continue
        
        last_visited_obstacles = collections.deque()
        shadow_visited_obstacles = collections.deque()
        possible_obstacle_pos = Point(x, y)
        new_obstacles = obstacle_pos.copy()
        new_obstacles.add(possible_obstacle_pos)
        guard = Guard(Point(guard_starting_pos.x, guard_starting_pos.y), guard_starting_direction)

        # now solve the map checking for the last 4 obstacles, if a loop is found add the new possible obstacle to the set
        while guard.pos.x >= 0 and guard.pos.x < len(m[0]) and guard.pos.y >= 0 and guard.pos.y < len(m):
            # save the last 4 obstacles
            current_obstacle = guard.move(new_obstacles)
            if current_obstacle is not None:
                if current_obstacle in last_visited_obstacles:
                    if current_obstacle in shadow_visited_obstacles:
                        # guard is in a loop!
                        loop_obstacle_pos.add(possible_obstacle_pos)
                        break
                    else:
                        shadow_visited_obstacles.append(current_obstacle)
                else:
                    last_visited_obstacles.append(current_obstacle)
                    shadow_visited_obstacles = collections.deque()
            
print(len(loop_obstacle_pos))