Advent of Code 2022 - Day 14 - Python

from __future__ import annotations
from enum import Enum
from dataclasses import dataclass


class T(Enum):
    Rock = "#"
    Sand = "+"
    Air = "."


@dataclass(unsafe_hash=True, frozen=True)
class Coordinate:
    x: int
    y: int


class Sand:
    def __init__(self, coordinate: Coordinate, cave: Cave) -> None:
        self.coordinate = coordinate
        self.cave = cave
        self.stopped = False

    def drop_infinite(self) -> Coordinate:
        while not self.stopped:
            next_coordinate = Coordinate(self.coordinate.x, self.coordinate.y + 1)
            if self.cave.rock_coordinates.get(next_coordinate, T.Air) == T.Air:
                self.coordinate = next_coordinate
                if self.coordinate.y == self.cave.y_max:
                    self.stopped = True
                continue

            next_coordinate = Coordinate(self.coordinate.x - 1, self.coordinate.y + 1)
            if self.cave.rock_coordinates.get(next_coordinate, T.Air) == T.Air:
                self.coordinate = next_coordinate
                if self.coordinate.y == self.cave.y_max:
                    self.stopped = True
                continue

            next_coordinate = Coordinate(self.coordinate.x + 1, self.coordinate.y + 1)
            if self.cave.rock_coordinates.get(next_coordinate, T.Air) == T.Air:
                self.coordinate = next_coordinate
                if self.coordinate.y == self.cave.y_max:
                    self.stopped = True
                continue

            self.stopped = True

        return self.coordinate


    def drop_to_floor(self) -> Coordinate:
        while not self.stopped:
            next_coordinate = Coordinate(self.coordinate.x, self.coordinate.y + 1)
            if self.cave.rock_coordinates.get(next_coordinate, T.Air) == T.Air:
                self.coordinate = next_coordinate
                if self.coordinate.y == self.cave.y_max - 1:
                    break
                continue

            next_coordinate = Coordinate(self.coordinate.x - 1, self.coordinate.y + 1)
            if self.cave.rock_coordinates.get(next_coordinate, T.Air) == T.Air:
                self.coordinate = next_coordinate
                if self.coordinate.y == self.cave.y_max - 1:
                    break
                continue

            next_coordinate = Coordinate(self.coordinate.x + 1, self.coordinate.y + 1)
            if self.cave.rock_coordinates.get(next_coordinate, T.Air) == T.Air:
                self.coordinate = next_coordinate
                if self.coordinate.y == self.cave.y_max -1:
                    break
                continue

            self.stopped = True

        return self.coordinate

class Cave:
    def __init__(self):
        self.origin = Coordinate(500, 0)  # sand starts here
        self.rock_coordinates: dict[Coordinate, T] = {}
        self.x_min = 0
        self.x_max = 0
        self.y_min = 0
        self.y_max = 0
        self.rock_paths = []
        self.sand_count = 0

    @classmethod
    def from_file(cls, filepath) -> Cave:
        c = Cave()
        with open(filepath, "r") as f:
            for line in f.readlines():
                c.add_rock_path_from_str(line)
        c.set_bounds()
        return c

    def set_bounds(self):
        for coordinate in self.rock_coordinates:
            self.y_max = max(coordinate.y, self.y_max)
            self.x_min = min(coordinate.x, self.x_min)
            self.x_max = max(coordinate.x, self.x_max)

    def add_rock_path_from_str(self, input_line: str) -> None:
        path: list[Coordinate] = []
        for pair in input_line.split(" -> "):
            path.append(Coordinate(*map(int, pair.split(","))))

        for i, coordinate in enumerate(path[1:], 1):
            self.add_rock_path(path[i - 1], coordinate)

    def add_rock_path(self, start: Coordinate, end: Coordinate):
        self.rock_paths.append((start, end))
        if start.x == end.x:
            x = start.x
            min_y, max_y = min(start.y, end.y), max(start.y, end.y)
            self.draw_vertical_rock(min_y, max_y, x)
        else:
            y = start.y
            min_x, max_x = min(start.x, end.x), max(start.x, end.x)
            self.draw_horizontal_rock(min_x, max_x, y)

    def draw_vertical_rock(self, start, stop, x):
        for y in range(start, stop + 1):
            self.rock_coordinates[Coordinate(x=x, y=y)] = T.Rock

    def draw_horizontal_rock(self, start, stop, y):
        for x in range(start, stop + 1):
            self.rock_coordinates[Coordinate(x=x, y=y)] = T.Rock

    def draw_sand(self, coordinate) -> None:
        self.rock_coordinates[coordinate] = T.Sand

    def simulate_sand(self, print_=False) -> int:
        while True:
            sand = Sand(self.origin, self)
            coordinate = sand.drop_infinite()
            if coordinate.y >= self.y_max:
                break
            else:
                self.draw_sand(coordinate)
            self.sand_count += 1
            if print_:
                print()
                print(f"{self.sand_count}")
                print(f"{self}")

        return self.sand_count

    def simulate_infinite_floor(self, print_=False) -> int:
        while True:
            sand = Sand(self.origin, self)
            coordinate = sand.drop_to_floor()
            self.sand_count += 1
            if coordinate == self.origin:
                break
            else:
                self.draw_sand(coordinate)
            if print_:
                print()
                print(f"{self.sand_count}")
                print(f"{self}")
        return self.sand_count


def part_one(filepath) -> int:
    c = Cave.from_file(filepath)
    sand_count = c.simulate_sand()
    return sand_count


def part_two(filepath) -> int:
    c = Cave.from_file(filepath)
    c.y_max += 2
    sand_count = c.simulate_infinite_floor()
    return sand_count