from tools.aoc import AOCDay
from tools.grid import Grid
from typing import Any


SORT_KEYS = {
    (0, -1): lambda c: c.y,
    (-1, 0): None,
    (0, 1): lambda c: -c.y,
    (1, 0): lambda c: -c.x,
}


def load_after_roll(platform: Grid, direction: tuple[int, int]) -> int:
    load = 0
    for c in sorted(platform.getActiveCells(), key=SORT_KEYS[direction]):
        if platform.get(c) == "#":
            continue

        t = c
        while not platform.get(t + direction) and platform.isWithinBoundaries(t + direction):
            t += direction

        platform.set(c, False)
        platform.set(t, "O")
        load += platform.maxY - t[1] + 1

    return load


class Day(AOCDay):
    inputs = [
        [
            (136, "input14_test"),
            (110565, "input14_dennis"),
            (110407, "input14"),
        ],
        [
            (64, "input14_test"),
            (89845, "input14_dennis"),
            (87273, "input14"),
        ],
    ]

    def part1(self) -> Any:
        platform = Grid.from_data(self.getInput(), translate={".": False})
        return load_after_roll(platform, (0, -1))

    def part2(self) -> Any:
        platform = Grid.from_data(self.getInput(), translate={".": False})
        DP = {}
        cycle_loads = {}
        directions = [(0, -1), (-1, 0), (0, 1), (1, 0)]
        max_cycles = 1_000_000_000
        cycle = 0

        while cycle < max_cycles:
            cycle += 1
            load = 0
            for d in directions:
                load = load_after_roll(platform, d)
            cycle_loads[cycle] = load

            if (load, platform) in DP:
                repeat_start = DP[(load, platform)]
                lookup_cycle = (max_cycles - repeat_start) % (cycle - repeat_start) + repeat_start
                return cycle_loads[lookup_cycle]
            else:
                DP[(load, platform)] = cycle

        return ""


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