from copy import deepcopy
from functools import cache

from tqdm import tqdm


def parse_input(filedata):
    return filedata.splitlines()


def solve1(inputs):
    grid = list(map(list, inputs))

    def _drop_stone(y, x):
        ny = y
        while ny - 1 >= 0 and grid[ny - 1][x] == ".":
            ny -= 1
        grid[ny][x], grid[y][x] = grid[y][x], grid[ny][x]

    for y in range(len(grid)):
        for x in range(len(grid[0])):
            if grid[y][x] == "O":
                _drop_stone(y, x)

    return sum(value * row.count("O") for value, row in enumerate(grid[::-1], start=1))


def solve2(inputs):
    @cache
    def _rotate_right(grid):
        return tuple(map(lambda x: "".join(x), zip(*grid[::-1])))

    grid = tuple(inputs)
    for _ in range(3):
        grid = _rotate_right(grid)

    @cache
    def _drop_row(row: str):
        return "#".join("".join(sorted(x, reverse=True)) for x in row.split("#"))

    @cache
    def _drop_stones(grid):
        return tuple(map(_drop_row, grid))

    @cache
    def _cycle(grid):
        for _ in range(4):
            grid = _drop_stones(grid)
            grid = _rotate_right(grid)
        return grid

    for _ in tqdm(range(1_000_000_000), dynamic_ncols=True):
        grid = _cycle(grid)

    grid = _rotate_right(grid)
    return sum(value * row.count("O") for value, row in enumerate(grid[::-1], start=1))


def main():
    with open("data/input.txt") as f:
        inputs = parse_input(f.read())

    print(f"Solution 1: {solve1(deepcopy(inputs))}")
    print(f"Solution 2: {solve2(deepcopy(inputs))}")


if __name__ == "__main__":
    main()