from utils import open_day, sliding_window
from enum import Enum, auto
from itertools import count

def point_from_str(s):
    return tuple(int(d) for d in s.split(','))

class Tile(Enum):
    WALL = auto()
    SAND = auto()

class Sandbox:
    def __init__(self, source=(500, 0)):
        self.tiles = dict()
        self.source = source
        self.minx, self.maxx = source[0], source[0]
        self.miny, self.maxy = source[1], source[1]
    def __setitem__(self, key, value):
        if   key[0] < self.minx: self.minx = key[0]
        elif key[0] > self.maxx: self.maxx = key[0]
        if   key[1] < self.miny: self.miny = key[1]
        elif key[1] > self.maxy: self.maxy = key[1]
        self.tiles[key] = value
    def __getitem__(self, key):
        return self.tiles[key]
    def __str__(self):
        lines = []
        for y in range(self.miny, self.maxy + 1):
            line = []
            for x in range(self.minx, self.maxx + 1):
                p = x, y
                match self.tiles.get(p):
                    case Tile.WALL: line.append('#')
                    case Tile.SAND: line.append('o')
                    case None:      line.append('.')
                if x == self.source[0] and y == self.source[1]:
                    line[-1] = '+'
            lines.append(''.join(line))
        return '\n'.join(lines)
    def simulate(self):
        for units in count():
            sandpos = self.source
            while self.minx <= sandpos[0] <= self.maxx and \
                    self.miny <= sandpos[1] <= self.maxy and \
                    self.source not in self.tiles:
                for dx in (0, -1, 1):
                    npos = sandpos[0] + dx, sandpos[1] + 1
                    if npos not in self.tiles:
                        sandpos = (sandpos[0] + dx, sandpos[1] + 1)
                        break
                else:
                    self.tiles[sandpos] = Tile.SAND
                    break
            else:
                break
        return units

sandbox = Sandbox()

with open_day(14) as f:
    for line in f:
        for b, e in sliding_window(map(point_from_str, line.rstrip().split(' -> ')), 2):
            minlx, maxlx = min(b[0], e[0]), max(b[0], e[0])
            minly, maxly = min(b[1], e[1]), max(b[1], e[1])
            for x in range(minlx, maxlx + 1):
                for y in range(minly, maxly + 1):
                    sandbox[x, y] = Tile.WALL

p1 = sandbox.simulate()
print(p1)

floory = sandbox.maxy + 2
height = floory - sandbox.source[1]
for x in range(500 - height, 500 + height + 1):
    sandbox[x, floory] = Tile.WALL

print(p1 + sandbox.simulate())