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


def load_after_roll(platform: Grid, direction: Coordinate) -> int:
    load = 0
    sort_order = {
        Coordinate(0, -1): sorted(platform.getActiveCells(), key=lambda c: c.y),
        Coordinate(-1, 0): sorted(platform.getActiveCells()),
        Coordinate(0, 1): sorted(platform.getActiveCells(), key=lambda c: -c.y),
        Coordinate(1, 0): sorted(platform.getActiveCells(), key=lambda c: -c.x),
    }
    for c in sort_order[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"),
            (110407, "input14"),
        ],
        [
            (64, "input14_test"),
            (87273, "input14"),
        ],
    ]

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

    def part2(self) -> Any:
        platform = Grid.from_data(self.getInput(), translate={".": False})
        DP = {}
        cycle_loads = {}
        directions = [Coordinate(0, -1), Coordinate(-1, 0), Coordinate(0, 1), Coordinate(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

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

        return ""


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