from collections import defaultdict
from tools.aoc import AOCDay
from typing import Any


def get_space_below(brick: tuple[tuple], max_y: dict[tuple, int]) -> int:
    below_y = 0
    for x in range(brick[0][0], brick[-1][0] + 1):
        for z in range(brick[0][2], brick[-1][2] + 1):
            below_y = max(below_y, max_y[(x, z)])

    return brick[0][1] - below_y - 1


def drop(bricks: set[tuple]) -> (set[tuple], int):
    dropped_bricks = set()
    max_y = defaultdict(int)
    count = 0
    for brick in sorted(bricks, key=lambda b: b[0][1]):
        space_below = get_space_below(brick, max_y)
        new_brick = tuple()
        for c in brick:
            new_y = c[1] - space_below
            max_y[(c[0], c[2])] = max(max_y[(c[0], c[2])], new_y)
            new_brick += ((c[0], new_y, c[2]),)

        dropped_bricks.add(new_brick)
        if space_below:
            count += 1

    return dropped_bricks, count


def can_disintegrate(brick: tuple, bricks: set[tuple]) -> bool:
    _, count = drop(bricks - {brick})
    return count == 0


def would_fall(brick: tuple, bricks: set[tuple]) -> int:
    _, count = drop(bricks - {brick})
    return count


class Day(AOCDay):
    inputs = [
        [
            (5, "input22_test"),
            (527, "input22"),
        ],
        [
            (7, "input22_test"),
            (100376, "input22"),
        ]
    ]

    def parse_input(self) -> set[tuple]:
        bricks = set()
        for line in self.getInput():
            a, b = line.split("~")
            x1, z1, y1 = map(int, a.split(","))
            x2, z2, y2 = map(int, b.split(","))
            bricks.add(
                tuple(
                    (x, y, z)
                    for x in range(x1, x2 + 1)
                    for y in range(y1, y2 + 1)
                    for z in range(z1, z2 + 1)
                )
            )

        return bricks

    def part1(self) -> Any:
        bricks = self.parse_input()
        dropped_bricks, _ = drop(bricks)
        return sum(can_disintegrate(x, dropped_bricks) for x in dropped_bricks)

    def part2(self) -> Any:
        bricks = self.parse_input()
        dropped_bricks, _ = drop(bricks)
        return sum(would_fall(x, dropped_bricks) for x in dropped_bricks)


if __name__ == '__main__':
    day = Day(2023, 22)
    day.run(verbose=True)