#!/usr/bin/env python3
import time

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

def draw(m, visited_pos, obstacle_pos, guard):
    print("\033c") # clear screen

    # draw the guard in blue, the obstacles in red, and the path in yellow
    for y in range(len(m)):
        for x in range(len(m[y])):
            if guard.pos.x == x and guard.pos.y == y:
                print("\033[34m", end="")
                if guard.direction == UP:
                    print("^", end="")
                elif guard.direction == DOWN:
                    print("v", end="")
                elif guard.direction == LEFT:
                    print("<", end="")
                elif guard.direction == RIGHT:
                    print(">", end="")
                print("\033[0m", end="")
            else:
                if Point(x, y) in visited_pos:
                    print("\033[33m", end="")
                    print(".", end="")
                    print("\033[0m", end="")
                elif Point(x, y) in obstacle_pos:
                    print("\033[31m", end="")
                    print("#", end="")
                    print("\033[0m", end="")
                else:
                    print(" ", end="")
    
        print()

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 = 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))

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)
    draw(m, visited_pos, obstacle_pos, guard)
    time.sleep(0.1)