zerosleeps

from itertools import pairwise
from unittest import TestCase
from utils import get_raw_input


class Cave:
    def __init__(self, raw_input):
        self.build(raw_input)
        self.sand = set()
        self.overflowing = False

    def plot_rocks(self, x1, y1, x2, y2):
        if y1 == y2:
            return [(i, y1) for i in range(min([x1, x2]), max([x1, x2]) + 1)]
        else:
            return [(x1, i) for i in range(min([y1, y2]), max([y1, y2]) + 1)]

    def build(self, raw_input):
        self.rocks = set()
        for line in raw_input:
            for pair in pairwise(line.split(" -> ")):
                self.rocks.update(
                    self.plot_rocks(
                        int(pair[0].split(",")[0]),
                        int(pair[0].split(",")[1]),
                        int(pair[1].split(",")[0]),
                        int(pair[1].split(",")[1]),
                    )
                )

    @property
    def boundaries(self):
        # y_min is highest rock, y_max is lowest rock
        return {
            "x_min": min([rock[0] for rock in self.rocks]),
            "x_max": max([rock[0] for rock in self.rocks]),
            "y_min": min([rock[1] for rock in self.rocks]),
            "y_max": max([rock[1] for rock in self.rocks]),
        }

    def drop_sand(self, floor=False):
        sand_position = (500, 0)
        while True:
            if (
                not floor
                and (
                    sand_position[1] > self.boundaries["y_max"]
                    or sand_position[0] < self.boundaries["x_min"]
                    or sand_position[0] > self.boundaries["x_max"]
                )
            ) or (floor and (500, 0) in self.sand):
                # Out of bounds
                self.overflowing = True
                return
            elif (
                (sand_position[0], sand_position[1] + 1) not in self.rocks
                and (sand_position[0], sand_position[1] + 1) not in self.sand
                and (
                    not floor
                    or (floor and sand_position[1] + 1 < self.boundaries["y_max"] + 2)
                )
            ):
                # Can fall straight down
                sand_position = (sand_position[0], sand_position[1] + 1)
            elif (
                (sand_position[0] - 1, sand_position[1] + 1) not in self.rocks
                and (sand_position[0] - 1, sand_position[1] + 1) not in self.sand
                and (
                    (not floor)
                    or (floor and sand_position[1] + 1 < self.boundaries["y_max"] + 2)
                )
            ):
                # Can fall down-and-left
                sand_position = (sand_position[0] - 1, sand_position[1] + 1)
            elif (
                (sand_position[0] + 1, sand_position[1] + 1) not in self.rocks
                and (sand_position[0] + 1, sand_position[1] + 1) not in self.sand
                and (
                    (not floor)
                    or (floor and sand_position[1] + 1 < self.boundaries["y_max"] + 2)
                )
            ):
                # Can fall down-and-right
                sand_position = (sand_position[0] + 1, sand_position[1] + 1)
            else:
                # Can't move
                self.sand.add(sand_position)
                return


class TestExamples(TestCase):
    def setUp(self):
        self.example_input = [
            "498,4 -> 498,6 -> 496,6",
            "503,4 -> 502,4 -> 502,9 -> 494,9",
        ]
        self.cave = Cave(self.example_input)

    def test_part_one(self):
        self.assertEqual(run(self.cave), 24)

    def test_part_two(self):
        self.assertEqual(run(self.cave, True), 93)


def run(cave, floor=False):
    while not cave.overflowing:
        cave.drop_sand(floor)
    return len(cave.sand)


if __name__ == "__main__":
    print(f"Part one: {run(Cave(get_raw_input('day_14_input.txt')))}")
    print(f"Part two: {run(Cave(get_raw_input('day_14_input.txt')), True)}")