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


def get_space_below(brick: tuple[tuple[int, int, int], ...], max_y: dict[tuple[int, int], 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[tuple[int, int, int], ...]]) -> (set[tuple[tuple[int, int, int], ...]], 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


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

    def parse_input(self) -> set[tuple[tuple[int, int, int], ...]]:
        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:
        dropped_bricks, _ = drop(self.parse_input())
        return sum(c[1] == 0 for c in (drop(dropped_bricks - {x}) for x in dropped_bricks))

    def part2(self) -> Any:
        dropped_bricks, _ = drop(self.parse_input())
        return sum(c[1] for c in (drop(dropped_bricks - {x}) for x in dropped_bricks))


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