from collections import deque
from enum import Enum
from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from tools.tools import compare
from typing import Any


SAND_SOURCE = Coordinate(500, 0)


class TType(Enum):
    AIR = ' '
    ROCK = '#'
    SAND = 'o'


class Day(AOCDay):
    inputs = [
        [
            (24, "input14_test"),
            (994, "input14_dennis"),
            (873, "input14"),
        ],
        [
            (93, "input14_test"),
            (26283, "input14_dennis"),
            (24813, "input14"),
        ]
    ]

    def get_cave(self) -> Grid:
        cave = Grid(TType.AIR)

        for line in self.getInput():
            corners = []
            for c in line.split(" -> "):
                x, y = map(int, c.split(","))
                corners.append(Coordinate(x, y))
            for i in range(1, len(corners)):
                for c in corners[i].getLineTo(corners[i-1]):
                    cave.set(c, TType.ROCK)

        cave.minY = 0
        return cave

    def get_cave_rocks_and_height(self) -> (set, int):
        rocks = set()
        max_y = 0
        for line in self.getInput():
            corners = [tuple(map(int, c.split(","))) for c in line.split(" -> ")]
            for i in range(len(corners)):
                if corners[i][1] > max_y:
                    max_y = corners[i][1]
                if i == 0:
                    continue
                from_x = corners[i][0] if corners[i][0] < corners[i-1][0] else corners[i-1][0]
                to_x = (corners[i][0] if corners[i][0] > corners[i-1][0] else corners[i-1][0]) + 1
                from_y = corners[i][1] if corners[i][1] < corners[i-1][1] else corners[i-1][1]
                to_y = (corners[i][1] if corners[i][1] > corners[i-1][1] else corners[i-1][1]) + 1
                for x in range(from_x, to_x):
                    for y in range(from_y, to_y):
                        rocks.add((x, y))

        return rocks, max_y


    def part1(self) -> Any:
        cave = self.get_cave()

        while True:
            tx, ty = 0, 0
            while ty <= cave.maxY:
                ty += 1
                for c in [Coordinate(tx, ty), Coordinate(tx - 1, ty), Coordinate(tx + 1, ty)]:
                    if cave.get(SAND_SOURCE + c) == TType.AIR:
                        tx += compare(c.x, tx)
                        break
                else:
                    cave.set(SAND_SOURCE + Coordinate(tx, ty - 1), TType.SAND)
                    break
            else:
                break

        return cave.count(TType.SAND)

    def part2(self) -> Any:
        rocks, max_y = self.get_cave_rocks_and_height()
        v = set()
        q = deque()
        q.append((SAND_SOURCE.x, SAND_SOURCE.y))
        while q:
            current = q.pop()
            if current in v or current[1] > max_y + 1:
                continue
            v.add(current)
            for dx in [-1, 0, 1]:
                if (current[0] + dx, current[1] + 1) not in rocks:
                    q.append((current[0] + dx, current[1] + 1))

        return len(v)


if __name__ == '__main__':
    day = Day(2022, 14)
    day.run(verbose=True)